From 9cf157c2521bf9a1f866fd2a8c4a9ff83696085a Mon Sep 17 00:00:00 2001 From: Jialin Ouyang Date: Mon, 31 Aug 2026 09:26:20 -0700 Subject: [PATCH] [Radix Cache] Add Rust TreeCore backend with shared parity tests (#32710) Co-authored-by: alphabetc1 <2508695655@qq.com> Co-authored-by: ispobock --- .github/actions/download-rust-ext/action.yml | 14 +- .github/workflows/_pr-test-rust-ext-build.yml | 120 +- .github/workflows/release-pypi-nightly.yml | 42 +- .github/workflows/release-pypi-pr.yml | 42 +- .github/workflows/release-pypi.yml | 14 +- .github/workflows/seed-rust-ext-cache.yml | 3 + .gitignore | 1 + .pre-commit-config.yaml | 4 +- docker/Dockerfile | 5 +- docker/Dockerfile.cu134 | 3 + python/pyproject.toml | 14 +- python/setup.py | 76 +- .../srt/arg_groups/pd_disaggregation_hook.py | 1 + .../disaggregation/decode_hicache_mixin.py | 15 +- python/sglang/srt/managers/schedule_policy.py | 42 +- .../sglang/srt/mem_cache/base_prefix_cache.py | 41 + .../srt/mem_cache/buffer_mode/pipeline.py | 242 +- .../srt/mem_cache/rust_tree_core/.gitignore | 1 + .../srt/mem_cache/rust_tree_core/__init__.py | 1 + .../srt/mem_cache/rust_tree_core/adapter.py | 971 +++ .../srt/mem_cache/rust_tree_core/extension.py | 40 + .../unified_cache/tree_core_registry.py | 10 + .../unified_cache/unified_tree_core.py | 61 +- .../unified_tree_core_interface.py | 35 + .../srt/mem_cache/unified_radix_cache.py | 12 +- python/sglang/srt/rust_extensions/loader.py | 115 +- .../sglang/srt/rust_extensions/torch_build.py | 115 + python/sglang/test/test_utils.py | 11 + rust/Cargo.toml | 6 + rust/mem-cache/.gitignore | 1 + rust/mem-cache/Cargo.lock | 905 ++ rust/mem-cache/Cargo.toml | 40 + rust/mem-cache/README.md | 42 + rust/mem-cache/src/components/full.rs | 493 ++ rust/mem-cache/src/components/mamba.rs | 751 ++ rust/mem-cache/src/components/mod.rs | 491 ++ rust/mem-cache/src/components/swa.rs | 1197 +++ rust/mem-cache/src/lib.rs | 27 + rust/mem-cache/src/node.rs | 1571 ++++ rust/mem-cache/src/python_bindings.rs | 3261 +++++++ rust/mem-cache/src/tests/components/base.rs | 160 + rust/mem-cache/src/tests/components/full.rs | 2828 ++++++ rust/mem-cache/src/tests/components/mamba.rs | 1784 ++++ rust/mem-cache/src/tests/components/swa.rs | 4894 +++++++++++ rust/mem-cache/src/tests/node.rs | 1942 +++++ rust/mem-cache/src/tests/test_utils.rs | 46 + rust/mem-cache/src/tests/unified_lru_list.rs | 708 ++ rust/mem-cache/src/tests/unified_tree_core.rs | 7739 +++++++++++++++++ rust/mem-cache/src/unified_lru_list.rs | 501 ++ rust/mem-cache/src/unified_tree_core.rs | 4811 ++++++++++ rust/mem-cache/torch_2_13_compat.h | 15 + scripts/ci/cuda/ci_install_dependency.sh | 34 +- scripts/ci/utils/stage_rust_ext_modules.sh | 23 +- scripts/release/prepare_sglang_wheel.py | 203 + .../test_unified_radix_cache_kl_dsv4.py | 68 +- .../test_unified_radix_cache_kl_full.py | 12 +- ..._unified_radix_cache_kl_hybrid_bitexact.py | 36 +- .../rust/test_run_mem_cache_rust_tests.py | 64 + test/registered/rust/test_run_rust_tests.py | 37 +- test/registered/rust/test_rust_extension.py | 215 +- .../test_decode_hicache_tree_core.py | 107 + .../test_schedule_policy_dfs_weight.py | 55 + .../rust_unified_tree_core_inspector.py | 228 + .../unit/mem_cache/test_rust_tree_core.py | 149 + .../test_rust_tree_core_integration.py | 1973 +++++ .../test_rust_unified_radix_cache_bench.py | 29 + .../test_rust_unified_radix_cache_unittest.py | 30 + .../test_unified_radix_cache_bench.py | 34 +- .../test_unified_radix_cache_unittest.py | 766 +- .../unified_tree_core_inspection_interface.py | 5 + .../mem_cache/unified_tree_core_inspector.py | 10 + .../unit/server_args/test_server_args.py | 27 +- 72 files changed, 39973 insertions(+), 396 deletions(-) create mode 100644 python/sglang/srt/mem_cache/rust_tree_core/.gitignore create mode 100644 python/sglang/srt/mem_cache/rust_tree_core/__init__.py create mode 100644 python/sglang/srt/mem_cache/rust_tree_core/adapter.py create mode 100644 python/sglang/srt/mem_cache/rust_tree_core/extension.py create mode 100644 python/sglang/srt/rust_extensions/torch_build.py create mode 100644 rust/mem-cache/.gitignore create mode 100644 rust/mem-cache/Cargo.lock create mode 100644 rust/mem-cache/Cargo.toml create mode 100644 rust/mem-cache/README.md create mode 100644 rust/mem-cache/src/components/full.rs create mode 100644 rust/mem-cache/src/components/mamba.rs create mode 100644 rust/mem-cache/src/components/mod.rs create mode 100644 rust/mem-cache/src/components/swa.rs create mode 100644 rust/mem-cache/src/lib.rs create mode 100644 rust/mem-cache/src/node.rs create mode 100644 rust/mem-cache/src/python_bindings.rs create mode 100644 rust/mem-cache/src/tests/components/base.rs create mode 100644 rust/mem-cache/src/tests/components/full.rs create mode 100644 rust/mem-cache/src/tests/components/mamba.rs create mode 100644 rust/mem-cache/src/tests/components/swa.rs create mode 100644 rust/mem-cache/src/tests/node.rs create mode 100644 rust/mem-cache/src/tests/test_utils.rs create mode 100644 rust/mem-cache/src/tests/unified_lru_list.rs create mode 100644 rust/mem-cache/src/tests/unified_tree_core.rs create mode 100644 rust/mem-cache/src/unified_lru_list.rs create mode 100644 rust/mem-cache/src/unified_tree_core.rs create mode 100644 rust/mem-cache/torch_2_13_compat.h create mode 100755 scripts/release/prepare_sglang_wheel.py create mode 100644 test/registered/rust/test_run_mem_cache_rust_tests.py create mode 100644 test/registered/unit/disaggregation/test_decode_hicache_tree_core.py create mode 100644 test/registered/unit/managers/test_schedule_policy_dfs_weight.py create mode 100644 test/registered/unit/mem_cache/rust_unified_tree_core_inspector.py create mode 100644 test/registered/unit/mem_cache/test_rust_tree_core.py create mode 100644 test/registered/unit/mem_cache/test_rust_tree_core_integration.py create mode 100644 test/registered/unit/mem_cache/test_rust_unified_radix_cache_bench.py create mode 100644 test/registered/unit/mem_cache/test_rust_unified_radix_cache_unittest.py diff --git a/.github/actions/download-rust-ext/action.yml b/.github/actions/download-rust-ext/action.yml index 78893e1ad..ee5183102 100644 --- a/.github/actions/download-rust-ext/action.yml +++ b/.github/actions/download-rust-ext/action.yml @@ -1,8 +1,9 @@ name: 'Download prebuilt Rust extensions' description: > Put rust-ext-build's PyO3 extension modules in the checkout and set - SGLANG_BUILD_RUST_EXTS=none for the job, so install skips the cargo build. - Sets nothing when neither source has them, leaving install to compile. + SGLANG_BUILD_RUST_EXTS=none and SGLANG_RUST_BUILD_MODE=never for the job, so + install skips the cargo build and runtime trusts these fingerprinted modules. + Sets nothing when neither source has them, leaving source builds enabled. inputs: artifact_name: @@ -49,8 +50,10 @@ runs: if: steps.artifact.outcome != 'success' uses: actions/cache/restore@v4 with: - path: python/sglang/srt/rust_extensions/_*.so - key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py') }} + path: | + python/sglang/srt/rust_extensions/_*.so + python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so + key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py', 'python/pyproject.toml', 'python/sglang/srt/rust_extensions/torch_build.py', '.github/workflows/_pr-test-rust-ext-build.yml', 'scripts/ci/utils/stage_rust_ext_modules.sh') }} # Job-wide, but only setup.py reads it, and only while building. # Whether the modules suit this interpreter is not decided here: @@ -65,6 +68,9 @@ runs: || [ "${{ steps.cache.outputs.cache-hit }}" = "true" ]; then echo "hit=true" >> "$GITHUB_OUTPUT" echo "SGLANG_BUILD_RUST_EXTS=none" >> "$GITHUB_ENV" + # A source checkout normally ignores in-package native artifacts + # because they may be stale. These are fingerprint-keyed CI bytes. + echo "SGLANG_RUST_BUILD_MODE=never" >> "$GITHUB_ENV" else echo "hit=false" >> "$GITHUB_OUTPUT" fi diff --git a/.github/workflows/_pr-test-rust-ext-build.yml b/.github/workflows/_pr-test-rust-ext-build.yml index 20a9cfae9..7631186e2 100644 --- a/.github/workflows/_pr-test-rust-ext-build.yml +++ b/.github/workflows/_pr-test-rust-ext-build.yml @@ -82,6 +82,8 @@ jobs: sparse-checkout: | rust python/setup.py + python/pyproject.toml + python/sglang/srt/rust_extensions/torch_build.py .github scripts/ci/utils sparse-checkout-cone-mode: false @@ -93,14 +95,16 @@ jobs: - name: Ensure zstd so the saved entry is readable run: bash scripts/ci/utils/ensure_zstd.sh - # setup.py counts because it selects which crates get built. pyproject.toml - # is left out - it churns on bumps that cannot affect these modules. + # The setup hook and torch helper select and configure what gets built; + # pyproject.toml pins the libtorch ABI used by mem-cache. - name: Restore built modules id: cache uses: actions/cache/restore@v4 with: - path: python/sglang/srt/rust_extensions/_*.so - key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py') }} + path: | + python/sglang/srt/rust_extensions/_*.so + python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so + key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py', 'python/pyproject.toml', 'python/sglang/srt/rust_extensions/torch_build.py', '.github/workflows/_pr-test-rust-ext-build.yml', 'scripts/ci/utils/stage_rust_ext_modules.sh') }} # On a miss: different hash = rust/setup.py moved; no entries = evicted. - name: Report cache lookup @@ -112,7 +116,8 @@ jobs: run: | if [ -n "${MATCHED_KEY}" ]; then echo "hit: ${MATCHED_KEY}" - ls -l python/sglang/srt/rust_extensions/_*.so + ls -l python/sglang/srt/rust_extensions/_*.so \ + python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so else echo "miss: ${PRIMARY_KEY}" echo "entries under ${KEY_PREFIX}- (created / ref / size / key):" @@ -134,7 +139,7 @@ jobs: uses: actions/upload-artifact@v4 with: name: ${{ inputs.artifact_name }} - # Archive holds rust_extensions/_*.so, so it unpacks into python/sglang/srt/. + # Archive holds package-relative paths, so it unpacks into python/sglang/srt/. path: rust-ext-staging/ if-no-files-found: error retention-days: 1 @@ -182,23 +187,21 @@ jobs: export PATH="${CARGO_HOME:-$HOME/.cargo}/bin:${PATH}" # Per-interpreter subdirs (set in the loop): PyO3's fingerprint tracks # the interpreter, so a shared dir rebuilds on every ABI switch. - cargo_target_root="${HOME}/.cache/sglang-cargo-target" + # ci_install_dependency.sh drops ${HOME}/.cache/sglang-cargo-target at 85% + # disk and unlocks before its own build, so a CUDA job sharing this host + # can delete the tree mid compile. Build in a per run dir nothing else + # touches. The .so cache above still carries results across runs. + cargo_target_root="${RUNNER_TEMP:-/tmp}/sglang-cargo-target-${GITHUB_RUN_ID:-norun}-$$" mkdir -p "${cargo_target_root}" - used="$(df --output=pcent "${cargo_target_root}" 2>/dev/null | tr -dc '0-9')" - if [ "${used:-0}" -ge 85 ]; then - echo "cargo target dir filesystem at ${used}%; dropping ${cargo_target_root}" - rm -rf "${cargo_target_root}" - mkdir -p "${cargo_target_root}" - fi python3 -m pip install --upgrade pip command -v uv >/dev/null 2>&1 || pip install uv - # build_rust needs only the build backend, not sglang's ~294 runtime deps. + # build_rust needs the build backend and torch, not sglang's ~294 other runtime deps. # Per-job path: these runners are persistent and shared, so a fixed one # both inherits the previous job's venv and races a concurrent build. venv_root="${RUNNER_TEMP:-/tmp}/sglang-ci-rust-ext-${GITHUB_RUN_ID:-norun}-$$" # Best-effort, like ci_cleanup_venv.sh: under set -e a failing EXIT trap # would fail the step, and nothing here is worth keeping for a postmortem. - trap 'rm -rf "${venv_root}" || true' EXIT + trap 'rm -rf "${venv_root}" "${cargo_target_root}" || true' EXIT for python_bin in "${{ steps.py310.outputs.python-path }}" "${{ steps.py312.outputs.python-path }}"; do minor="$("${python_bin}" -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')" export CARGO_TARGET_DIR="${cargo_target_root}/py${minor}" @@ -206,8 +209,87 @@ jobs: uv venv "${venv}" --python "${python_bin}" --seed # shellcheck disable=SC1091 source "${venv}/bin/activate" - uv pip install "setuptools>=61.0" "setuptools-rust>=1.10" "setuptools-scm>=8.0" wheel + uv pip install "setuptools>=61.0" "setuptools-rust>=1.11" "setuptools-scm>=8.0" "torch==2.13.0" wheel + # torch-sys resolves libtorch from the active interpreter and bakes that + # path into the persistent cargo cache. venv_root is per-run and deleted on + # exit, so a later run reuses a "Fresh torch-sys" whose -L points at a gone + # directory and fails with "unable to find library -ltorch". Pin LIBTORCH and + # drop the torch-shim units so they rebuild against this run's venv. + LIBTORCH="$(python -c 'import pathlib, torch; print(pathlib.Path(torch.__file__).parent)')" + export LIBTORCH + export LD_LIBRARY_PATH="${LIBTORCH}/lib:${LD_LIBRARY_PATH:-}" + cargo clean --release --manifest-path rust/mem-cache/Cargo.toml \ + -p torch-sys -p mem_cache 2>/dev/null || true (cd python && SGLANG_BUILD_RUST_EXTS=all python setup.py build_rust --inplace) + python - <<'PY' + import importlib.util + import os + import pathlib + import runpy + import shutil + import subprocess + import sysconfig + + import torch # noqa: F401 - preload libtorch before the extension + + root = pathlib.Path.cwd() + suffix = sysconfig.get_config_var("EXT_SUFFIX") + production_path = pathlib.Path( + "python/sglang/srt/mem_cache/rust_tree_core/mem_cache" + suffix + ).resolve() + inspection_path = production_path.with_name( + "mem_cache_inspection" + suffix + ) + + helper = runpy.run_path( + root / "python/sglang/srt/rust_extensions/torch_build.py" + ) + build = helper["torch_build_configuration"]( + compat_header=root / "rust/mem-cache/torch_2_13_compat.h", + python_module="sglang.srt.mem_cache.rust_tree_core.mem_cache", + torch_module=torch, + include_absolute_rpath=False, + ) + subprocess.run( + [ + "cargo", + "build", + "--release", + "--locked", + "--manifest-path", + "rust/mem-cache/Cargo.toml", + "--features", + "python-extension,inspection", + ], + env=build.environment, + check=True, + ) + release_dir = pathlib.Path(os.environ["CARGO_TARGET_DIR"]) / "release" + if target := os.environ.get("CARGO_BUILD_TARGET"): + release_dir = pathlib.Path(os.environ["CARGO_TARGET_DIR"]) / target / "release" + shutil.copy2(release_dir / "libmem_cache.so", inspection_path) + + def load(name, path): + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + production = load( + "sglang.srt.mem_cache.rust_tree_core.mem_cache", production_path + ) + inspection = load( + "sglang.srt.mem_cache.rust_tree_core.mem_cache_inspection", + inspection_path, + ) + assert not hasattr( + production.RustUnifiedTreeCoreBinding, "inspect_contains_node" + ) + assert hasattr( + inspection.RustUnifiedTreeCoreBinding, "inspect_contains_node" + ) + PY deactivate done @@ -226,8 +308,10 @@ jobs: - name: Save built modules uses: actions/cache/save@v4 with: - path: python/sglang/srt/rust_extensions/_*.so - key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py') }} + path: | + python/sglang/srt/rust_extensions/_*.so + python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so + key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py', 'python/pyproject.toml', 'python/sglang/srt/rust_extensions/torch_build.py', '.github/workflows/_pr-test-rust-ext-build.yml', 'scripts/ci/utils/stage_rust_ext_modules.sh') }} - name: Upload extension modules uses: actions/upload-artifact@v4 diff --git a/.github/workflows/release-pypi-nightly.yml b/.github/workflows/release-pypi-nightly.yml index 86d1f65d0..b072a8316 100644 --- a/.github/workflows/release-pypi-nightly.yml +++ b/.github/workflows/release-pypi-nightly.yml @@ -24,7 +24,8 @@ jobs: if: github.repository == 'sgl-project/sglang' runs-on: ubuntu-latest outputs: - nightly_version: ${{ steps.build.outputs.nightly_version }} + nightly_version: ${{ steps.wheel.outputs.wheel_version }} + wheel_filename: ${{ steps.wheel.outputs.wheel_filename }} commit_hash: ${{ steps.build.outputs.commit_hash }} build_date: ${{ steps.build.outputs.build_date }} steps: @@ -41,7 +42,9 @@ jobs: - name: Install build dependencies run: | - pip install build wheel setuptools setuptools-scm + pip install \ + auditwheel build patchelf wheel "setuptools>=61.0" \ + "setuptools-rust>=1.11" "setuptools-scm>=8.0" "torch==2.13.0" # Needed by setuptools-rust to build the bundled native gRPC extension # (rust/sglang-grpc) when `python -m build` builds the sglang wheel. @@ -66,7 +69,7 @@ jobs: MINOR=$(echo "$VERSION" | cut -d. -f2) PATCH_RAW=$(echo "$VERSION" | cut -d. -f3) # Strip pre-release suffixes (rc0, post1, etc.) to get numeric patch - PATCH=$(echo "$PATCH_RAW" | sed 's/[^0-9].*//') + PATCH=${PATCH_RAW%%[^0-9]*} NEXT_PATCH=$((PATCH + 1)) NEXT_VERSION="${MAJOR}.${MINOR}.${NEXT_PATCH}" @@ -77,24 +80,25 @@ jobs: export SETUPTOOLS_SCM_PRETEND_VERSION="$FORCE_VERSION" # Build wheel - python3 -m build --wheel - - # Extract version from built wheel filename - WHEEL_FILE=$(ls dist/*.whl) - NIGHTLY_VERSION=$(echo "$WHEEL_FILE" | sed 's/.*sglang-\(.*\)-py3.*/\1/') + python3 -m build --wheel --no-isolation # Get commit info COMMIT_HASH=$(git rev-parse --short HEAD) BUILD_DATE=$(date -u +%Y-%m-%d) - echo "Built wheel: $WHEEL_FILE" - echo "Nightly version: ${NIGHTLY_VERSION}" echo "Commit: ${COMMIT_HASH}" echo "Build date: ${BUILD_DATE}" - echo "nightly_version=${NIGHTLY_VERSION}" >> $GITHUB_OUTPUT - echo "commit_hash=${COMMIT_HASH}" >> $GITHUB_OUTPUT - echo "build_date=${BUILD_DATE}" >> $GITHUB_OUTPUT + { + echo "commit_hash=${COMMIT_HASH}" + echo "build_date=${BUILD_DATE}" + } >> "$GITHUB_OUTPUT" + + - name: Repair and smoke-test wheel + id: wheel + run: | + python3 scripts/release/prepare_sglang_wheel.py python/dist \ + --github-output "$GITHUB_OUTPUT" - name: Upload wheel artifact uses: actions/upload-artifact@v4 @@ -139,7 +143,7 @@ jobs: token: ${{ secrets.GH_PAT_FOR_WHL_RELEASE }} prerelease: true body: | - Nightly build from commit ${{ github.sha }} + Nightly build from commit ${{ needs.build-nightly-wheel.outputs.commit_hash }} Build date: ${{ needs.build-nightly-wheel.outputs.build_date }} Version: ${{ needs.build-nightly-wheel.outputs.nightly_version }} files: | @@ -147,7 +151,7 @@ jobs: - name: Clone wheel index repository run: | - git clone https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git sgl-whl + git clone "https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git" sgl-whl cd sgl-whl git config --local user.name "sglang-bot" git config --local user.email "sglangbot@gmail.com" @@ -162,10 +166,10 @@ jobs: - name: Update wheel index run: | python3 scripts/update_nightly_whl_index.py \ - --commit-hash ${{ needs.build-nightly-wheel.outputs.commit_hash }} \ - --nightly-version ${{ needs.build-nightly-wheel.outputs.nightly_version }} \ - --cuda-version ${{ matrix.cuda_version }} \ - --build-date ${{ needs.build-nightly-wheel.outputs.build_date }} + --commit-hash "${{ needs.build-nightly-wheel.outputs.commit_hash }}" \ + --nightly-version "${{ needs.build-nightly-wheel.outputs.nightly_version }}" \ + --cuda-version "${{ matrix.cuda_version }}" \ + --build-date "${{ needs.build-nightly-wheel.outputs.build_date }}" - name: Push wheel index run: | diff --git a/.github/workflows/release-pypi-pr.yml b/.github/workflows/release-pypi-pr.yml index edcca5c06..627f3750c 100644 --- a/.github/workflows/release-pypi-pr.yml +++ b/.github/workflows/release-pypi-pr.yml @@ -17,7 +17,8 @@ jobs: if: github.repository == 'sgl-project/sglang' runs-on: ubuntu-latest outputs: - wheel_version: ${{ steps.gen_version.outputs.wheel_version }} + wheel_version: ${{ steps.wheel.outputs.wheel_version }} + wheel_filename: ${{ steps.wheel.outputs.wheel_filename }} commit_hash: ${{ steps.gen_version.outputs.commit_hash }} build_date: ${{ steps.gen_version.outputs.build_date }} steps: @@ -34,7 +35,7 @@ jobs: - name: Generate PR wheel version id: gen_version run: | - LATEST_TAG=$(python3 scripts/release/get_version_tag.py) + LATEST_TAG=$(python3 scripts/release/get_version_tag.py --tag-only) BASE_VERSION=${LATEST_TAG#v} echo "Latest release tag: ${LATEST_TAG}" @@ -57,10 +58,12 @@ jobs: echo "Commit: ${COMMIT_HASH}" echo "Build date: ${BUILD_DATE}" - echo "wheel_version=${WHEEL_VERSION}" >> $GITHUB_OUTPUT - echo "commit_hash=${COMMIT_HASH}" >> $GITHUB_OUTPUT - echo "base_version=${BASE_VERSION}" >> $GITHUB_OUTPUT - echo "build_date=${BUILD_DATE}" >> $GITHUB_OUTPUT + { + echo "wheel_version=${WHEEL_VERSION}" + echo "commit_hash=${COMMIT_HASH}" + echo "base_version=${BASE_VERSION}" + echo "build_date=${BUILD_DATE}" + } >> "$GITHUB_OUTPUT" - name: Update pyproject.toml with PR wheel version run: | @@ -79,19 +82,26 @@ jobs: - name: Install build dependencies run: | - cd python - pip install build wheel setuptools + pip install \ + auditwheel build patchelf wheel "setuptools>=61.0" \ + "setuptools-rust>=1.11" "setuptools-scm>=8.0" "torch==2.13.0" - name: Build wheel run: | cd python cp ../README.md ../LICENSE . - python3 -m build --wheel + python3 -m build --wheel --no-isolation # List built wheels echo "Built wheel:" ls -lh dist/ + - name: Repair and smoke-test wheel + id: wheel + run: | + python3 scripts/release/prepare_sglang_wheel.py python/dist \ + --github-output "$GITHUB_OUTPUT" + - name: Upload wheel artifact uses: actions/upload-artifact@v4 with: @@ -127,7 +137,7 @@ jobs: prerelease: true body: | PR wheel build from PR #${{ inputs.pr_number }} - Commit: ${{ github.sha }} + Commit: ${{ needs.build-pr-wheel.outputs.commit_hash }} Build date: ${{ needs.build-pr-wheel.outputs.build_date }} Version: ${{ needs.build-pr-wheel.outputs.wheel_version }} @@ -143,14 +153,14 @@ jobs: **Direct installation:** ```bash - pip install https://github.com/sgl-project/whl/releases/download/pr-${{ inputs.pr_number }}-${{ needs.build-pr-wheel.outputs.build_date }}-${{ needs.build-pr-wheel.outputs.commit_hash }}/sglang-${{ needs.build-pr-wheel.outputs.wheel_version }}-py3-none-any.whl + pip install https://github.com/sgl-project/whl/releases/download/pr-${{ inputs.pr_number }}-${{ needs.build-pr-wheel.outputs.build_date }}-${{ needs.build-pr-wheel.outputs.commit_hash }}/${{ needs.build-pr-wheel.outputs.wheel_filename }} ``` files: | dist/*.whl - name: Clone wheel index repository run: | - git clone https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git sgl-whl + git clone "https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git" sgl-whl cd sgl-whl git config --local user.name "sglang-bot" git config --local user.email "sglangbot@gmail.com" @@ -165,10 +175,10 @@ jobs: - name: Update wheel index run: | python3 scripts/update_pr_whl_index.py \ - --pr-number ${{ inputs.pr_number }} \ - --commit-hash ${{ needs.build-pr-wheel.outputs.commit_hash }} \ - --wheel-version ${{ needs.build-pr-wheel.outputs.wheel_version }} \ - --build-date ${{ needs.build-pr-wheel.outputs.build_date }} + --pr-number "${{ inputs.pr_number }}" \ + --commit-hash "${{ needs.build-pr-wheel.outputs.commit_hash }}" \ + --wheel-version "${{ needs.build-pr-wheel.outputs.wheel_version }}" \ + --build-date "${{ needs.build-pr-wheel.outputs.build_date }}" - name: Push wheel index run: | diff --git a/.github/workflows/release-pypi.yml b/.github/workflows/release-pypi.yml index 6fffb2f6b..591bcde68 100644 --- a/.github/workflows/release-pypi.yml +++ b/.github/workflows/release-pypi.yml @@ -61,25 +61,23 @@ jobs: run: | cd python cp ../README.md ../LICENSE . - pip install build wheel setuptools setuptools-scm setuptools-rust + pip install \ + build wheel "setuptools>=61.0" "setuptools-rust>=1.11" \ + "setuptools-scm>=8.0" "torch==2.13.0" if [ -n "$RELEASE_VERSION" ]; then export SETUPTOOLS_SCM_PRETEND_VERSION="${RELEASE_VERSION#v}" echo "Pinning wheel version to $SETUPTOOLS_SCM_PRETEND_VERSION" fi - python3 -m build --wheel + python3 -m build --wheel --no-isolation # PyPI rejects plain `linux_x86_64` / `linux_aarch64` platform tags; # auditwheel rewrites the wheel's platform tag to a `manylinux_*` tag # and bundles any external native deps. The runner's glibc determines # the lowest acceptable manylinux policy. - - name: Repair wheel for manylinux + - name: Repair and smoke-test wheel run: | - cd python pip install auditwheel patchelf - mkdir -p dist-repaired - python3 -m auditwheel repair dist/*.whl -w dist-repaired/ - rm dist/*.whl - mv dist-repaired/*.whl dist/ + python3 scripts/release/prepare_sglang_wheel.py python/dist - name: Upload artifacts uses: actions/upload-artifact@v4 diff --git a/.github/workflows/seed-rust-ext-cache.yml b/.github/workflows/seed-rust-ext-cache.yml index 6ca4fb90d..67ce07453 100644 --- a/.github/workflows/seed-rust-ext-cache.yml +++ b/.github/workflows/seed-rust-ext-cache.yml @@ -10,8 +10,11 @@ on: paths: - 'rust/**' - 'python/setup.py' + - 'python/pyproject.toml' + - 'python/sglang/srt/rust_extensions/torch_build.py' # The key's prefix lives in this file, so a bump there also moves the key. - '.github/workflows/_pr-test-rust-ext-build.yml' + - 'scripts/ci/utils/stage_rust_ext_modules.sh' workflow_dispatch: # Only the newest merge needs to seed; earlier ones are already stale. diff --git a/.gitignore b/.gitignore index cd1c6325d..d0dcf786d 100644 --- a/.gitignore +++ b/.gitignore @@ -231,6 +231,7 @@ work_dirs/ # Rust lib Cargo.lock !rust/Cargo.lock +!rust/mem-cache/Cargo.lock # Generated vision test fixtures (regenerate with: python scripts/generate_vision_golden.py) sgl-model-gateway/tests/fixtures/golden/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index f7a14d03e..eba95baa0 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -151,13 +151,13 @@ repos: # cover sglang-mm's PyO3 bindings + rayon fan-out — both sit behind # non-default features, so `--workspace` alone never compiles them. protoc # is not required: sglang-grpc's build.rs falls back to a vendored binary. - entry: bash -c 'cd rust && cargo clippy --workspace --fix --allow-dirty --allow-staged && cargo clippy --workspace -- -D warnings && cargo clippy -p sglang-mm --features python,parallel --lib -- -D warnings' + entry: bash -c 'cd rust && cargo clippy --workspace --fix --allow-dirty --allow-staged && cargo clippy --workspace -- -D warnings && cargo clippy -p sglang-mm --features python,parallel --lib -- -D warnings && cargo clippy --manifest-path mem-cache/Cargo.toml --all-targets --no-default-features --features tch/doc-only -- -D warnings' language: system files: ^rust/.*\.rs$ pass_filenames: false - id: rustfmt-rust-workspace name: rustfmt rust/ workspace - entry: bash -c 'cd rust && cargo fmt' + entry: bash -c 'cd rust && cargo fmt && cargo fmt --manifest-path mem-cache/Cargo.toml' language: system files: ^rust/.*\.rs$ pass_filenames: false diff --git a/docker/Dockerfile b/docker/Dockerfile index d7d9a07a6..ceedf8dae 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -671,6 +671,8 @@ RUN --mount=type=cache,target=/root/.cache/pip \ && ( if [ -f python/kernels.lock ]; then mv python/kernels.lock /root/.cache/sglang/; fi ) \ && ( find /opt/sglang/lib/python3.12/site-packages -type d -name "__pycache__" -exec rm -rf {} + 2>/dev/null || true ) +ENV SGLANG_RUST_BUILD_MODE=never + # Install pre-built gateway artifacts from parallel builder COPY --from=gateway_builder /build/sgl-model-gateway-bin /opt/sglang/bin/sgl-model-gateway @@ -723,7 +725,8 @@ ARG GDRCOPY_VERSION=2.5.1 ENV DEBIAN_FRONTEND=noninteractive \ CUDA_HOME=/usr/local/cuda \ - GDRCOPY_HOME=/usr/src/gdrdrv-${GDRCOPY_VERSION}/ + GDRCOPY_HOME=/usr/src/gdrdrv-${GDRCOPY_VERSION}/ \ + SGLANG_RUST_BUILD_MODE=never # Add GKE default lib and bin locations + CUDA compiler paths for FlashInfer JIT ENV PATH="${PATH}:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/cuda/nvvm/bin" \ diff --git a/docker/Dockerfile.cu134 b/docker/Dockerfile.cu134 index 38fd5e653..d6d1f9a00 100644 --- a/docker/Dockerfile.cu134 +++ b/docker/Dockerfile.cu134 @@ -944,6 +944,9 @@ ARG CUDA_VERSION ARG SGL_VERSION ARG USE_LATEST_SGLANG +# TODO(Jialin): Set SGLANG_RUST_BUILD_MODE=never after Rust TreeCore supports +# the PyTorch 2.15 nightly used by this preview image. + WORKDIR /sgl-workspace COPY --from=local_src /src /tmp/local_src diff --git a/python/pyproject.toml b/python/pyproject.toml index 95fb5a935..b82df0fb2 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -1,5 +1,11 @@ [build-system] -requires = ["setuptools>=61.0", "setuptools-rust>=1.10", "setuptools-scm>=8.0", "wheel"] +requires = [ + "setuptools>=61.0", + "setuptools-rust>=1.11", + "setuptools-scm>=8.0", + "torch==2.13.0", + "wheel", +] build-backend = "setuptools.build_meta" [project] @@ -214,6 +220,7 @@ killall_sglang = "sglang.cli.killall:main" "sglang" = [ "kernels/aot/*", "kernels/aot/**/*", + "srt/mem_cache/rust_tree_core/mem_cache_inspection*.so", ] [tool.setuptools.packages.find] @@ -247,8 +254,9 @@ git_describe_command = ["python3", "scripts/release/get_version_tag.py"] # Allow editable installs even when .git metadata is not available. fallback_version = "0.0.0.dev0" -# Rust extension modules are auto-discovered by setup.py from the cargo -# workspace in ../rust ([package.metadata.sglang] python-module in each crate). +# Rust extension modules are auto-discovered by setup.py from the Cargo +# workspace in ../rust and its declared extension manifests +# ([package.metadata.sglang] python-module in each crate). # This CUDA pyproject builds all of them; platform variants restrict the set # via [tool.sglang] rust-extensions (see pyproject_other.toml). diff --git a/python/setup.py b/python/setup.py index f9676f358..96b489d95 100644 --- a/python/setup.py +++ b/python/setup.py @@ -1,7 +1,8 @@ """sglang build hooks. -Rust extensions are auto-discovered from the cargo workspace in ../rust: every -crate whose Cargo.toml declares +Rust extensions are auto-discovered from the Cargo workspace in ../rust and +the extension manifests declared by its workspace metadata. Every crate whose +Cargo.toml declares [package.metadata.sglang] python-module = "sglang.srt.._core" # import path inside the wheel @@ -28,6 +29,7 @@ Two filters can narrow the discovered set: import json import os import re +import runpy import subprocess from pathlib import Path @@ -45,14 +47,17 @@ except ModuleNotFoundError as exc: _BUILD_RUST_EXTS_ENV = "SGLANG_BUILD_RUST_EXTS" _PYTHON_DIR = Path(__file__).resolve().parent _RUST_WORKSPACE_DIR = _PYTHON_DIR.parent / "rust" +_RUST_BUILD_HELPERS = runpy.run_path( + os.fspath(_PYTHON_DIR / "sglang" / "srt" / "rust_extensions" / "torch_build.py") +) +_torch_build_configuration = _RUST_BUILD_HELPERS["torch_build_configuration"] -def _cargo_workspace_metadata(): - """The rust/ cargo workspace as JSON, straight from cargo's own parser.""" - manifest_path = _RUST_WORKSPACE_DIR / "Cargo.toml" +def _cargo_metadata(manifest_path): + """One Cargo workspace/package manifest as Cargo's own JSON metadata.""" if not manifest_path.is_file(): raise RuntimeError( - f"no cargo workspace at {manifest_path} (building outside a repo " + f"no Cargo manifest at {manifest_path} (building outside a repo " f"checkout?); set {_BUILD_RUST_EXTS_ENV}=none to build without " "Rust extensions" ) @@ -83,6 +88,27 @@ def _cargo_workspace_metadata(): return json.loads(out.stdout) +def _cargo_workspace_metadata(): + """Root workspace metadata plus explicitly declared extension workspaces.""" + root_manifest = _RUST_WORKSPACE_DIR / "Cargo.toml" + document = _cargo_metadata(root_manifest) + external_manifests = ( + (document.get("metadata") or {}) + .get("sglang", {}) + .get("extension-manifests", []) + ) + packages = list(document["packages"]) + for relative_manifest in external_manifests: + external = (_RUST_WORKSPACE_DIR / relative_manifest).resolve() + if _RUST_WORKSPACE_DIR not in external.parents: + raise RuntimeError( + f"external Rust extension manifest escapes rust/: {relative_manifest}" + ) + packages.extend(_cargo_metadata(external)["packages"]) + document["packages"] = packages + return document + + def _match_by_substring(declared, tokens, source): """Match tokens as case-insensitive substrings of extension names.""" matched = set() @@ -111,17 +137,22 @@ def _discovered_rust_extensions(): sglang_meta = (package["metadata"] or {}).get("sglang", {}) if "python-module" not in sglang_meta: continue - extensions.append( - RustExtension( - target=sglang_meta["python-module"], - path=package["manifest_path"], - binding=Binding.PyO3, - debug=sglang_meta.get("debug"), - # Crates that gate their PyO3 bindings behind a non-default - # feature (so the pure-Rust core stays pyo3-free) declare it here. - features=sglang_meta.get("features"), - ) + extension = RustExtension( + target=sglang_meta["python-module"], + path=package["manifest_path"], + binding=Binding.PyO3, + debug=sglang_meta.get("debug"), + # Crates that gate their PyO3 bindings behind a non-default + # feature (so the pure-Rust core stays pyo3-free) declare it here. + features=sglang_meta.get("features"), + cargo_manifest_args=["--locked"], ) + # Preserve Cargo metadata until the selected extension is actually + # built. Alternate platform pyprojects filter mem-cache out before + # this point and therefore do not need torch as a build dependency. + extension._sglang_metadata = sglang_meta + extension._sglang_manifest_path = package["manifest_path"] + extensions.append(extension) if not extensions: raise RuntimeError( f"no crate under {_RUST_WORKSPACE_DIR} declares " @@ -188,6 +219,19 @@ if build_rust is not None: class BuildRust(build_rust): """Build only the Rust extensions selected by SGLANG_BUILD_RUST_EXTS.""" + def run_for_extension(self, extension) -> None: + metadata = extension._sglang_metadata + compat_header = metadata.get("torch-compat-header") + if compat_header is not None: + manifest = Path(extension._sglang_manifest_path) + build = _torch_build_configuration( + compat_header=manifest.parent / compat_header, + python_module=extension.name, + include_absolute_rpath=False, + ) + extension.env.env = build.environment + super().run_for_extension(extension) + def run(self) -> None: rust_extensions = _selected_rust_extensions(self.extensions or []) self.extensions = rust_extensions diff --git a/python/sglang/srt/arg_groups/pd_disaggregation_hook.py b/python/sglang/srt/arg_groups/pd_disaggregation_hook.py index 78bf2f755..7e05b9176 100644 --- a/python/sglang/srt/arg_groups/pd_disaggregation_hook.py +++ b/python/sglang/srt/arg_groups/pd_disaggregation_hook.py @@ -22,6 +22,7 @@ logger = logging.getLogger(__name__) def handle_pd_disaggregation(server_args: ServerArgs) -> None: """Validate and normalize PD-disaggregation server args.""" cfg = resolving_view(server_args) + # "mooncake_tcp" is mooncake with the TCP transport forced: set MC_FORCE_TCP # so mooncake installs TcpTransport instead of RDMA, rewrite the backend to # mooncake, and skip RDMA HCA selection. Must run before backend-name checks. diff --git a/python/sglang/srt/disaggregation/decode_hicache_mixin.py b/python/sglang/srt/disaggregation/decode_hicache_mixin.py index fb39cf3e4..5827e3b99 100644 --- a/python/sglang/srt/disaggregation/decode_hicache_mixin.py +++ b/python/sglang/srt/disaggregation/decode_hicache_mixin.py @@ -71,13 +71,15 @@ class DecodeHiCachePreallocMixin: l3_storage_hit_length = 0 last_host_node = None if self.scheduler.enable_decode_hicache: - last_host_node = self.tree_cache.resolve_node_handle(result.last_host_node) - if last_host_node.backuped or last_host_node is self.tree_cache.root_node: + last_host_node = result.last_host_node + if self.tree_cache.is_backuped(last_host_node) or self.tree_cache.is_root( + last_host_node + ): matched_len = l1_prefix_len + l2_host_hit_length suffix_tokens = req.origin_input_ids[matched_len:] - last_hash = last_host_node.get_last_hash_value() + last_hash = self.tree_cache.get_last_hash_value(last_host_node) prefix_keys = ( - last_host_node.get_prefix_hash_values(last_host_node.parent) + self.tree_cache.get_prefix_hash_values(last_host_node) if self.tree_cache.hicache_storage_pass_prefix_keys else None ) @@ -112,14 +114,13 @@ class DecodeHiCachePreallocMixin: ): return try: - node = self.tree_cache.resolve_node_handle(prefix_match.last_host_node) matched_len = prefix_match.l1_prefix_len + prefix_match.l2_host_hit_length suffix = req.origin_input_ids[ matched_len : matched_len + prefix_match.l3_storage_hit_length ] - last_hash = node.get_last_hash_value() + last_hash = self.tree_cache.get_last_hash_value(prefix_match.last_host_node) prefix_keys = ( - node.get_prefix_hash_values(node.parent) + self.tree_cache.get_prefix_hash_values(prefix_match.last_host_node) if self.tree_cache.hicache_storage_pass_prefix_keys else None ) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 2c4286d39..5bf710ad0 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -31,7 +31,7 @@ logger = logging.getLogger(__name__) import os import random -from collections import Counter, defaultdict +from collections import Counter from contextlib import contextmanager from enum import Enum, auto from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union @@ -397,23 +397,8 @@ class SchedulePolicy: waiting_queue: List[Req], tree_cache: BasePrefixCache ) -> None: """Sorts the waiting queue based on a depth-first search weighting.""" - last_node_to_reqs = defaultdict(list) - for req in waiting_queue: - last_node = tree_cache.resolve_node_handle(req.last_node) - last_node_to_reqs[last_node].append(req) - - node_to_weight = defaultdict(int) - for node in last_node_to_reqs: - node_to_weight[node] = len(last_node_to_reqs[node]) - SchedulePolicy._calc_weight(tree_cache.root_node, node_to_weight) - - waiting_queue.clear() - SchedulePolicy._get_dfs_priority( - tree_cache.root_node, - node_to_weight, - last_node_to_reqs, - waiting_queue, - ) + order = tree_cache.dfs_weight_order([req.last_node for req in waiting_queue]) + waiting_queue[:] = [waiting_queue[index] for index in order] @staticmethod def _sort_by_longest_output( @@ -482,27 +467,6 @@ class SchedulePolicy: waiting_keys_after = [r.routing_key for r in waiting_queue] logger.info(f"waiting_keys_after={waiting_keys_after}") - @staticmethod - def _calc_weight(cur_node: TreeNode, node_to_weight: Dict[TreeNode, int]) -> None: - for child in cur_node.children.values(): - SchedulePolicy._calc_weight(child, node_to_weight) - node_to_weight[cur_node] += node_to_weight[child] - - @staticmethod - def _get_dfs_priority( - cur_node: TreeNode, - node_to_priority: Dict[TreeNode, int], - last_node_to_reqs: Dict[TreeNode, List[Req]], - q: List, - ) -> None: - children = [child for child in cur_node.children.values()] - children.sort(key=lambda x: -node_to_priority[x]) - for child in children: - SchedulePolicy._get_dfs_priority( - child, node_to_priority, last_node_to_reqs, q - ) - q.extend(last_node_to_reqs[cur_node]) - class AddReqResult(Enum): CONTINUE = auto() # Continue to add requests diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 2dabaa1b2..8338aeffd 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -6,6 +6,7 @@ from abc import ABC, abstractmethod from typing import ( TYPE_CHECKING, Any, + Callable, NamedTuple, Optional, Protocol, @@ -244,6 +245,42 @@ def zero_match_result( ) +def _dfs_weight_order( + root_node: Any, + node_handles: Sequence[Any], + resolve_node_handle: Callable[[Any], Any], +) -> list[int]: + last_node_to_indices: dict[Any, list[int]] = {} + for index, node_handle in enumerate(node_handles): + node = resolve_node_handle(node_handle) + last_node_to_indices.setdefault(node, []).append(index) + + node_to_weight: dict[Any, int] = { + node: len(indices) for node, indices in last_node_to_indices.items() + } + + def calc_weight(node: Any) -> None: + for child in node.children.values(): + calc_weight(child) + node_to_weight[node] = node_to_weight.get(node, 0) + node_to_weight.get( + child, 0 + ) + + calc_weight(root_node) + + order: list[int] = [] + + def append_dfs(node: Any) -> None: + children = list(node.children.values()) + children.sort(key=lambda child: -node_to_weight.get(child, 0)) + for child in children: + append_dfs(child) + order.extend(last_node_to_indices.get(node, ())) + + append_dfs(root_node) + return order + + class BasePrefixCache(ABC, PrefixCacheTrait): """Cache can be indexed by either rid or key.""" @@ -289,6 +326,10 @@ class BasePrefixCache(ABC, PrefixCacheTrait): def supports_fast_match_prefix(self) -> bool: return False + def dfs_weight_order(self, node_handles: Sequence[Any]) -> list[int]: + """Return request indices in depth-first, subtree-weight order.""" + return _dfs_weight_order(self.root_node, node_handles, self.resolve_node_handle) + def resolve_node_handle(self, node_handle: Any) -> Any: """Map a node handle to its node -- e.g. UnifiedRadixCache looks up the node object from its NodeId. Temporary API for the Unified Radix Cache diff --git a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py b/python/sglang/srt/mem_cache/buffer_mode/pipeline.py index efa532772..cf63442f9 100644 --- a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py +++ b/python/sglang/srt/mem_cache/buffer_mode/pipeline.py @@ -46,12 +46,13 @@ from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName, PoolTr from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.unified_cache.cache_action import RebuildFullToSWAMapping from sglang.srt.mem_cache.unified_cache.components import ( - BASE_COMPONENT_TYPE, + CacheTransferPhase, ComponentType, ) -from sglang.srt.mem_cache.unified_cache.unified_tree_core import ( +from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import ( + BufferBackupSnapshot, + BufferBackupState, NodeId, - UnifiedTreeNode, ) if TYPE_CHECKING: @@ -65,16 +66,12 @@ class _UnifiedBackupIntent(msgspec.Struct): """Buffer-mode backup intent, unpinned while queued. Snapshots node identity at enqueue time: a split rewrites the node's - key/hash in place while these copies stay intact, so - ``node.hash_value != hash_values`` doubles as split detection and a None - FULL device value as eviction detection (``_backup_intent_stale``). + key/hash in place while these copies stay intact, so a key-length change + detects a split and a missing FULL device value detects eviction + (``_validate_backup_intent``). """ - node: UnifiedTreeNode - node_id: int - hash_values: list[str] - key: RadixKey - prefix_keys: Optional[list[str]] = None + snapshot: BufferBackupSnapshot class _UnifiedBufferBackupEntry(msgspec.Struct): @@ -289,18 +286,19 @@ class BufferModePipeline: # ---- backup pipeline (device -> staging -> storage) ---- - def _backup_parent_covered(self, node: UnifiedTreeNode) -> bool: + def _backup_parent_covered(self, state: BufferBackupState) -> bool: """Only admit a node whose parent is stored/in-flight: writing above a dropped parent creates a permanent longest-prefix hole.""" - parent = node.parent if ( - parent is self._cache.root_node - or parent.id in self.inflight_backup_node_ids + state.parent_is_root + or state.parent_node_id in self.inflight_backup_node_ids ): return True - last_hash = parent.get_last_hash_value() - return last_hash is not None and self._cache.storage_existence_cache.contains( - PoolName.KV, last_hash + return ( + state.parent_last_hash is not None + and self._cache.storage_existence_cache.contains( + PoolName.KV, state.parent_last_hash + ) ) def _log_backup_dropped(self, num_tokens: int) -> None: @@ -308,22 +306,29 @@ class BufferModePipeline: if cache.enable_storage_metrics and cache.storage_metrics_collector is not None: cache.storage_metrics_collector.log_backup_dropped_tokens(num_tokens) - def enqueue_backup_intent(self, node: UnifiedTreeNode) -> None: + def enqueue_backup_intent(self, node_id: NodeId) -> None: """Snapshot a backup intent and commit it to the write queue. Admission gates: belief skip, parent-cover, backlog cap, oversize. Drops are silent; the node re-triggers on a later hit.""" - if not self._cache.enable_storage or not node.hash_value: + if not self._cache.enable_storage: return - if node.id in self.inflight_backup_node_ids: + if node_id in self.inflight_backup_node_ids: + return + snapshot = self._cache.tree_core.snapshot_buffer_backup( + node_id, self._cache.hicache_storage_pass_prefix_keys + ) + if snapshot is None: return # Admission cover: beliefs plus content past its D2H launch. The # launched cover keeps republished content (fill inserts under new # node ids) from re-writing while the original write drains. if self._cache.storage_existence_cache.covers_all( - PoolName.KV, node.hash_value, extra_cover=self.inflight_backup_hashes + PoolName.KV, + snapshot.hash_values, + extra_cover=self.inflight_backup_hashes, ): return - intent_tokens = len(node.hash_value) * self._cache.page_size + intent_tokens = len(snapshot.hash_values) * self._cache.page_size if self.write_backlog_tokens_ >= self.write_backlog_cap: # The cap sits at 2x the intrinsic live-backlog ceiling (see # init_hicache), so reaching it means leaked accounting or a @@ -344,51 +349,59 @@ class BufferModePipeline: return # A span larger than any pool's whole staging capacity can never # stage; admitting it would wedge the head-of-line queue forever. - if not self._backup_parent_covered(node) or self._backup_oversize( - node, intent_tokens + state = BufferBackupState( + parent_node_id=snapshot.parent_node_id, + parent_is_root=snapshot.parent_is_root, + parent_last_hash=snapshot.parent_last_hash, + ) + if not self._backup_parent_covered(state) or self._backup_oversize( + snapshot.node_id, snapshot.hash_values, intent_tokens ): self._log_backup_dropped(intent_tokens) return - prefix_keys = ( - node.get_prefix_hash_values(node.parent) - if self._cache.hicache_storage_pass_prefix_keys - else None - ) - intent = _UnifiedBackupIntent( - node=node, - node_id=node.id, - hash_values=list(node.hash_value), - key=node.key, - prefix_keys=prefix_keys, - ) + intent = _UnifiedBackupIntent(snapshot=snapshot) self.pending_write_queue.append(intent) - self.inflight_backup_node_ids.add(node.id) + self.inflight_backup_node_ids.add(snapshot.node_id) self.write_backlog_tokens_ += intent_tokens def _build_aux_staging_transfers( - self, node: UnifiedTreeNode - ) -> Optional[list[PoolTransfer]]: + self, + node_id: NodeId, + hash_values: list[str], + comp_xfers: Optional[dict[ComponentType, list[PoolTransfer]]] = None, + ) -> list[PoolTransfer]: """Keys-only aux transfers mirroring what BACKUP_STORAGE would write; sizes the per-pool oversize gate (beliefs do not consult these).""" transfers: list[PoolTransfer] = [] if ComponentType.SWA in self._cache.components: - cd = node.component_data[ComponentType.SWA] - if cd.value is not None: - num_pages = len(cd.value) // self._cache.page_size + current = ( + comp_xfers.get(ComponentType.SWA) + if comp_xfers is not None + else self._cache.tree_core.build_hicache_transfers( + ComponentType.SWA, + node_id, + CacheTransferPhase.BACKUP_HOST, + ) + ) + for transfer in current or (): + if transfer.device_indices is None: + continue + num_pages = len(transfer.device_indices) // self._cache.page_size if num_pages > 0: transfers.append( PoolTransfer( name=PoolName.SWA, - keys=node.hash_value[-num_pages:], + keys=hash_values[-num_pages:], hit_policy=PoolHitPolicy.TRAILING_PAGES, ) ) - return transfers or None + return transfers def _backup_oversize( self, - node: UnifiedTreeNode, + node_id: NodeId, + hash_values: list[str], intent_tokens: int, aux_xfers: Optional[list[PoolTransfer]] = None, ) -> bool: @@ -400,7 +413,7 @@ class BufferModePipeline: if intent_tokens > cc.mem_pool_host.size: return True if aux_xfers is None: - aux_xfers = self._build_aux_staging_transfers(node) + aux_xfers = self._build_aux_staging_transfers(node_id, hash_values) for t in aux_xfers or (): entry = cc.mem_pool_host.entry_map.get(t.name) if entry is not None and ( @@ -420,35 +433,37 @@ class BufferModePipeline: host_pool.size // 10, ) - def _backup_intent_stale(self, intent: _UnifiedBackupIntent) -> bool: - # Arena-lookup failure = deleted, hash mismatch vs the enqueue-time - # snapshot = split, a None FULL device value = evicted. Stale + def _validate_backup_intent( + self, intent: _UnifiedBackupIntent + ) -> Optional[BufferBackupState]: + # Arena-lookup failure = deleted, key-length mismatch vs the snapshot + # = split, a None FULL device value = evicted. Stale # intents drop silently; the node re-triggers on a later hit. - node = intent.node - try: - self._cache.tree_core.node_by_id(intent.node_id) - except KeyError: - return True - return ( - node.component_data[BASE_COMPONENT_TYPE].value is None - or node.hash_value != intent.hash_values + snapshot = intent.snapshot + return self._cache.tree_core.validate_buffer_backup( + snapshot.node_id, len(snapshot.key) ) - def _sweep_stale_backup_intents(self) -> None: + def _sweep_stale_backup_intents(self) -> dict[NodeId, BufferBackupState]: """Cancel stale intents anywhere in the queue, not just at the head: a dead intent would otherwise inflate the backlog accounting and hold FIFO position ahead of live segments.""" if not self.pending_write_queue: - return + return {} page_size = self._cache.page_size survivors: deque[_UnifiedBackupIntent] = deque() + states: dict[NodeId, BufferBackupState] = {} for intent in self.pending_write_queue: - if self._backup_intent_stale(intent): - self.inflight_backup_node_ids.discard(intent.node_id) - self.write_backlog_tokens_ -= len(intent.hash_values) * page_size + snapshot = intent.snapshot + state = self._validate_backup_intent(intent) + if state is None: + self.inflight_backup_node_ids.discard(snapshot.node_id) + self.write_backlog_tokens_ -= len(snapshot.hash_values) * page_size continue survivors.append(intent) + states[snapshot.node_id] = state self.pending_write_queue = survivors + return states def flush_pending_writes(self) -> None: """Launch D2H transfers for admitted intents, head-of-line: device @@ -456,7 +471,7 @@ class BufferModePipeline: if not self.pending_write_queue: return cc = self._cache.cache_controller - self._sweep_stale_backup_intents() + states = self._sweep_stale_backup_intents() # Loads have priority (writes are deferrable): the write window is # the pool minus prefetch occupancy minus a 10% margin, floored at # the configured fraction. @@ -467,34 +482,56 @@ class BufferModePipeline: ) while self.pending_write_queue: intent = self.pending_write_queue[0] - intent_tokens = len(intent.hash_values) * self._cache.page_size - if not self._backup_parent_covered(intent.node) or self._backup_oversize( - intent.node, intent_tokens - ): - # Unwritable intent (dropped parent or unstageable size): - # cascade the drop down the chain rather than creating a - # permanent storage hole / stalling the head-of-line queue. + snapshot = intent.snapshot + state = states[snapshot.node_id] + intent_tokens = len(snapshot.hash_values) * self._cache.page_size + if not self._backup_parent_covered(state): + # Cascade a dropped parent down the chain rather than creating + # a permanent storage hole. self.pending_write_queue.popleft() - self.inflight_backup_node_ids.discard(intent.node_id) + self.inflight_backup_node_ids.discard(snapshot.node_id) self.write_backlog_tokens_ -= intent_tokens self._log_backup_dropped(intent_tokens) continue if self.write_staged_tokens_ >= live_cap: # Yield to live fetch demand; retry next round. break - if self._aux_budget_blocked(intent): + device_value, comp_xfers = self._cache.tree_core.build_backup_spec( + snapshot.node_id + ) + sizing_xfers = self._build_aux_staging_transfers( + snapshot.node_id, snapshot.hash_values, comp_xfers + ) + if self._backup_oversize( + snapshot.node_id, + snapshot.hash_values, + intent_tokens, + sizing_xfers, + ): + # A permanently unstageable head must not block the queue. + self.pending_write_queue.popleft() + self.inflight_backup_node_ids.discard(snapshot.node_id) + self.write_backlog_tokens_ -= intent_tokens + self._log_backup_dropped(intent_tokens) + continue + if self._aux_budget_blocked(intent, sizing_xfers): # An aux pool lacks staging headroom: yield at the gate # instead of failing the alloc inside cc.write; acks free # aux staging, retry next round. break - if not self._launch_backup_intent(intent): + if not self._launch_backup_intent(intent, device_value, comp_xfers): # Pool full of in-flight staging and nothing reclaimable # (the tree never holds host values in buffer mode): # defer, head-of-line; pending acks will free slots. break self.pending_write_queue.popleft() - def _launch_backup_intent(self, intent: _UnifiedBackupIntent) -> bool: + def _launch_backup_intent( + self, + intent: _UnifiedBackupIntent, + device_value: torch.Tensor, + comp_xfers: dict[ComponentType, list[PoolTransfer]], + ) -> bool: """Launch one admitted intent's D2H (staging alloc + device lock + async copy); the caller removes it from pending_write_queue. Returns False when staging cannot be allocated. From a successful launch the @@ -502,33 +539,34 @@ class BufferModePipeline: LAUNCHED cover consulted by admission.""" cache = self._cache cc = cache.cache_controller - node = intent.node - # Build aux transfers from the node's CURRENT state: a SWA span - # tombstoned since admission backs up FULL-only, as in cache mode. - device_value, comp_xfers = cache.tree_core.build_backup_spec(node.id) + snapshot = intent.snapshot aux_xfers = [x for xfers in comp_xfers.values() for x in xfers] host_indices = cc.write( device_value, - node_id=node.id, + node_id=snapshot.node_id, extra_pools=aux_xfers or None, ) if host_indices is None: return False - _track_content_refs(self.inflight_backup_hashes, intent.hash_values) + _track_content_refs(self.inflight_backup_hashes, snapshot.hash_values) # NOTE: no commit_backup — the node must never appear # host-resident in buffer mode; staging slots live in the entry. - lock_params = cache.inc_lock_ref(node.id).to_dec_params() - self.ongoing_write_through[node.id] = _UnifiedBufferBackupEntry( + lock_params = cache.inc_lock_ref(snapshot.node_id).to_dec_params() + self.ongoing_write_through[snapshot.node_id] = _UnifiedBufferBackupEntry( intent=intent, host_indices=host_indices, aux_xfers=aux_xfers, lock_params=lock_params, ) self.write_staged_tokens_ += len(host_indices) - self.write_backlog_tokens_ -= len(intent.hash_values) * cache.page_size + self.write_backlog_tokens_ -= len(snapshot.hash_values) * cache.page_size return True - def _aux_budget_blocked(self, intent: _UnifiedBackupIntent) -> bool: + def _aux_budget_blocked( + self, + intent: _UnifiedBackupIntent, + aux: Optional[list[PoolTransfer]] = None, + ) -> bool: """True when an aux pool cannot stage this intent right now (free minus the loads-priority margin falls short of the need): defer at the gate instead of failing the alloc inside cc.write and blocking @@ -536,7 +574,11 @@ class BufferModePipeline: loads-have-priority on aux pools the way live_cap does on the KV pool; avail already reflects prefetch-held slots, so no occupancy subtraction here.""" - aux = self._build_aux_staging_transfers(intent.node) + snapshot = intent.snapshot + if aux is None: + aux = self._build_aux_staging_transfers( + snapshot.node_id, snapshot.hash_values + ) if not aux: return False cc = self._cache.cache_controller @@ -574,14 +616,15 @@ class BufferModePipeline: (which reads from the staging copy, so device eviction may proceed).""" entry = self.ongoing_write_through.pop(ack_id) intent = entry.intent - self._cache.dec_lock_ref(intent.node_id, entry.lock_params) + snapshot = intent.snapshot + self._cache.dec_lock_ref(snapshot.node_id, entry.lock_params) # Every aux pool writes a trailing snapshot keyed by the last KV page # hashes it covers: the SWA window spans page_size-sized pages, the # Mamba state is a single slot (host pool page_size 1 -> one key). storage_xfers: list[PoolTransfer] = [] for staged in entry.aux_xfers: - keys = self._aux_window_keys(intent.hash_values, staged) + keys = self._aux_window_keys(snapshot.hash_values, staged) if keys is None: continue storage_xfers.append( @@ -594,9 +637,9 @@ class BufferModePipeline: ) operation_id = self._cache.cache_controller.write_storage( entry.host_indices, - intent.key.token_ids, - intent.hash_values, - intent.prefix_keys, + snapshot.key.token_ids, + snapshot.hash_values, + snapshot.prefix_keys, extra_pools=storage_xfers or None, ) self.ongoing_backup[operation_id] = entry @@ -611,11 +654,12 @@ class BufferModePipeline: if entry is None: return intent = entry.intent - self._cache.storage_existence_cache.add(PoolName.KV, intent.hash_values) + snapshot = intent.snapshot + self._cache.storage_existence_cache.add(PoolName.KV, snapshot.hash_values) self._free_staging_now(entry.host_indices, entry.aux_xfers) self.write_staged_tokens_ -= len(entry.host_indices) - self.inflight_backup_node_ids.discard(entry.intent.node_id) - _untrack_content_refs(self.inflight_backup_hashes, intent.hash_values) + self.inflight_backup_node_ids.discard(snapshot.node_id) + _untrack_content_refs(self.inflight_backup_hashes, snapshot.hash_values) def _free_staging_now( self, host_indices: torch.Tensor, aux_xfers: list[PoolTransfer] @@ -671,10 +715,18 @@ class BufferModePipeline: ) return "cap_skip" cache = self._cache + anchor_tokens = array("q", prefix_tokens) + if cache.tree_core.is_eagle: + # The suffix owns the boundary token shared with the last matched + # bigram, so include it when rebuilding the anchor key. + info = cache.ongoing_prefetch.get(req_id) + if info is None or not info.prefetch_key.token_ids: + return "anchor_lost" + anchor_tokens.append(info.prefetch_key.token_ids[0]) match = cache.match_prefix( MatchPrefixParams( key=RadixKey( - array("q", prefix_tokens), + anchor_tokens, extra_key=extra_key, is_bigram=cache.tree_core.is_eagle, cache_salt=cache_salt, diff --git a/python/sglang/srt/mem_cache/rust_tree_core/.gitignore b/python/sglang/srt/mem_cache/rust_tree_core/.gitignore new file mode 100644 index 000000000..82690f2de --- /dev/null +++ b/python/sglang/srt/mem_cache/rust_tree_core/.gitignore @@ -0,0 +1 @@ +mem_cache.so diff --git a/python/sglang/srt/mem_cache/rust_tree_core/__init__.py b/python/sglang/srt/mem_cache/rust_tree_core/__init__.py new file mode 100644 index 000000000..647854317 --- /dev/null +++ b/python/sglang/srt/mem_cache/rust_tree_core/__init__.py @@ -0,0 +1 @@ +"""The in-tree Rust TreeCore backend; the factory lives in tree_core_registry.""" diff --git a/python/sglang/srt/mem_cache/rust_tree_core/adapter.py b/python/sglang/srt/mem_cache/rust_tree_core/adapter.py new file mode 100644 index 000000000..55f20e39b --- /dev/null +++ b/python/sglang/srt/mem_cache/rust_tree_core/adapter.py @@ -0,0 +1,971 @@ +"""The Rust TreeCore adapter: satisfies ``UnifiedTreeCoreInterface`` over the +``mem_cache`` extension's ``RustUnifiedTreeCoreBinding``.""" + +from __future__ import annotations + +from array import array +from typing import TYPE_CHECKING, Optional, Sequence + +import torch + +from sglang.srt.disaggregation.kv_events import ( + AllBlocksCleared, + BlockRemoved, + BlockStored, + BlockStoredMetadata, + BlockStoredWithMetadata, + StorageMedium, +) +from sglang.srt.mem_cache.base_prefix_cache import ( + DecLockRefParams, + DecLockRefResult, + IncLockRefResult, + InsertParams, + InsertResult, + MatchPrefixParams, + MatchResult, +) +from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName, PoolTransfer +from sglang.srt.mem_cache.radix_cache import RadixKey +from sglang.srt.mem_cache.rust_tree_core.extension import bindings +from sglang.srt.mem_cache.unified_cache.cache_action import ( + BackupKV, + FreeComponentDeviceSlot, + FreeComponentHostSlot, + FreeDeviceKV, + FreeDeviceKVFullOnly, + MambaEvictExcessPathStates, + RebuildFullToSWAMapping, + RecoverSWAWithLockedFull, + ReplaceWriteThroughOnNodeSplit, + SWARebuild, +) +from sglang.srt.mem_cache.unified_cache.component_type import ComponentType +from sglang.srt.mem_cache.unified_cache.unified_tree_core import StorageBackupSpec +from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import ( + BufferBackupSnapshot, + BufferBackupState, + DecSwaLockOnlyResult, + DemoteResult, + DriveHostEvictionResult, + DropSubtreeNoHostResult, + EvictDeviceLeafResult, + EvictDeviceNextNodeResult, + InsertStepResult, + NodeId, + RadixCacheWalkResult, + UnifiedTreeCoreInterface, +) +from sglang.srt.runtime_context import get_exec, mamba_cache_chunk_size + +if TYPE_CHECKING: + from sglang.srt.managers.schedule_batch import Req + from sglang.srt.mem_cache.cache_init_params import CacheInitParams + from sglang.srt.mem_cache.hicache_storage import PoolTransferResult + from sglang.srt.mem_cache.unified_cache.cache_action import ( + CacheAction, + ComponentAction, + ) + from sglang.srt.mem_cache.unified_cache.components import CacheTransferPhase + from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeNode + + +def _radix_key_buffer(key: RadixKey) -> array: + """The key's token ids honoring `limit`; view-independent since the + binding derives its own atoms.""" + token_ids = key.raw_token_ids() + assert ( + isinstance(token_ids, array) and token_ids.typecode == "q" + ), f"tree keys must carry array('q') token ids, got {type(token_ids).__name__}" + return token_ids + + +def _kv_event_from_tagged(event: tuple): + """Build the Python KV cache event for one of the binding's tagged tuples.""" + tag = event[0] + if tag == "block_stored": + event_args = dict( + block_hashes=event[1], + parent_block_hash=event[2], + token_ids=event[3], + block_size=event[4], + lora_id=None, + medium=StorageMedium(event[5]), + ) + if event[6] is None: + return BlockStored(**event_args) + return BlockStoredWithMetadata( + **event_args, + metadata=BlockStoredMetadata(cache_salt=event[6]), + ) + if tag == "block_removed": + return BlockRemoved(block_hashes=event[1], medium=StorageMedium(event[2])) + if tag == "all_blocks_cleared": + return AllBlocksCleared() + raise ValueError(f"unknown kv event tag: {tag}") + + +def _cache_action_from_tagged(action: tuple) -> CacheAction: + """Build the Python CacheAction for one of the binding's tagged tuples.""" + tag = action[0] + if tag == "free_device_kv": + return FreeDeviceKV(indices=list(action[1])) + if tag == "free_device_kv_full_only": + return FreeDeviceKVFullOnly(indices=list(action[1])) + if tag == "backup_kv": + return BackupKV(node_ids=list(action[1])) + if tag == "mamba_evict_excess_path_states": + return MambaEvictExcessPathStates(tail_node_id=action[1]) + if tag == "replace_write_through_on_node_split": + return ReplaceWriteThroughOnNodeSplit( + ack_id=action[1], + old_node_id=action[2], + new_node_id=action[3], + new_child_node_id=action[4], + ) + if tag == "free_component_device_slot": + return FreeComponentDeviceSlot( + component_type=ComponentType(action[1]), indices=list(action[2]) + ) + if tag == "free_component_host_slot": + return FreeComponentHostSlot( + component_type=ComponentType(action[1]), host_indices=list(action[2]) + ) + if tag == "rebuild_full_to_swa_mapping": + return RebuildFullToSWAMapping( + full_indices=list(action[1]), swa_indices=list(action[2]) + ) + if tag == "recover_swa_with_locked_full": + return RecoverSWAWithLockedFull( + node_id=action[1], kept_full=action[2], incoming_full=action[3] + ) + if tag == "swa_rebuild": + return SWARebuild(node_id=action[1], source_value=action[2]) + raise ValueError(f"unknown cache action tag: {tag}") + + +def _cache_actions_from_tagged(actions: Sequence[tuple]) -> list[CacheAction]: + """Build the Python CacheActions for the binding's tagged tuples, in order.""" + return [_cache_action_from_tagged(action) for action in actions] + + +def _inc_lock_ref_result_from_binding(result) -> IncLockRefResult: + return IncLockRefResult( + delta=result.delta, + swa_uuid_for_lock=result.swa_uuid_for_lock, + swa_uuid_for_host_lock=result.swa_uuid_for_host_lock, + skip_lock_node_ids=_skip_lock_node_ids_from_binding(result.skip_lock_node_ids), + ) + + +def _transfer_to_binding(transfer: PoolTransfer) -> tuple: + """The binding's (name, host_indices, device_indices, nodes_to_load, keys, + hit_policy) tuple.""" + return ( + transfer.name.value, + transfer.host_indices, + transfer.device_indices, + transfer.nodes_to_load, + transfer.keys, + transfer.hit_policy.value, + ) + + +def _transfer_from_binding(transfer: tuple) -> PoolTransfer: + """Build the Python PoolTransfer for one of the binding's transfer tuples.""" + name, host_indices, device_indices, nodes_to_load, keys, hit_policy = transfer + return PoolTransfer( + name=PoolName(name), + host_indices=host_indices, + device_indices=device_indices, + keys=keys, + hit_policy=PoolHitPolicy(hit_policy), + nodes_to_load=nodes_to_load, + ) + + +def _comp_xfers_to_binding( + comp_xfers: dict[ComponentType, list[PoolTransfer]], +) -> dict[int, list[tuple]]: + """Rekey per-component transfers by the binding's component values.""" + return { + int(ct): [_transfer_to_binding(x) for x in xfers] + for ct, xfers in comp_xfers.items() + } + + +def _comp_xfers_from_binding( + comp_xfers: dict[int, list[tuple]], +) -> dict[ComponentType, list[PoolTransfer]]: + """Rekey the binding's per-component transfer tuples by ComponentType.""" + return { + ComponentType(ct): [_transfer_from_binding(x) for x in xfers] + for ct, xfers in comp_xfers.items() + } + + +def _insert_step_from_binding(step) -> InsertStepResult: + """Build the interface step for the binding's step (result on the final one).""" + result = None + if step.result is not None: + # A stepped insert delivers all actions through steps, never the result. + assert not step.result.cache_actions + result = InsertResult( + prefix_len=step.result.prefix_len, + last_device_node=step.result.last_device_node, + mamba_exist=step.result.mamba_exist, + host_insert_dropped=step.result.host_insert_dropped, + adopted_ranges=( + { + ComponentType(component_type): list(ranges) + for component_type, ranges in step.result.adopted_ranges.items() + } + if step.result.adopted_ranges is not None + else None + ), + ) + return InsertStepResult( + actions=_cache_actions_from_tagged(step.actions), result=result + ) + + +def _match_result_from_binding(result) -> MatchResult: + """Build the Python MatchResult for the binding's match result.""" + return MatchResult( + device_indices=result.device_indices, + last_device_node=result.last_device_node_id, + last_host_node=result.last_host_node_id, + best_match_node=result.best_match_node_id, + host_hit_length=result.host_hit_length, + swa_host_hit_length=result.swa_host_hit_length, + mamba_host_hit_length=result.mamba_host_hit_length, + mamba_branching_seqlen=result.mamba_branching_seqlen, + full_kv_hit_length=result.full_kv_hit_length, + cache_actions=_cache_actions_from_tagged(result.cache_actions), + ) + + +def _skip_lock_node_ids_from_binding( + skip_lock_node_ids: dict[int, set[int]], +) -> dict[ComponentType, set[int]]: + """Rekey the binding's component-value skip map by ComponentType.""" + return { + ComponentType(component): set(node_ids) + for component, node_ids in skip_lock_node_ids.items() + } + + +def _skip_lock_node_ids_to_binding( + skip_lock_node_ids: dict[ComponentType, set[int]], +) -> dict[int, set[int]]: + """Rekey a ComponentType skip map by the binding's component values.""" + return { + int(component): set(node_ids) + for component, node_ids in skip_lock_node_ids.items() + } + + +def _tracker_to_binding(tracker: dict[ComponentType, int]) -> dict[int, int]: + """Rekey a ComponentType tracker by the binding's component values.""" + return {int(component): freed for component, freed in tracker.items()} + + +def _fill_evict_result(binding_result, result): + """Map a binding eviction step into an interface step result; both carry + this step's per-component deltas and freed tensors.""" + for component, delta in binding_result.tracker.items(): + result.tracker[ComponentType(component)] = delta + for component, tensors in binding_result.new_device_frees.items(): + result.device_frees[ComponentType(component)].extend(tensors) + for component, tensors in binding_result.new_host_frees.items(): + result.host_frees[ComponentType(component)].extend(tensors) + return result + + +class _RustKVCacheEventRecorder: + """Expose the Rust event queue through the Python recorder interface.""" + + def __init__(self, binding, enabled: bool): + self._binding = binding + self.enabled = enabled + + def record_all_cleared(self) -> None: + self._binding.record_all_cleared_event() + + def take(self) -> list: + return [_kv_event_from_tagged(event) for event in self._binding.take_events()] + + +class RustUnifiedTreeCore(UnifiedTreeCoreInterface): + """A TreeCore backed by the Rust extension binding.""" + + _bindings = bindings + + def __init__(self, params: CacheInitParams): + assert params.tree_components is not None + self.tree_components = tuple(params.tree_components) + + # TODO(Jialin): Port session-reference-aware TreeCore support from #29173. + if params.enable_session_radix_cache: + raise ValueError( + "--enable-session-radix-cache is not supported by the Rust TreeCore" + ) + + # TODO(Jialin): Port custom component registration from #25754 and + # C128 support from #33676. + unsupported_components = set(self.tree_components) - { + ComponentType.FULL, + ComponentType.SWA, + ComponentType.MAMBA, + } + if unsupported_components: + names = ", ".join( + sorted(component.name for component in unsupported_components) + ) + raise ValueError(f"Rust TreeCore does not support components: {names}") + if params.component_registry_override: + raise ValueError( + "Rust TreeCore does not support component_registry_override" + ) + + self._page_size = params.page_size + self.is_eagle = ( + params.is_eagle and ComponentType.MAMBA not in self.tree_components + ) + + # ``device`` is derived from the construction-time allocator; the + # allocator/pool themselves are owned by the cache, not the tree. + if params.token_to_kv_pool_allocator: + device = torch.device(params.token_to_kv_pool_allocator.device) + # A bare "cuda" means the process's current device, not cuda:0. + if device.type == "cuda" and device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + self.device = device + else: + self.device = torch.device("cpu") + + self.enable_kv_cache_events = params.enable_kv_cache_events + has_mamba = ComponentType.MAMBA in self.tree_components + mamba_max_states_per_path = ( + get_exec().mamba.mamba_max_states_per_path if has_mamba else -1 + ) + + self._binding = self._binding_class()( + self._bindings.TreeCoreInitParamsBinding( + eviction_policy=params.eviction_policy, + page_size=params.page_size, + is_write_back=False, + enable_hicache=False, + write_through_threshold=256, + device=str(self.device), + swa_sliding_window_size=params.sliding_window_size, + enable_kv_cache_events=params.enable_kv_cache_events, + mamba_cache_chunk_size=( + mamba_cache_chunk_size() if has_mamba else None + ), + mamba_max_states_per_path=( + mamba_max_states_per_path + if mamba_max_states_per_path >= 0 + else None + ), + ), + [int(component) for component in self.tree_components], + ) + self.kv_events = _RustKVCacheEventRecorder( + self._binding, params.enable_kv_cache_events + ) + # The default-root empty result, prebuilt once from the binding. + self._empty_match_result = _match_result_from_binding( + self._binding.empty_match_result() + ) + + def _binding_class(self) -> type: + """The extension binding class this core constructs.""" + if self.is_eagle: + return self._bindings.RustBigramUnifiedTreeCoreBinding + return self._bindings.RustUnifiedTreeCoreBinding + + # ==== Tree API ==== + + def reset(self) -> None: + self._binding.reset() + # Node handles are never re-minted, so the fresh root gets a new one. + self._empty_match_result = _match_result_from_binding( + self._binding.empty_match_result() + ) + + def node_by_id(self, node_id: NodeId) -> UnifiedTreeNode: + # TODO(Jialin): Move the remaining Python-node consumers to + # backend-neutral APIs: sessions (#29173), C128 (#33676). + raise NotImplementedError("node_by_id: not yet ported to the Rust tree core") + + @property + def root_node(self) -> UnifiedTreeNode: + raise NotImplementedError("root_node: not yet ported to the Rust tree core") + + def inc_lock_ref( + self, + node_id: NodeId, + skip_lock_components: Sequence[ComponentType] = (), + ) -> IncLockRefResult: + result = self._binding.inc_lock_ref( + node_id, [int(component) for component in skip_lock_components] + ) + return _inc_lock_ref_result_from_binding(result) + + def dec_lock_ref( + self, + node_id: NodeId, + params: Optional[DecLockRefParams] = None, + skip_swa: bool = False, + ) -> DecLockRefResult: + binding_params = ( + self._bindings.DecLockRefParamsBinding( + swa_uuid_for_lock=params.swa_uuid_for_lock, + swa_uuid_for_host_lock=params.swa_uuid_for_host_lock, + skip_lock_node_ids=_skip_lock_node_ids_to_binding( + params.skip_lock_node_ids + ), + ) + if params is not None + else None + ) + self._binding.dec_lock_ref(node_id, binding_params, skip_swa) + return DecLockRefResult() + + def dec_swa_lock_only( + self, + node_id: NodeId, + swa_uuid_for_lock: Optional[int], + skip_lock_node_ids: Optional[dict] = None, + ) -> DecSwaLockOnlyResult: + result = DecSwaLockOnlyResult() + new_device_frees, new_host_frees = self._binding.dec_swa_lock_only( + node_id, + swa_uuid_for_lock, + ( + _skip_lock_node_ids_to_binding(skip_lock_node_ids) + if skip_lock_node_ids + else None + ), + ) + for component, tensors in new_device_frees.items(): + result.device_frees[ComponentType(component)].extend(tensors) + for component, tensors in new_host_frees.items(): + result.host_frees[ComponentType(component)].extend(tensors) + return result + + # ==== Device eviction (driven step-wise by the Controller's evict()) ==== + + def evict_device_start( + self, component_type: ComponentType, request_cnt: int + ) -> None: + self._binding.evict_device_start(int(component_type), request_cnt) + + def evict_device_next_node( + self, component_type: ComponentType, tracker: dict[ComponentType, int] + ) -> EvictDeviceNextNodeResult: + binding_result = self._binding.evict_device_next_node( + int(component_type), _tracker_to_binding(tracker) + ) + result = EvictDeviceNextNodeResult( + node_id=binding_result.node_id, + made_progress=binding_result.made_progress, + ) + return _fill_evict_result(binding_result, result) + + def evict_device_leaf( + self, node_id: NodeId, is_write_back: bool + ) -> EvictDeviceLeafResult: + # The binding reads is_write_back from the core's construction config. + assert ( + is_write_back == self.is_write_back + ), "is_write_back must match the core's construction config" + binding_result = self._binding.evict_device_leaf(node_id) + backup = binding_result.backup_kv + result = EvictDeviceLeafResult( + backup_kv=_cache_action_from_tagged(backup) if backup is not None else None + ) + return _fill_evict_result(binding_result, result) + + def demote(self, node_id: NodeId) -> DemoteResult: + binding_result = self._binding.demote(node_id) + return _fill_evict_result(binding_result, DemoteResult()) + + def evict_device_end(self, component_type: ComponentType) -> None: + self._binding.evict_device_end(int(component_type)) + + def inc_host_lock_ref(self, node_id: NodeId) -> IncLockRefResult: + result = self._binding.inc_host_lock_ref(node_id) + return IncLockRefResult( + delta=result.delta, + swa_uuid_for_lock=result.swa_uuid_for_lock, + swa_uuid_for_host_lock=result.swa_uuid_for_host_lock, + skip_lock_node_ids=_skip_lock_node_ids_from_binding( + result.skip_lock_node_ids + ), + ) + + def dec_host_lock_ref( + self, node_id: NodeId, params: Optional[DecLockRefParams] = None + ) -> DecLockRefResult: + binding_params = ( + self._bindings.DecLockRefParamsBinding( + swa_uuid_for_lock=params.swa_uuid_for_lock, + swa_uuid_for_host_lock=params.swa_uuid_for_host_lock, + skip_lock_node_ids=_skip_lock_node_ids_to_binding( + params.skip_lock_node_ids + ), + ) + if params is not None + else None + ) + self._binding.dec_host_lock_ref(node_id, binding_params) + return DecLockRefResult() + + def evictable_size(self) -> int: + return self._binding.evictable_size() + + def protected_size(self) -> int: + return self._binding.protected_size() + + def component_evictable_size(self, component_type: ComponentType) -> int: + return self._binding.component_evictable_size(int(component_type)) + + def full_evictable_size(self) -> int: + return self._binding.full_evictable_size() + + def full_protected_size(self) -> int: + return self._binding.full_protected_size() + + def swa_evictable_size(self) -> int: + return self._binding.component_evictable_size(int(ComponentType.SWA)) + + def mamba_evictable_size(self) -> int: + return self._binding.component_evictable_size(int(ComponentType.MAMBA)) + + def swa_protected_size(self) -> int: + return self._binding.component_protected_size(int(ComponentType.SWA)) + + def mamba_protected_size(self) -> int: + return self._binding.component_protected_size(int(ComponentType.MAMBA)) + + def total_size(self) -> tuple[int, int]: + return self._binding.total_size() + + def all_values_flatten(self) -> torch.Tensor: + return self._binding.all_values_flatten() + + def walk_for_kv_canary( + self, unlocked_only: bool, swa_resident_only: bool + ) -> RadixCacheWalkResult: + result = self._binding.walk_for_kv_canary(unlocked_only, swa_resident_only) + return RadixCacheWalkResult( + slot_indices=result.slot_indices, + positions=result.positions, + prev_slot_indices=result.prev_slot_indices, + ) + + def _record_all_cleared_event(self) -> None: + self.kv_events.record_all_cleared() + + def take_events(self) -> list: + return self.kv_events.take() + + def all_mamba_values_flatten(self) -> torch.Tensor: + return self._binding.all_mamba_values_flatten() + + def match_prefix(self, params: MatchPrefixParams) -> MatchResult: + key = params.key + result = self._binding.match_prefix( + self._bindings.MatchParamsBinding( + key=_radix_key_buffer(key), + extra_key=key.extra_key, + cache_salt=key.cache_salt, + ) + ) + return _match_result_from_binding(result) + + @property + def empty_match_result(self) -> MatchResult: + return self._empty_match_result + + def is_full_device_evicted(self, node_id: NodeId) -> bool: + return self._binding.is_full_device_evicted(node_id) + + def collect_full_device_indices( + self, from_node_id: NodeId, until_node_id: NodeId + ) -> torch.Tensor: + return self._binding.collect_full_device_indices(from_node_id, until_node_id) + + def begin_insert(self, params: InsertParams) -> InsertStepResult: + key = params.key + key_buffer = _radix_key_buffer(key) + value = params.value + if value is None: + # The binding always receives a value tensor; fall back to the + # token ids materialized on the core's device. + value = torch.tensor(key_buffer, dtype=torch.int64, device=self.device) + step = self._binding.begin_insert( + self._bindings.InsertParamsBinding( + key=key_buffer, + value=value, + extra_key=key.extra_key, + cache_salt=key.cache_salt, + mamba_value=params.mamba_value, + prev_prefix_len=params.prev_prefix_len, + swa_evicted_seqlen=params.swa_evicted_seqlen, + chunked=params.chunked, + priority=0 if params.priority is None else params.priority, + track_adopted_ranges=params.track_adopted_ranges, + ) + ) + return _insert_step_from_binding(step) + + def resume_insert(self) -> InsertStepResult: + return _insert_step_from_binding(self._binding.resume_insert()) + + def has_ongoing_insert(self) -> bool: + return self._binding.has_ongoing_insert() + + def end_insert(self) -> list[CacheAction | ComponentAction]: + return _cache_actions_from_tagged(self._binding.end_insert()) + + def drive_host_eviction( + self, component_type: ComponentType, num_tokens: int + ) -> DriveHostEvictionResult: + binding_result = self._binding.drive_host_eviction( + int(component_type), num_tokens + ) + return _fill_evict_result(binding_result, DriveHostEvictionResult()) + + def evict_excess_path_states( + self, + tail_node_id: NodeId, + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + ) -> None: + binding_result = self._binding.evict_excess_path_states(tail_node_id) + for component, tensors in binding_result.new_device_frees.items(): + device_frees[ComponentType(component)].extend(tensors) + for component, tensors in binding_result.new_host_frees.items(): + host_frees[ComponentType(component)].extend(tensors) + + # ==== HiCache ==== + + def set_hicache_enabled(self) -> None: + self._binding.set_hicache_enabled() + + @property + def page_size(self) -> int: + # Read-only: the Rust core freezes it at construction. + return self._page_size + + @property + def enable_hicache(self) -> bool: + return self._binding.enable_hicache() + + @property + def has_swa_host_pool(self) -> bool: + return self._binding.has_swa_host_pool() + + @has_swa_host_pool.setter + def has_swa_host_pool(self, value: bool) -> None: + # The Rust core has no unset path; reject a True -> False transition. + assert value or not self.has_swa_host_pool + if value: + self._binding.set_has_swa_host_pool() + + @property + def write_through_threshold(self) -> int: + return self._binding.write_through_threshold() + + @write_through_threshold.setter + def write_through_threshold(self, value: int) -> None: + # The cache assigns tree_core.write_through_threshold at HiCache init. + self._binding.set_write_through_threshold(value) + + @property + def is_write_back(self) -> bool: + return self._binding.is_write_back() + + @is_write_back.setter + def is_write_back(self, value: bool) -> None: + # The cache assigns tree_core.is_write_back at HiCache init; forward it. + self._binding.set_is_write_back(value) + + @property + def enable_storage(self) -> bool: + return self._binding.enable_storage() + + @enable_storage.setter + def enable_storage(self, value: bool) -> None: + # The cache assigns tree_core.enable_storage at storage init; forward it. + self._binding.set_enable_storage(value) + + @property + def enable_external_cache_linker(self) -> bool: + return False + + @enable_external_cache_linker.setter + def enable_external_cache_linker(self, value: bool) -> None: + # TODO(Jialin): Port external cache linker support from #37091 and #37151. + if value: + raise ValueError( + "External cache linker is not supported by the Rust TreeCore" + ) + + def insert_host( + self, + node_id: NodeId, + key: RadixKey, + host_value: torch.Tensor, + hash_value: list[str], + ) -> InsertResult: + result = self._binding.insert_host( + node_id, + key.extra_key, + _radix_key_buffer(key), + host_value, + list(hash_value), + key.cache_salt, + ) + return InsertResult( + prefix_len=result.prefix_len, + total_len=result.total_len, + last_device_node=result.last_device_node, + inserted_host_node=result.inserted_host_node, + host_insert_dropped=result.host_insert_dropped, + mamba_exist=result.mamba_exist, + cache_actions=_cache_actions_from_tagged(result.cache_actions), + ) + + def build_backup_spec( + self, node_id: NodeId + ) -> tuple[torch.Tensor, dict[ComponentType, list[PoolTransfer]]]: + device_value, comp_xfers = self._binding.build_backup_spec(node_id) + return device_value, _comp_xfers_from_binding(comp_xfers) + + def build_storage_backup_spec( + self, node_id: NodeId, pass_prefix_keys: bool + ) -> Optional[StorageBackupSpec]: + spec = self._binding.build_storage_backup_spec(node_id, pass_prefix_keys) + if spec is None: + return None + # Token ids cross the boundary as raw int64 bytes, not per-token ints. + token_ids = array("q") + token_ids.frombytes(spec.token_ids) + return StorageBackupSpec( + host_value=spec.host_value, + token_ids=token_ids, + hash_value=spec.hash_value, + prefix_keys=spec.prefix_keys, + comp_xfers=_comp_xfers_from_binding(spec.comp_xfers), + ) + + def build_hicache_transfers( + self, + component_type: ComponentType, + node_id: NodeId, + phase: CacheTransferPhase, + *, + host_indices: Optional[torch.Tensor] = None, + token_ids: Optional[Sequence[int]] = None, + prefetch_tokens: int = 0, + last_hash: Optional[str] = None, + ) -> Optional[list[PoolTransfer]]: + transfers = self._binding.build_hicache_transfers( + int(component_type), + node_id, + phase.value, + host_indices, + # TODO: Forward token ids when Rust Mamba prefetch consumes them. + None, + prefetch_tokens, + last_hash, + ) + if transfers is None: + return None + return [_transfer_from_binding(transfer) for transfer in transfers] + + def build_load_back_spec( + self, node_id: NodeId, req: Optional[Req] = None + ) -> tuple[PoolTransfer, dict[ComponentType, list[PoolTransfer]]]: + # Component hooks take primitives, not Req: extract its fields here. + mamba_pool_idx = req.kv.mamba_pool_idx if req is not None else None + kv_xfer, comp_xfers = self._binding.build_load_back_spec( + node_id, mamba_pool_idx + ) + return _transfer_from_binding(kv_xfer), _comp_xfers_from_binding(comp_xfers) + + def prefetch_anchor_info( + self, node_id: NodeId + ) -> tuple[Optional[str], Optional[str]]: + return self._binding.prefetch_anchor_info(node_id) + + def is_backuped(self, node_id: NodeId) -> bool: + return self._binding.node_backuped(node_id) + + def is_root(self, node_id: NodeId) -> bool: + return self._binding.is_root(node_id) + + def get_last_hash_value(self, node_id: NodeId) -> Optional[str]: + return self._binding.get_last_hash_value(node_id) + + def get_prefix_hash_values(self, node_id: NodeId) -> list[str]: + return self._binding.get_prefix_hash_values(node_id) + + def get_hash_values(self, node_id: NodeId) -> list[str]: + return self._binding.get_hash_values(node_id) + + def snapshot_buffer_backup( + self, node_id: NodeId, pass_prefix_keys: bool + ) -> Optional[BufferBackupSnapshot]: + snapshot = self._binding.snapshot_buffer_backup(node_id, pass_prefix_keys) + if snapshot is None: + return None + token_ids = array("q") + token_ids.frombytes(snapshot.key_token_ids) + return BufferBackupSnapshot( + node_id=snapshot.node_id, + parent_node_id=snapshot.parent_node_id, + parent_is_root=snapshot.parent_is_root, + parent_last_hash=snapshot.parent_last_hash, + hash_values=snapshot.hash_values, + key=RadixKey( + token_ids, + extra_key=snapshot.extra_key, + is_bigram=snapshot.is_bigram, + cache_salt=snapshot.cache_salt, + ), + prefix_keys=snapshot.prefix_keys, + ) + + def validate_buffer_backup( + self, node_id: NodeId, expected_key_length: int + ) -> Optional[BufferBackupState]: + state = self._binding.validate_buffer_backup(node_id, expected_key_length) + if state is None: + return None + return BufferBackupState( + parent_node_id=state.parent_node_id, + parent_is_root=state.parent_is_root, + parent_last_hash=state.parent_last_hash, + ) + + def backfill_missing_hash_values(self) -> int: + return self._binding.backfill_missing_hash_values() + + def root_node_handle(self, extra_key: Optional[str] = None) -> NodeId: + return self._binding.root_node_handle(extra_key) + + def dfs_weight_order(self, node_ids: Sequence[NodeId]) -> list[int]: + return self._binding.dfs_weight_order(list(node_ids)) + + def commit_hicache_transfers( + self, + node_id: NodeId, + phase: CacheTransferPhase, + comp_xfers: dict[ComponentType, list[PoolTransfer]], + *, + cache_actions: list[CacheAction | ComponentAction], + insert_result: Optional[InsertResult] = None, + pool_storage_result: Optional[PoolTransferResult] = None, + ) -> None: + actions, mamba_exist = self._binding.commit_hicache_transfers( + node_id, + phase.value, + _comp_xfers_to_binding(comp_xfers), + ( + None + if insert_result is None + else ( + insert_result.total_len, + insert_result.inserted_host_node, + insert_result.mamba_exist, + ) + ), + ( + None + if pool_storage_result is None + else ( + pool_storage_result.kv_hit_pages, + dict(pool_storage_result.extra_pool_hit_pages), + ) + ), + ) + if insert_result is not None and mamba_exist is not None: + insert_result.mamba_exist = mamba_exist + cache_actions.extend(_cache_actions_from_tagged(actions)) + + def commit_backup( + self, + node_id: NodeId, + host_indices: torch.Tensor, + comp_xfers: dict[ComponentType, list[PoolTransfer]], + ) -> None: + self._binding.commit_backup( + node_id, host_indices, _comp_xfers_to_binding(comp_xfers) + ) + + def commit_load_back( + self, + node_id: NodeId, + device_indices: torch.Tensor, + kv_xfer: PoolTransfer, + comp_xfers: dict[ComponentType, list[PoolTransfer]], + ) -> list[CacheAction | ComponentAction]: + actions = self._binding.commit_load_back( + node_id, + device_indices, + _transfer_to_binding(kv_xfer), + _comp_xfers_to_binding(comp_xfers), + ) + return _cache_actions_from_tagged(actions) + + def drop_subtree_no_host(self, node_id: NodeId) -> DropSubtreeNoHostResult: + binding_result = self._binding.drop_subtree_no_host(node_id) + result = DropSubtreeNoHostResult(is_dropped=binding_result.dropped) + return _fill_evict_result(binding_result, result) + + def mark_write_through_pending(self, node_id: NodeId) -> None: + self._binding.mark_write_through_pending(node_id) + + def finish_write_through(self, node_ids: list[NodeId], ack_id: int) -> None: + self._binding.finish_write_through(list(node_ids), ack_id) + + def finish_load_back(self, anchor_node_id: NodeId) -> None: + self._binding.finish_load_back(anchor_node_id) + + @property + def write_back_duplicate_reclaim_digest(self) -> int: + return self._binding.write_back_duplicate_reclaim_digest() + + def set_component_device_value( + self, node_id: NodeId, component_type: ComponentType, value: torch.Tensor + ) -> None: + self._binding.set_component_device_value( + node_id, int(component_type), value.to(torch.int64) + ) + + def get_component_device_value( + self, node_id: NodeId, component_type: ComponentType + ) -> Optional[torch.Tensor]: + return self._binding.get_component_device_value(node_id, int(component_type)) + + def component_has_host_value_only( + self, node_id: NodeId, component_type: ComponentType + ) -> bool: + return self._binding.component_has_host_value_only(node_id, int(component_type)) + + # ==== Others ==== + + def sanity_check( + self, + ongoing_write_through: list[tuple[int, NodeId]], + ongoing_load_back: list[tuple[int, NodeId]], + ) -> None: + self._binding.sanity_check(ongoing_write_through, ongoing_load_back) + + def pretty_print(self) -> None: + self._binding.pretty_print() diff --git a/python/sglang/srt/mem_cache/rust_tree_core/extension.py b/python/sglang/srt/mem_cache/rust_tree_core/extension.py new file mode 100644 index 000000000..955fe7906 --- /dev/null +++ b/python/sglang/srt/mem_cache/rust_tree_core/extension.py @@ -0,0 +1,40 @@ +"""Load the bundled Rust TreeCore extension or a fingerprinted local build.""" + +from pathlib import Path + +# Loading torch first makes its libtorch dependencies resident before dlopen. +import torch + +from sglang.srt.rust_extensions import load_rust_extension +from sglang.srt.rust_extensions.torch_build import torch_build_configuration + +_PYTHON_MODULE = "sglang.srt.mem_cache.rust_tree_core.mem_cache" +_INSPECTION_MODULE = "sglang.srt.mem_cache.rust_tree_core.mem_cache_inspection" +_CRATE_DIR = Path(__file__).resolve().parents[5] / "rust" / "mem-cache" +_TORCH_COMPAT_HEADER = _CRATE_DIR / "torch_2_13_compat.h" + + +def load_tree_core_extension(*, inspection: bool = False): + """Load the production binding or the test-only inspection variant.""" + build = torch_build_configuration( + compat_header=_TORCH_COMPAT_HEADER, + python_module=_PYTHON_MODULE, + torch_module=torch, + ) + return load_rust_extension( + _PYTHON_MODULE, + additional_features=("inspection",) if inspection else (), + extension_module=_INSPECTION_MODULE if inspection else None, + build_environment=build.environment, + build_fingerprint=build.fingerprint, + ) + + +bindings = load_tree_core_extension() + +DecLockRefParamsBinding = bindings.DecLockRefParamsBinding +InsertParamsBinding = bindings.InsertParamsBinding +MatchParamsBinding = bindings.MatchParamsBinding +RustBigramUnifiedTreeCoreBinding = bindings.RustBigramUnifiedTreeCoreBinding +RustUnifiedTreeCoreBinding = bindings.RustUnifiedTreeCoreBinding +TreeCoreInitParamsBinding = bindings.TreeCoreInitParamsBinding diff --git a/python/sglang/srt/mem_cache/unified_cache/tree_core_registry.py b/python/sglang/srt/mem_cache/unified_cache/tree_core_registry.py index f8a30bc3f..47bd6e1b1 100644 --- a/python/sglang/srt/mem_cache/unified_cache/tree_core_registry.py +++ b/python/sglang/srt/mem_cache/unified_cache/tree_core_registry.py @@ -54,7 +54,17 @@ def _python_tree_core_factory( return UnifiedTreeCore(params, components) +def _rust_tree_core_factory( + params: CacheInitParams, components: dict[ComponentType, TreeComponent] +) -> UnifiedTreeCoreInterface: + """Load and construct the in-tree Rust TreeCore only when selected.""" + from sglang.srt.mem_cache.rust_tree_core.adapter import RustUnifiedTreeCore + + return RustUnifiedTreeCore(params) + + register_tree_core_backend("python", _python_tree_core_factory) +register_tree_core_backend("rust", _rust_tree_core_factory) def create_tree_core( diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index b9e3c8a9c..5d3a02cd3 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -35,6 +35,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( InsertResult, MatchPrefixParams, MatchResult, + _dfs_weight_order, ) from sglang.srt.mem_cache.events import KVCacheEventRecorder from sglang.srt.mem_cache.hicache_storage import ( @@ -63,6 +64,8 @@ from sglang.srt.mem_cache.unified_cache.components import ( get_and_increase_time_counter, ) from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import ( + BufferBackupSnapshot, + BufferBackupState, DecSwaLockOnlyResult, DemoteResult, DriveHostEvictionResult, @@ -518,6 +521,55 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): """The hash values owned by this node, excluding its ancestors.""" return self.node_by_id(node_id).hash_value or [] + def snapshot_buffer_backup( + self, node_id: NodeId, pass_prefix_keys: bool + ) -> Optional[BufferBackupSnapshot]: + node = self._node_arena.get(node_id) + if ( + node is None + or node is self.root_node + or not node.hash_value + or node.component_data[BASE_COMPONENT_TYPE].value is None + ): + return None + parent = node.parent + assert parent is not None and node.key is not None + return BufferBackupSnapshot( + node_id=node.id, + parent_node_id=parent.id, + parent_is_root=parent is self.root_node, + parent_last_hash=parent.get_last_hash_value(), + hash_values=list(node.hash_value), + key=RadixKey( + array("q", node.key.raw_token_ids()), + extra_key=node.key.extra_key, + is_bigram=node.key.is_bigram, + cache_salt=node.key.cache_salt, + ), + prefix_keys=( + node.get_prefix_hash_values(parent) if pass_prefix_keys else None + ), + ) + + def validate_buffer_backup( + self, node_id: NodeId, expected_key_length: int + ) -> Optional[BufferBackupState]: + node = self._node_arena.get(node_id) + if ( + node is None + or node.component_data[BASE_COMPONENT_TYPE].value is None + or len(node.key) != expected_key_length + ): + return None + parent = node.parent + if parent is None: + return None + return BufferBackupState( + parent_node_id=parent.id, + parent_is_root=parent is self.root_node, + parent_last_hash=parent.get_last_hash_value(), + ) + def backfill_missing_hash_values(self) -> int: """Hash every node that was built while storage was disabled. @@ -543,6 +595,9 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): """The NodeId anchoring matches; the single root serves every namespace.""" return self.root_node.id + def dfs_weight_order(self, node_ids: Sequence[NodeId]) -> list[int]: + return _dfs_weight_order(self.root_node, node_ids, self.node_by_id) + def _new_node(self, priority: int = 0) -> UnifiedTreeNode: """Create and register a tree node in the arena.""" node = UnifiedTreeNode(self.component_types, priority=priority) @@ -872,7 +927,8 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): def begin_insert(self, params: InsertParams) -> InsertStepResult: """Start the insert, running to its first barrier or completion.""" # Insert walks are single-flight; a live walk means re-entrancy. - assert self._ongoing_insert_walk_state is None, "concurrent insert walks" + if self._ongoing_insert_walk_state is not None: + raise RuntimeError("concurrent insert walks") key = params.key value = params.value key, value = key.maybe_to_bigram_view(self.is_eagle, value) @@ -913,7 +969,8 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): def resume_insert(self) -> InsertStepResult: """Continue the suspended insert after its step actions were executed.""" - assert self._ongoing_insert_walk_state is not None, "no in-flight insert" + if self._ongoing_insert_walk_state is None: + raise RuntimeError("no in-flight insert") return self._advance_insert() def has_ongoing_insert(self) -> bool: diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py index 53b2f420f..a497f34c7 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py @@ -80,6 +80,22 @@ class RadixCacheWalkResult(msgspec.Struct, frozen=True, kw_only=True): prev_slot_indices: torch.Tensor +class BufferBackupSnapshot(msgspec.Struct, frozen=True): + node_id: NodeId + parent_node_id: NodeId + parent_is_root: bool + parent_last_hash: Optional[str] + hash_values: list[str] + key: RadixKey + prefix_keys: Optional[list[str]] + + +class BufferBackupState(msgspec.Struct, frozen=True): + parent_node_id: NodeId + parent_is_root: bool + parent_last_hash: Optional[str] + + class InsertStepResult(msgspec.Struct, frozen=True): """One step of a resumable insert: the Controller executes ``actions``, then resumes while ``result`` is None; ``result`` is set on the final step.""" @@ -181,6 +197,20 @@ class UnifiedTreeCoreInterface(ABC): """The hash values owned by this node, excluding its ancestors.""" ... + @abstractmethod + def snapshot_buffer_backup( + self, node_id: NodeId, pass_prefix_keys: bool + ) -> Optional[BufferBackupSnapshot]: + """Snapshot an eligible buffer-only backup node.""" + ... + + @abstractmethod + def validate_buffer_backup( + self, node_id: NodeId, expected_key_length: int + ) -> Optional[BufferBackupState]: + """Validate a queued backup and return its current parent state.""" + ... + @abstractmethod def backfill_missing_hash_values(self) -> int: """Hash every node built while storage was disabled; return how many. @@ -196,6 +226,11 @@ class UnifiedTreeCoreInterface(ABC): """The NodeId anchoring matches for the namespace.""" ... + @abstractmethod + def dfs_weight_order(self, node_ids: Sequence[NodeId]) -> list[int]: + """Return input indices in depth-first, subtree-weight order.""" + ... + @abstractmethod def inc_lock_ref( self, node_id: NodeId, skip_lock_components: Sequence[ComponentType] = () diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index b8a2cb69b..925eac13f 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -199,8 +199,9 @@ class UnifiedRadixCache(BasePrefixCache): ) # The TreeCore owns the tree member-var state (structure, LRUs, sizes, # evictable leaves) and drives the components' tree-level hooks. + self._tree_core_backend = envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.get() self.tree_core = create_tree_core( - name=envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.get(), + name=self._tree_core_backend, params=params, components=self.components, ) @@ -386,6 +387,8 @@ class UnifiedRadixCache(BasePrefixCache): """Initialize HiCache infrastructure.""" self.host_memory_mode = get_memory().hicache_host_memory_mode if self.host_memory_mode == "buffer_only": + # TODO(Jialin): Extend buffer-only state handoff to Mamba in a + # follow-up to #34798 and #35769. # FULL and FULL+SWA only: Mamba has no state-handoff channel on # the admission-time load-back read path and is not layer-gated. # Lifting the fence also needs the admission charge: a staged @@ -1340,9 +1343,7 @@ class UnifiedRadixCache(BasePrefixCache): # FIFO ordering instead (BackupKV chains are parent-before-child # and every pipeline stage drains in order). for node_id in action.node_ids: - self.buffer_pipeline.enqueue_backup_intent( - self.tree_core.node_by_id(node_id) - ) + self.buffer_pipeline.enqueue_backup_intent(node_id) return 0 written = 0 for node_id in action.node_ids: @@ -3120,3 +3121,6 @@ class UnifiedRadixCache(BasePrefixCache): def root_node_handle(self, extra_key: Optional[str] = None) -> NodeId: """The root's NodeId -- URC match results carry NodeIds.""" return self.tree_core.root_node_handle(extra_key) + + def dfs_weight_order(self, node_handles: Sequence[NodeId]) -> list[int]: + return self.tree_core.dfs_weight_order(node_handles) diff --git a/python/sglang/srt/rust_extensions/loader.py b/python/sglang/srt/rust_extensions/loader.py index 654576dbf..bc76bb1af 100644 --- a/python/sglang/srt/rust_extensions/loader.py +++ b/python/sglang/srt/rust_extensions/loader.py @@ -19,7 +19,7 @@ from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path from types import ModuleType -from typing import Iterator, Literal +from typing import Iterator, Literal, Mapping try: import tomllib @@ -52,6 +52,7 @@ class _CrateSpec: package: str library: str python_module: str + manifest: Path workspace: Path features: tuple[str, ...] @@ -69,6 +70,10 @@ def load_rust_extension( mode: RustBuildMode | None = None, cache_dir: Path | None = None, workspace: Path | None = None, + additional_features: tuple[str, ...] = (), + extension_module: str | None = None, + build_environment: Mapping[str, str] | None = None, + build_fingerprint: Mapping[str, object] | None = None, ) -> ModuleType: """Import a PyO3 extension, compiling it locally when permitted and needed. @@ -77,9 +82,13 @@ def load_rust_extension( to ``python_module`` (the same metadata setup.py uses for wheel builds), so new crates need no registration here. - ``auto`` prefers a module bundled in the installed wheel, then a cached - local build, and finally Cargo. ``never`` permits the first two but never - invokes Cargo. ``force`` rebuilds from source and replaces the cache entry. + ``auto`` prefers a module bundled in an installed wheel. In a source tree, + it ignores unverified in-package artifacts and uses the fingerprinted cache + before invoking Cargo. ``never`` explicitly trusts a bundled module, then + permits the cache but never invokes Cargo. ``force`` rebuilds from source. + A same-name feature variant is always sourced from the fingerprinted cache. + A distinctly named variant may be supplied by test infrastructure and is + otherwise built into that cache after its bundled import misses. ``mode`` defaults to ``SGLANG_RUST_BUILD_MODE``. """ if mode is None: @@ -89,29 +98,45 @@ def load_rust_extension( f"invalid Rust extension build mode {mode!r}; expected auto, never, or force" ) - if mode != "force": - module = _import_bundled_extension(python_module) - if module is not None: - return module - elif python_module in sys.modules: + load_module = extension_module or python_module + same_name_feature_variant = ( + bool(additional_features) and load_module == python_module + ) + if loaded := sys.modules.get(load_module): + if mode != "force": + return loaded raise RuntimeError( - f"cannot force-build {python_module} after it has been imported; " + f"cannot force-build {load_module} after it has been imported; " "start a new Python process" ) - if workspace is None: workspace = _RUST_WORKSPACE + source_checkout = (Path(workspace) / "Cargo.toml").is_file() + trust_bundled = mode == "never" or not source_checkout + if mode != "force" and trust_bundled and not same_name_feature_variant: + module = _import_bundled_extension(load_module) + if module is not None: + return module + crate = _discover_crate(workspace, python_module) - context = _build_context(crate) + features = tuple(dict.fromkeys((*crate.features, *additional_features))) + context = _build_context( + crate, + features=features, + build_fingerprint=build_fingerprint, + extension_module=load_module, + ) cache_root = _cache_root(cache_dir) - extension_path = _cached_extension_path(cache_root, crate, context.fingerprint) + extension_path = _cached_extension_path( + cache_root, crate, context.fingerprint, load_module + ) lock_path = ( cache_root / "locks" / f"{crate.package}-{context.target_fingerprint}.lock" ) with _filesystem_lock(lock_path): if mode != "force" and extension_path.is_file(): - return _load_extension_from_path(crate.python_module, extension_path) + return _load_extension_from_path(load_module, extension_path) if mode == "never": raise ModuleNotFoundError( @@ -121,14 +146,19 @@ def load_rust_extension( ) target_dir = cache_root / "targets" / context.target_fingerprint - artifact = _cargo_build(crate, target_dir) + artifact = _cargo_build( + crate, + target_dir, + features=features, + build_environment=build_environment, + ) if _source_digest(crate.workspace) != context.source_digest: raise RuntimeError( f"Rust sources under {crate.workspace} changed during the build; " "the result was not cached" ) _stage_atomically(artifact, extension_path) - return _load_extension_from_path(crate.python_module, extension_path) + return _load_extension_from_path(load_module, extension_path) def _import_bundled_extension(module_name: str) -> ModuleType | None: @@ -143,16 +173,10 @@ def _import_bundled_extension(module_name: str) -> ModuleType | None: def _discover_crate(workspace: Path, python_module: str) -> _CrateSpec: workspace = Path(workspace).resolve() workspace_manifest = workspace / "Cargo.toml" - lockfile = workspace / "Cargo.lock" if not workspace_manifest.is_file(): raise FileNotFoundError( f"Rust workspace for {python_module} was not found at {workspace}" ) - if not lockfile.is_file(): - raise FileNotFoundError( - f"{lockfile} is required for reproducible `cargo build --locked` builds" - ) - matches: list[_CrateSpec] = [] declared_modules: list[str] = [] for manifest in _source_files(workspace): @@ -178,12 +202,19 @@ def _discover_crate(workspace: Path, python_module: str) -> _CrateSpec: f"{manifest} declares python-module {python_module!r} but must " "also set `package.name` and `lib.name`" ) + crate_workspace = manifest.parent if "workspace" in document else workspace + lockfile = crate_workspace / "Cargo.lock" + if not lockfile.is_file(): + raise FileNotFoundError( + f"{lockfile} is required for reproducible `cargo build --locked` builds" + ) matches.append( _CrateSpec( package=package_name, library=library, python_module=python_module, - workspace=workspace, + manifest=manifest, + workspace=crate_workspace, features=tuple(sglang_metadata.get("features", ())), ) ) @@ -203,7 +234,17 @@ def _discover_crate(workspace: Path, python_module: str) -> _CrateSpec: return matches[0] -def _build_context(crate: _CrateSpec) -> _BuildContext: +def _build_context( + crate: _CrateSpec, + *, + features: tuple[str, ...] | None = None, + build_fingerprint: Mapping[str, object] | None = None, + extension_module: str | None = None, +) -> _BuildContext: + if features is None: + features = crate.features + if extension_module is None: + extension_module = crate.python_module source_digest = _source_digest(crate.workspace) toolchain = { "cargo": _command_version( @@ -224,6 +265,7 @@ def _build_context(crate: _CrateSpec) -> _BuildContext: } target_inputs = { "build_environment": build_environment, + "extension_build": dict(build_fingerprint or {}), "python_abi": python_abi, "toolchain": toolchain, } @@ -234,6 +276,8 @@ def _build_context(crate: _CrateSpec) -> _BuildContext: "package": crate.package, "library": crate.library, "python_module": crate.python_module, + "extension_module": extension_module, + "features": features, "source_digest": source_digest, **target_inputs, } @@ -301,12 +345,15 @@ def _cache_root(cache_dir: Path | None) -> Path: def _cached_extension_path( - cache_root: Path, crate: _CrateSpec, fingerprint: str + cache_root: Path, + crate: _CrateSpec, + fingerprint: str, + extension_module: str | None = None, ) -> Path: extension_suffix = sysconfig.get_config_var("EXT_SUFFIX") if not extension_suffix: raise RuntimeError("Python did not report an EXT_SUFFIX for native extensions") - module_leaf = crate.python_module.rsplit(".", 1)[-1] + module_leaf = (extension_module or crate.python_module).rsplit(".", 1)[-1] return ( cache_root / "artifacts" @@ -327,7 +374,15 @@ def _filesystem_lock(path: Path) -> Iterator[None]: fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) -def _cargo_build(crate: _CrateSpec, target_dir: Path) -> Path: +def _cargo_build( + crate: _CrateSpec, + target_dir: Path, + *, + features: tuple[str, ...] | None = None, + build_environment: Mapping[str, str] | None = None, +) -> Path: + if features is None: + features = crate.features target_dir.mkdir(parents=True, exist_ok=True) command = [ "cargo", @@ -337,10 +392,10 @@ def _cargo_build(crate: _CrateSpec, target_dir: Path) -> Path: "--package", crate.package, ] - if crate.features: - command.extend(("--features", ",".join(crate.features))) + if features: + command.extend(("--features", ",".join(features))) - environment = os.environ.copy() + environment = dict(os.environ if build_environment is None else build_environment) environment["CARGO_TARGET_DIR"] = os.fspath(target_dir) environment["PYO3_PYTHON"] = sys.executable logger.info("Building %s with `%s`", crate.python_module, " ".join(command)) diff --git a/python/sglang/srt/rust_extensions/torch_build.py b/python/sglang/srt/rust_extensions/torch_build.py new file mode 100644 index 000000000..f4fd40adb --- /dev/null +++ b/python/sglang/srt/rust_extensions/torch_build.py @@ -0,0 +1,115 @@ +"""Build settings for Rust extensions that link against the active PyTorch.""" + +from __future__ import annotations + +import hashlib +import os +import re +import shlex +import sys +from dataclasses import dataclass +from pathlib import Path +from types import ModuleType +from typing import Mapping + +_MIN_SUPPORTED_TORCH = (2, 11) +_MAX_SUPPORTED_TORCH = (2, 13) + + +@dataclass(frozen=True) +class TorchBuildConfiguration: + """Environment overrides plus stable inputs for the artifact fingerprint.""" + + environment: dict[str, str] + fingerprint: dict[str, object] + + +def torch_build_configuration( + *, + compat_header: Path, + python_module: str, + torch_module: ModuleType | None = None, + base_environment: Mapping[str, str] | None = None, + include_absolute_rpath: bool = True, +) -> TorchBuildConfiguration: + """Describe a build against the torch package loaded by this interpreter.""" + if sys.platform != "linux": + raise RuntimeError("the Rust TreeCore extension currently supports Linux only") + + if torch_module is None: + try: + import torch as torch_module + except ModuleNotFoundError as exc: + raise RuntimeError( + "PyTorch must be installed before building the Rust TreeCore extension" + ) from exc + + version = str(torch_module.__version__) + match = re.match(r"^(\d+)\.(\d+)", version) + if match is None: + raise RuntimeError(f"could not parse PyTorch version {version!r}") + major_minor = (int(match.group(1)), int(match.group(2))) + if not _MIN_SUPPORTED_TORCH <= major_minor <= _MAX_SUPPORTED_TORCH: + minimum = ".".join(map(str, _MIN_SUPPORTED_TORCH)) + maximum = ".".join(map(str, _MAX_SUPPORTED_TORCH)) + raise RuntimeError( + f"the Rust TreeCore supports PyTorch {minimum} through {maximum}; " + f"found {version}" + ) + + torch_file = getattr(torch_module, "__file__", None) + if torch_file is None: + raise RuntimeError("the active PyTorch package has no filesystem location") + torch_root = Path(torch_file).resolve().parent + torch_lib = torch_root / "lib" + if not torch_lib.is_dir(): + raise RuntimeError( + f"the active PyTorch package has no library dir at {torch_lib}" + ) + + cxx11_abi_fn = getattr(torch_module, "compiled_with_cxx11_abi", None) + if cxx11_abi_fn is not None: + cxx11_abi = bool(cxx11_abi_fn()) + else: + cxx11_abi = bool(torch_module._C._GLIBCXX_USE_CXX11_ABI) + + environment = dict(os.environ if base_environment is None else base_environment) + environment["LIBTORCH_USE_PYTORCH"] = "1" + # tch 0.24 targets Torch 2.11. The compatibility header below covers the + # API removals in the supported 2.12/2.13 builds, after this explicit gate. + environment["LIBTORCH_BYPASS_VERSION_CHECK"] = "1" + environment["PYO3_PYTHON"] = sys.executable + environment["PATH"] = os.pathsep.join( + filter(None, (os.fspath(Path(sys.executable).parent), environment.get("PATH"))) + ) + environment["LD_LIBRARY_PATH"] = os.pathsep.join( + filter(None, (os.fspath(torch_lib), environment.get("LD_LIBRARY_PATH"))) + ) + + cxxflags = environment.get("CXXFLAGS", "") + environment["CXXFLAGS"] = ( + f"{cxxflags} -include {shlex.quote(os.fspath(compat_header.resolve()))}" + ).strip() + + package_depth = len(python_module.split(".")) - 1 + bundled_torch_lib = "$ORIGIN/" + "../" * package_depth + "torch/lib" + rustflags = environment.get("RUSTFLAGS", "") + rpath_flags = [f"-C link-arg=-Wl,-rpath,{bundled_torch_lib}"] + if include_absolute_rpath: + rpath_flags.append(f"-C link-arg=-Wl,-rpath,{torch_lib}") + environment["RUSTFLAGS"] = " ".join(filter(None, (rustflags, *rpath_flags))) + + fingerprint = { + "torch_version": version, + "torch_root": os.fspath(torch_root), + "torch_cxx11_abi": cxx11_abi, + "torch_cuda": getattr(torch_module.version, "cuda", None), + "torch_hip": getattr(torch_module.version, "hip", None), + "include_absolute_rpath": include_absolute_rpath, + "compat_header_sha256": ( + hashlib.sha256(compat_header.read_bytes()).hexdigest() + if compat_header.is_file() + else None + ), + } + return TorchBuildConfiguration(environment=environment, fingerprint=fingerprint) diff --git a/python/sglang/test/test_utils.py b/python/sglang/test/test_utils.py index fd96535de..aefa11251 100644 --- a/python/sglang/test/test_utils.py +++ b/python/sglang/test/test_utils.py @@ -658,6 +658,17 @@ def _wait_for_server_health( return False, "Server failed to start within the timeout period" +def unified_radix_tree_server_env( + tree_core_backend: str, **extra_env: str +) -> dict[str, str]: + return { + **os.environ, + **extra_env, + "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1", + "SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND": tree_core_backend, + } + + def popen_launch_server( model: str, base_url: str, diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 914eac885..553ae362a 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -5,6 +5,12 @@ members = [ "sglang-mm", "sglang-server" ] +exclude = ["mem-cache"] + +[workspace.metadata.sglang] +# Extension crates that intentionally cannot share this workspace's dependency +# graph (mem-cache currently uses the tch-compatible PyO3 0.22 API). +extension-manifests = ["mem-cache/Cargo.toml"] [workspace.package] version = "0.1.0" diff --git a/rust/mem-cache/.gitignore b/rust/mem-cache/.gitignore new file mode 100644 index 000000000..2f7896d1d --- /dev/null +++ b/rust/mem-cache/.gitignore @@ -0,0 +1 @@ +target/ diff --git a/rust/mem-cache/Cargo.lock b/rust/mem-cache/Cargo.lock new file mode 100644 index 000000000..d07a7c7e1 --- /dev/null +++ b/rust/mem-cache/Cargo.lock @@ -0,0 +1,905 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + +[[package]] +name = "aes" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" +dependencies = [ + "cfg-if", + "cipher", + "cpufeatures", +] + +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "base64ct" +version = "1.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + +[[package]] +name = "bzip2" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bdb116a6ef3f6c3698828873ad02c3014b3c85cadb88496095628e3ef1e347f8" +dependencies = [ + "bzip2-sys", + "libc", +] + +[[package]] +name = "bzip2-sys" +version = "0.1.13+1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "225bff33b2141874fe80d71e07d6eec4f85c5c216453dd96388240f96e1acc14" +dependencies = [ + "cc", + "pkg-config", +] + +[[package]] +name = "cc" +version = "1.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ad534f4357a5264cce5019c989cf66a4f0dc4e0d1b1d15f8aacec0ff7360273" +dependencies = [ + "find-msvc-tools", + "jobserver", + "libc", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "cipher" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" +dependencies = [ + "crypto-common", + "inout", +] + +[[package]] +name = "constant_time_eq" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "245097e9a4535ee1e3e3931fcfcd55a796a44c643e8596ff6566d68f09b87bbc" + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8498c871161e1742aaa9d52551b2d6ebdd4c3d45a3be423e3728f33b955be550" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "deranged" +version = "0.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", + "subtle", +] + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "find-msvc-tools" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d45db016d36b838f563236e9193d0ee6ce38f3f68b6c94e914b4929c96bbb890" + +[[package]] +name = "flate2" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb" +dependencies = [ + "crc32fast", + "miniz_oxide", +] + +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi", +] + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "hmac" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e" +dependencies = [ + "digest", +] + +[[package]] +name = "indoc" +version = "2.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" +dependencies = [ + "rustversion", +] + +[[package]] +name = "inout" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" +dependencies = [ + "generic-array", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "jobserver" +version = "0.1.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3" +dependencies = [ + "getrandom 0.4.3", + "libc", +] + +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[package]] +name = "libc" +version = "0.2.189" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" + +[[package]] +name = "matrixmultiply" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f607c237553f086e7043417a51df26b2eb899d3caff94e6a67592ff992fedc7" +dependencies = [ + "autocfg", + "rawpointer", +] + +[[package]] +name = "mem_cache" +version = "0.1.0" +dependencies = [ + "hashbrown", + "pyo3", + "sha2", + "tch", + "thiserror", +] + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg", +] + +[[package]] +name = "miniz_oxide" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b63fbc4a50860e98e7b2aa7804ded1db5cbc3aff9193adaff57a6931bf7c4b4c" +dependencies = [ + "adler2", + "simd-adler32", +] + +[[package]] +name = "ndarray" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "882ed72dce9365842bf196bdeedf5055305f11fc8c03dee7bb0194a6cad34841" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-conv" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441" + +[[package]] +name = "num-integer" +version = "0.1.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ce2d95d4b3734dc35aa2f45e1aa22cd416814592a4f9d9205e11affd5b8e10b" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "password-hash" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7676374caaee8a325c9e7a2ae557f216c5563a171d6997b0ef8a65af35147700" +dependencies = [ + "base64ct", + "rand_core", + "subtle", +] + +[[package]] +name = "pbkdf2" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83a0692ec44e4cf1ef28ca317f14f8f07da2d95ec3fa01f86e4467b725e60917" +dependencies = [ + "digest", + "hmac", + "password-hash", + "sha2", +] + +[[package]] +name = "pkg-config" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6b464fbc74e149a392436b17d523f769e057cb6877f6a5c4618bc6f11800548" + +[[package]] +name = "portable-atomic" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85" + +[[package]] +name = "portable-atomic-util" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" +dependencies = [ + "portable-atomic", +] + +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "pyo3" +version = "0.22.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f402062616ab18202ae8319da13fa4279883a2b8a9d9f83f20dbade813ce1884" +dependencies = [ + "cfg-if", + "indoc", + "libc", + "memoffset", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", + "unindent", +] + +[[package]] +name = "pyo3-build-config" +version = "0.22.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b14b5775b5ff446dd1056212d778012cbe8a0fbffd368029fd9e25b514479c38" +dependencies = [ + "once_cell", + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.22.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ab5bcf04a2cdcbb50c7d6105de943f543f9ed92af55818fd17b660390fc8636" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.22.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fd24d897903a9e6d80b968368a34e1525aeb719d568dba8b3d4bfa5dc67d453" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.22.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36c011a03ba1e50152b4b394b479826cad97e7a21eb52df179cd91ac411cbfbe" +dependencies = [ + "heck", + "proc-macro2", + "pyo3-build-config", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "rand" +version = "0.8.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e058c7de0b26af77780c769414d6257830bb240f3c38477dbc2c16e5f54d6d4c" +dependencies = [ + "libc", + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.17", +] + +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "safetensors" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d93279b86b3de76f820a8854dd06cbc33cfa57a417b19c47f6a25280112fb1df" +dependencies = [ + "serde", + "serde_json", +] + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.4", +] + +[[package]] +name = "serde_json" +version = "1.0.151" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha1" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a978451301f4db1d02937a4ab3ccce137717b81826e79b7d49ffe3244a13c3b8" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", + "sha2-asm", +] + +[[package]] +name = "sha2-asm" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b845214d6175804686b2bd482bcffe96651bb2d1200742b712003504a2dac1ab" +dependencies = [ + "cc", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "simd-adler32" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea" + +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6275cddf4610d1775e6d1fe9469b2e77d0f39fd98fb7450901b821e0c53649f" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "target-lexicon" +version = "0.12.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" + +[[package]] +name = "tch" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d3f84a069d8ba16dbf720b61e8bf131d90ffb8e958a664eae8e4993c5c2fa6f" +dependencies = [ + "half", + "lazy_static", + "libc", + "ndarray", + "rand", + "safetensors", + "thiserror", + "torch-sys", + "zip", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "time" +version = "0.3.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdb87b95ec50ddfa440816d227a17b2ccbdda963a316a727fda0fc4334f7d134" +dependencies = [ + "deranged", + "num-conv", + "powerfmt", + "serde_core", + "time-core", +] + +[[package]] +name = "time-core" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109" + +[[package]] +name = "torch-sys" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4ba78777379cf09aaa79708c63e477cf0f95e021d04360c6821f1a9f56173f7" +dependencies = [ + "anyhow", + "cc", + "libc", + "zip", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unindent" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "zerocopy" +version = "0.8.56" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "556764e583adb45a9f8d413c2a147fa7e8d821e48e12b14fd560b607998b75eb" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.56" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ab42fc20575779bd240faa45f94a74256f755c0fa9e89f0ede20d91d0cdfc1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zip" +version = "0.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "760394e246e4c28189f19d488c058bf16f564016aefac5d32bb1f3b51d5e9261" +dependencies = [ + "aes", + "byteorder", + "bzip2", + "constant_time_eq", + "crc32fast", + "crossbeam-utils", + "flate2", + "hmac", + "pbkdf2", + "sha1", + "time", + "zstd", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" + +[[package]] +name = "zstd" +version = "0.11.2+zstd.1.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20cc960326ece64f010d2d2107537f26dc589a6573a316bd5b1dba685fa5fde4" +dependencies = [ + "zstd-safe", +] + +[[package]] +name = "zstd-safe" +version = "5.0.2+zstd.1.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d2a5585e04f9eea4b2a3d1eca508c4dee9592a89ef6f450c11719da0726f4db" +dependencies = [ + "libc", + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.0.16+zstd.1.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748" +dependencies = [ + "cc", + "pkg-config", +] diff --git a/rust/mem-cache/Cargo.toml b/rust/mem-cache/Cargo.toml new file mode 100644 index 000000000..69057cb0e --- /dev/null +++ b/rust/mem-cache/Cargo.toml @@ -0,0 +1,40 @@ +[package] +name = "mem_cache" +version = "0.1.0" +edition = "2024" +description = "Rust tree core for the Unified Radix Cache" +license = "Apache-2.0" +repository = "https://github.com/sgl-project/sglang" + +[package.metadata.sglang] +python-module = "sglang.srt.mem_cache.rust_tree_core.mem_cache" +debug = false +features = ["python-extension"] +torch-compat-header = "torch_2_13_compat.h" + +[lib] +name = "mem_cache" +crate-type = ["cdylib"] + +# Keep this crate in its own workspace until its PyO3 version can be aligned +# with the root workspace without conflicting native `links = "python"` crates. +[workspace] + +[dependencies] +hashbrown = "0.16" +pyo3 = { version = "0.22", optional = true } +sha2 = "0.10" +tch = "=0.24.0" +thiserror = "1" + +# Hardware SHA extensions; identical digests. +[target.'cfg(all(target_arch = "aarch64", target_os = "linux"))'.dependencies] +sha2 = { version = "0.10", features = ["asm"] } + +[features] +# Keep the native core as the default so workspace tests do not link a Python +# extension. Wheel and source builds select python-extension through package +# metadata; shared white-box tests additionally select inspection. +default = [] +python-extension = ["dep:pyo3", "pyo3/extension-module", "tch/python-extension"] +inspection = [] diff --git a/rust/mem-cache/README.md b/rust/mem-cache/README.md new file mode 100644 index 000000000..6479e493a --- /dev/null +++ b/rust/mem-cache/README.md @@ -0,0 +1,42 @@ +# mem-cache + +Rust tree core for the Unified Radix Cache, covering Full attention, sliding window attention, and Mamba components. It implements the tree side of the `UnifiedTreeCoreInterface` split — match/insert walks, node arena, locks, eviction walks, HiCache backup/load-back specs, and KV events — behind a PyO3 binding, while the cache orchestration stays in Python. + +## Usage + +Select the backend with: + +```bash +SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND=rust +``` + +SGLang wheels bundle the production extension. A source checkout falls back to +the shared fingerprinted Rust-extension cache; it never writes a shared object +into the Python package. LibTorch and the Python headers come from the running +interpreter's PyTorch install. PyTorch 2.11 through 2.13 are accepted explicitly, +and `torch_2_13_compat.h` covers two alignment APIs removed in PyTorch 2.13. + +## Development + +```bash +# Build (libtorch from the installed torch package): +cd rust/mem-cache +LIBTORCH_USE_PYTORCH=1 \ + LIBTORCH_BYPASS_VERSION_CHECK=1 \ + CXXFLAGS="-include $PWD/torch_2_13_compat.h" \ + cargo build --release --locked --features python-extension + +# Native tests do not enable pyo3's extension-module feature: +TORCH_ROOT=$(python3 -c 'import pathlib, torch; print(pathlib.Path(torch.__file__).parent)') +LIBTORCH_USE_PYTORCH=1 LIBTORCH_BYPASS_VERSION_CHECK=1 \ + CXXFLAGS="-include $PWD/torch_2_13_compat.h" \ + LD_LIBRARY_PATH="$TORCH_ROOT/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" \ + cargo test --locked +``` + +The `inspection` Cargo feature adds white-box methods for the shared Python/Rust +cache suite. Production wheels do not enable it. + +Unit tests live in `src/tests/`, mirroring the source layout one file per module (wired via `#[cfg(test)] #[path = ...]`), so implementation files stay free of inline test blocks. + +Supported component sets are `[Full]`, `[Full, SWA]`, `[Full, Mamba]`, and `[Full, SWA, Mamba]`. diff --git a/rust/mem-cache/src/components/full.rs b/rust/mem-cache/src/components/full.rs new file mode 100644 index 000000000..be58d7158 --- /dev/null +++ b/rust/mem-cache/src/components/full.rs @@ -0,0 +1,493 @@ +//! FULL attention component driver: overrides the methods FULL customizes and inherits +//! the rest from the `TreeComponent` defaults. + +use std::cmp::Reverse; +use std::collections::{BinaryHeap, HashMap, HashSet}; + +use tch::{Kind, Tensor}; + +use crate::components::TreeComponent; +use crate::components::{ComponentType, FULL}; +use crate::node::ChildKeyType; +use crate::node::Node; +use crate::node::{NodeId, NodeIdx_, TreeCoreRuntimeError, ValueSlotIdx}; +use crate::unified_lru_list::PriorityKey; +use crate::unified_tree_core::{ + CacheAction, CacheTransferPhase, DecLockRefParams, EvictLayer, IncLockRefResult, InsertResult, + MatchPrefixParams, MatchResult, PoolName, PoolTransfer, PoolTransferResult, UnifiedTreeCore, +}; + +/// FULL attention component driver; owns the FULL device/host value slots. +pub struct FullComponent; + +impl FullComponent { + /// The component's device value slot. + pub const DEVICE: ValueSlotIdx = ValueSlotIdx::device(FULL); + /// The component's host value slot. + pub const HOST: ValueSlotIdx = ValueSlotIdx::host(FULL); +} + +impl TreeComponent for FullComponent { + fn component_type(&self) -> ComponentType { + FULL + } + + fn create_match_validator( + &self, + _tree_core: &UnifiedTreeCore, + match_device_only: bool, + ) -> Box, NodeIdx_) -> bool> { + // Device value present -> always a boundary; otherwise a backuped (host-resident) + // node also matches, unless the match is restricted to device. + Box::new(move |tree_core: &UnifiedTreeCore, node_id: NodeIdx_| { + let node = tree_core.arena.node(node_id); + node.has_device_value(FULL) || (!match_device_only && node.has_host_value(FULL)) + }) + } + + fn finalize_match_result_in_tree_core( + &self, + tree_core: &UnifiedTreeCore, + mut result: MatchResult, + params: &MatchPrefixParams<'_, K>, + value_chunks: &[Tensor], + best_value_len: usize, + ) -> MatchResult { + // Compute Full KV host hit length: walk from last_host_node up to + // last_device_node, summing host_value lengths of evicted nodes. + let mut kv_host_hit = 0; + let mut node_idx = tree_core.arena.resolve(result.best_match_node_id); + let last_device_idx = tree_core.arena.resolve(result.last_device_node_id); + while node_idx != last_device_idx { + let node = tree_core.arena.node(node_idx); + let parent = node.try_parent().unwrap_or_else(|| { + panic!( + "finalize walk from best_match_node {} hit root {} before \ + last_device_node {}", + result.best_match_node_id, node.id, result.last_device_node_id + ) + }); + kv_host_hit += node.host_value_len(FULL); + node_idx = parent; + } + if kv_host_hit > 0 { + result.host_hit_length = result.host_hit_length.max(kv_host_hit); + } + result + } + + fn redistribute_on_node_split( + &self, + tree_core: &mut UnifiedTreeCore, + new_parent_id: NodeIdx_, + child_id: NodeIdx_, + ) { + let (new_parent, child) = tree_core.arena.node_pair_mut(new_parent_id, child_id); + let split_len = new_parent.key.atom_len() as i64; + new_parent.copy_device_lock_ref(FULL, child); + if child.has_device_value(FULL) { + Node::redistribute_child_device_value(new_parent, child, FULL, split_len); + } + if child.has_host_value(FULL) { + Node::redistribute_child_host_value(new_parent, child, FULL, split_len); + } + } + + fn evict_component( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + target: EvictLayer, + ) -> (usize, usize) { + let node = tree_core.arena.node_mut(node_id); + let mut freed = 0; + let mut host_freed = 0; + if target.contains(EvictLayer::Device) && node.has_device_value(FULL) { + let value = node.device_value(FULL); + freed = node.device_value_len(FULL); + device_frees + .entry(FULL) + .or_default() + .push(value.shallow_clone()); + // NOTE: cd.value = None is deferred to _cascade_evict (Full as trigger) + // because SWA's free_swa still needs to read Full.value. + } + if target.contains(EvictLayer::Host) && node.has_host_value(FULL) { + host_freed = node.host_value_len(FULL); + host_frees + .entry(FULL) + .or_default() + .push(node.take_host_value(FULL)); + } + if freed > 0 { + tree_core.dec_evictable_size(FULL, freed); + } + (freed, host_freed) + } + + fn eviction_priority(&self, is_leaf: bool) -> i64 { + if is_leaf { 0 } else { 2 } + } + + fn evict_device_start(&self, tree_core: &mut UnifiedTreeCore, request_cnt: usize) { + tree_core.set_evict_device_start(FULL, request_cnt); + tree_core.full_evict_device_heap.clear(); + let arena = &tree_core.arena; + let strategy = &tree_core.eviction_strategy; + tree_core.full_evict_device_heap.extend( + tree_core + .evictable_device_leaves + .iter() + .map(|id| Reverse((strategy.get_priority(arena.node(id)), id))), + ); + } + + fn evict_device_next_node( + &self, + tree_core: &mut UnifiedTreeCore, + tracker: &mut HashMap, + _device_frees: &mut HashMap>, + _host_frees: &mut HashMap>, + ) -> Option { + let ct = FULL; + assert!( + tree_core.component_state(FULL).is_evict_device_ongoing, + "Full device eviction not started" + ); + // Re-admit the previously returned leaf's parent once it became a + // D-leaf; the parent id was captured at return time because the leaf + // itself may have been freed by the eviction in between. + if let Some(last_node_parent) = tree_core.component_state(FULL).evict_device_cursor + && tree_core.evictable_device_leaves.contains(last_node_parent) + { + let key = tree_core + .eviction_strategy + .get_priority(tree_core.arena.node(last_node_parent)); + tree_core + .full_evict_device_heap + .push(Reverse((key, last_node_parent))); + } + tree_core.component_state_mut(FULL).evict_device_cursor = None; + // The budget only advances between calls (the driver's evictions fill + // the tracker), so it gates the walk once up front. + if tracker[&ct] >= tree_core.component_state(FULL).evict_device_request_cnt { + return None; + } + while let Some(Reverse((_, x))) = tree_core.full_evict_device_heap.pop() { + if !tree_core.evictable_device_leaves.contains(x) { + continue; + } + let last_node_parent = tree_core.arena.node(x).try_parent(); + tree_core.component_state_mut(FULL).evict_device_cursor = last_node_parent; + return Some(x); + } + None + } + + fn evict_device_end(&self, tree_core: &mut UnifiedTreeCore) { + tree_core.set_evict_device_end(FULL); + tree_core.full_evict_device_heap.clear(); + } + + fn reclaim_coexisting_host_values( + &self, + tree_core: &mut UnifiedTreeCore, + num_tokens: usize, + tracker: &mut HashMap, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + ) { + for spare_imminent_demotes in [true, false] { + if tracker[&FULL] >= num_tokens { + break; + } + let candidates: Vec = tree_core.full_coexisting_host_nodes.iter().collect(); + for node_id in candidates { + if tracker[&FULL] >= num_tokens { + break; + } + let node = tree_core.arena.node(node_id); + if !node.has_device_value(FULL) || !node.has_host_value(FULL) { + tree_core.full_coexisting_host_nodes.discard(node_id); + continue; + } + if spare_imminent_demotes && tree_core.evictable_device_leaves.contains(node_id) { + continue; + } + if !tree_core.can_reclaim_coexisting_host_value_(node_id, FULL) { + continue; + } + tree_core.release_coexisting_host_value_( + node_id, + FULL, + tracker, + device_frees, + host_frees, + ); + tree_core.full_coexisting_host_nodes.discard(node_id); + } + } + } + + /// Evict host leaves to free KV host pool space. + fn drive_host_eviction( + &self, + tree_core: &mut UnifiedTreeCore, + num_tokens: usize, + tracker: &mut HashMap, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + ) { + let ct = FULL; + let arena = &tree_core.arena; + let strategy = &tree_core.eviction_strategy; + let mut heap: BinaryHeap> = tree_core + .evictable_host_leaves + .iter() + .map(|id| Reverse((strategy.get_priority(arena.node(id)), id))) + .collect(); + while tracker[&ct] < num_tokens { + let Some(Reverse((_, x))) = heap.pop() else { + break; + }; + if !tree_core.evictable_host_leaves.contains(x) { + continue; + } + // The parent id is captured before the eviction frees the leaf. + let parent = tree_core.arena.node(x).try_parent(); + tree_core.evict_host_leaf_(x, tracker, device_frees, host_frees); + if let Some(parent) = parent + && tree_core.evictable_host_leaves.contains(parent) + { + let key = tree_core + .eviction_strategy + .get_priority(tree_core.arena.node(parent)); + heap.push(Reverse((key, parent))); + } + } + } + + fn acquire_component_lock( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + mut result: IncLockRefResult, + lock_host: bool, + ) -> IncLockRefResult { + let ct = FULL; + + // Only the last host node needs to be protected. + if lock_host { + let node = tree_core.arena.node_mut(node_id); + // write_back mode: the anchor may be device-only (no host_value); pin it anyway. + if !node.has_host_value(FULL) && !tree_core.is_write_back { + return result; + } + node.inc_host_lock_ref(FULL); + tree_core.update_evictable_leaf_sets_(node_id); + return result; + } + + // Skip the bottom evicted segment, recording it for the matching release. + let on_boundary = |node: &Node| node.is_root() || node.has_device_value(FULL); + let mut cur = node_id; + let mut node = tree_core.arena.node(cur); + if !on_boundary(node) { + let skip_lock_node_ids = result.skip_lock_node_ids.entry(ct).or_default(); + loop { + skip_lock_node_ids.insert(node.id); + cur = node.parent(); + node = tree_core.arena.node(cur); + if on_boundary(node) { + break; + } + } + } + + // Lock the device-on segment up to the root. + let mut delta = 0; + loop { + let node = tree_core.arena.node_mut(cur); + if node.is_root() { + break; + } + assert!( + node.has_device_value(FULL), + "FULL invariant broken: evicted ancestor {cur} above device-on segment" + ); + let parent = node.parent(); + let newly_locked_len = if node.device_lock_ref(FULL) == 0 { + Some(node.device_value_len(FULL)) + } else { + None + }; + node.inc_device_lock_ref(FULL); + if let Some(key_len) = newly_locked_len { + tree_core.dec_evictable_size(FULL, key_len); + tree_core.inc_protected_size(FULL, key_len); + delta += key_len; + } + tree_core.evictable_device_leaves.discard(cur); + cur = parent; + } + result.delta = Some(delta); + result + } + + fn release_component_lock( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + params: Option<&DecLockRefParams>, + lock_host: bool, + ) { + let ct = FULL; + + if lock_host { + let node = tree_core.arena.node_mut(node_id); + if node.host_lock_ref(FULL) == 0 { + return; + } + // Mirror of `acquire`. write_back uses a pure counter. + if !node.has_host_value(FULL) && !tree_core.is_write_back { + return; + } + node.dec_host_lock_ref(FULL); + tree_core.update_evictable_leaf_sets_(node_id); + return; + } + + let empty = HashSet::new(); + let skip_lock_node_ids = params + .and_then(|p| p.skip_lock_node_ids.get(&ct)) + .unwrap_or(&empty); + let mut cur = node_id; + loop { + let node = tree_core.arena.node_mut(cur); + if node.is_root() { + break; + } + let parent = node.parent(); + if skip_lock_node_ids.contains(&node.id) { + cur = parent; + continue; + } + assert!( + node.has_device_value(FULL), + "release_component_lock: node {cur} has no FULL device value" + ); + let old_lock_ref = node.device_lock_ref(FULL); + assert!( + old_lock_ref > 0, + "release_component_lock: node {cur} is not locked" + ); + let newly_unlocked_len = if old_lock_ref == 1 { + Some(node.device_value_len(FULL)) + } else { + None + }; + node.dec_device_lock_ref(FULL); + if let Some(key_len) = newly_unlocked_len { + tree_core.dec_protected_size(FULL, key_len); + tree_core.inc_evictable_size(FULL, key_len); + tree_core.update_evictable_leaf_sets_(cur); + } + cur = parent; + } + } + + fn build_hicache_transfers( + &self, + tree_core: &UnifiedTreeCore, + node_id: NodeIdx_, + phase: CacheTransferPhase, + _mamba_pool_idx: Option, + _host_indices: Option, + _token_ids: Option<&[i64]>, + _prefetch_tokens: usize, + _last_hash: Option<&str>, + ) -> Result>, TreeCoreRuntimeError> { + Ok(match phase { + // Full KV backup is handled by the main flow + // (cache_controller.write on host_value directly). + // No extra PoolTransfer needed. + CacheTransferPhase::BackupHost => None, + CacheTransferPhase::LoadBack => { + // `node` is best_match_node. FULL device evict only from leaves, + // so once we hit a device-on node, everything above is also device-on. + let mut backed_up: Vec = Vec::new(); + let mut nodes_to_load: Vec = Vec::new(); + let mut cur = tree_core.arena.node(node_id); + while cur.evicted() { + backed_up.push(cur.host_value(FULL).shallow_clone()); + nodes_to_load.push(cur.id); + cur = tree_core.arena.node(cur.parent()); + } + backed_up.reverse(); + nodes_to_load.reverse(); + let host_indices = if backed_up.is_empty() { + Tensor::empty([0], (Kind::Int64, tch::Device::Cpu)) + } else { + Tensor::cat(&backed_up, 0) + }; + Some(vec![PoolTransfer { + name: PoolName::Kv, + host_indices: Some(host_indices), + nodes_to_load: Some(nodes_to_load), + ..Default::default() + }]) + } + CacheTransferPhase::BackupStorage | CacheTransferPhase::Prefetch => None, + }) + } + + fn commit_hicache_transfer( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + phase: CacheTransferPhase, + transfers: Vec, + cache_actions: &mut Vec, + insert_result: Option<&mut InsertResult>, + pool_storage_result: Option<&PoolTransferResult>, + ) { + match phase { + CacheTransferPhase::BackupHost => { + if let Some(transfer) = transfers.first() + && let Some(host_indices) = &transfer.host_indices + { + tree_core + .arena + .set_host_value(node_id, FULL, host_indices.copy()); + } + } + CacheTransferPhase::LoadBack => { + if let Some(transfer) = transfers.first() + && let Some(device_indices) = &transfer.device_indices + { + let mut offset = 0i64; + for &loaded_id in transfer.nodes_to_load.iter().flatten() { + let loaded_idx = tree_core.arena.resolve(loaded_id); + let loaded = tree_core.arena.node_mut(loaded_idx); + let n_len = loaded.host_value_len(FULL) as i64; + loaded + .set_device_value(FULL, device_indices.narrow(0, offset, n_len).copy()); + offset += n_len; + // Full uses leaf sets, not LRU. + tree_core.inc_evictable_size(FULL, n_len as usize); + tree_core.update_evictable_leaf_sets_(loaded_idx); + } + } + tree_core.update_evictable_leaf_sets_(node_id); + } + // The Full component has no storage-phase commits. + CacheTransferPhase::BackupStorage | CacheTransferPhase::Prefetch => {} + } + } +} + +#[cfg(test)] +#[path = "../tests/components/full.rs"] +mod tests; diff --git a/rust/mem-cache/src/components/mamba.rs b/rust/mem-cache/src/components/mamba.rs new file mode 100644 index 000000000..d4f3c0749 --- /dev/null +++ b/rust/mem-cache/src/components/mamba.rs @@ -0,0 +1,751 @@ +//! Mamba (SSM state) component driver: overrides the methods Mamba customizes +//! and inherits the rest from the `TreeComponent` defaults. +//! Mamba data is per-leaf single-slot state; sizes count slots, not tokens. + +use std::collections::HashMap; + +use tch::Tensor; + +use crate::components::TreeComponent; +use crate::components::{ComponentType, MAMBA}; +use crate::node::ChildKeyType; +use crate::node::Node; +use crate::node::{NodeId, NodeIdx_, TreeCoreRuntimeError, ValueSlotIdx}; +use crate::unified_tree_core::{ + CacheAction, CacheInitParams, CacheTransferPhase, DecLockRefParams, EvictLayer, + IncLockRefResult, InsertParams, InsertResult, LRURefreshPhase, MatchPrefixParams, MatchResult, + PoolHitPolicy, PoolName, PoolTransfer, PoolTransferResult, UnifiedTreeCore, +}; + +/// Mamba component driver; owns the Mamba device/host value slots. +pub struct MambaComponent { + /// Joint chunk/tree-page alignment for the mamba branching seqlen. + mamba_checkpoint_grid: usize, + /// Per-root-path cap on cached Mamba states; None means unlimited. + mamba_max_states_per_path: Option, +} + +impl MambaComponent { + /// The component's device value slot. + pub const DEVICE: ValueSlotIdx = ValueSlotIdx::device(MAMBA); + /// The component's host value slot. + pub const HOST: ValueSlotIdx = ValueSlotIdx::host(MAMBA); +} + +impl MambaComponent { + /// Build the driver from the tree's init params. + pub fn new(params: &CacheInitParams) -> Self { + let mamba_cache_chunk_size = params + .mamba_cache_chunk_size + .expect("the Mamba component requires mamba_cache_chunk_size"); + MambaComponent { + // A donated checkpoint must land on both the model's chunk grid and + // a radix-node boundary. `params.page_size` is already widened by DCP. + mamba_checkpoint_grid: least_common_multiple(mamba_cache_chunk_size, params.page_size), + mamba_max_states_per_path: params.mamba_max_states_per_path, + } + } +} + +fn least_common_multiple(lhs: usize, rhs: usize) -> usize { + let mut a = lhs; + let mut b = rhs; + while b != 0 { + (a, b) = (b, a % b); + } + lhs / a * rhs +} + +impl MambaComponent { + // Tier-selected mamba slot read for the lock paths; `host` picks the host slot. + fn has_value(node: &Node, host: bool) -> bool { + if host { + node.has_host_value(MAMBA) + } else { + node.has_device_value(MAMBA) + } + } + + /// Defer the path-cap eviction so it runs after the insert's BackupKV. + fn emit_excess_path_states_eviction_( + &self, + tail_node_id: NodeId, + cache_actions: &mut Vec, + ) { + if self.mamba_max_states_per_path.is_none() { + return; + } + cache_actions.push(CacheAction::MambaEvictExcessPathStates { tail_node_id }); + } +} + +impl TreeComponent for MambaComponent { + fn component_type(&self) -> ComponentType { + MAMBA + } + + fn needs_incremental_backup(&self, tree_core: &UnifiedTreeCore, node_id: NodeIdx_) -> bool { + let node = tree_core.arena.node(node_id); + node.has_device_value(MAMBA) && !node.has_host_value(MAMBA) + } + + /// A match consumes only the best-match node's mamba state, so MATCH_END + /// touches just that node; new-leaf states enter the LRU at insert commit, + /// so WALKDOWN and INSERT_END are no-ops. + fn refresh_lru( + &self, + tree_core: &mut UnifiedTreeCore, + phase: LRURefreshPhase, + node_id: NodeIdx_, + ) { + match phase { + LRURefreshPhase::Walkdown => {} + LRURefreshPhase::MatchEnd => { + if tree_core.arena.has_device_value(node_id, MAMBA) { + tree_core.device_lru_list_mut(MAMBA).reset_node_mru(node_id); + } + } + LRURefreshPhase::InsertEnd => {} + } + } + + /// A per-match predicate accepting nodes that hold mamba data. + fn create_match_validator( + &self, + _tree_core: &UnifiedTreeCore, + match_device_only: bool, + ) -> Box, NodeIdx_) -> bool> { + // HiCache: evicted + backuped (host_value present) is also a valid match. + Box::new(move |tree_core: &UnifiedTreeCore, node_id: NodeIdx_| { + let node = tree_core.arena.node(node_id); + node.has_device_value(MAMBA) || (!match_device_only && node.has_host_value(MAMBA)) + }) + } + + /// The mamba branching seqlen and the host-only hit bump. + fn finalize_match_result_in_tree_core( + &self, + tree_core: &UnifiedTreeCore, + mut result: MatchResult, + _params: &MatchPrefixParams<'_, K>, + _value_chunks: &[Tensor], + _best_value_len: usize, + ) -> MatchResult { + let mamba_boundary_len = result.device_indices.size()[0] as usize + result.host_hit_length; + + // Full KV may extend beyond the latest reusable Mamba state. The branching + // point is the last checkpoint-grid-aligned position within the Full-KV hit + // that lies beyond the current Mamba boundary. + let aligned_seqlen = + result.full_kv_hit_length / self.mamba_checkpoint_grid * self.mamba_checkpoint_grid; + result.mamba_branching_seqlen = + (aligned_seqlen > mamba_boundary_len).then_some(aligned_seqlen); + + // HiCache: if mamba was evicted from device but has host backup, + // ensure mamba_host_hit_length >= 1 so load_back is triggered. + let last_node = tree_core + .arena + .node(tree_core.arena.resolve(result.best_match_node_id)); + if !last_node.has_device_value(MAMBA) && last_node.has_host_value(MAMBA) { + result.mamba_host_hit_length = result.mamba_host_hit_length.max(1); + } + result + } + + /// Attach the donated mamba slot to the insert target leaf. + fn commit_insert_component_data( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + is_new_leaf: bool, + params: &InsertParams<'_, K>, + result: &mut InsertResult, + cache_actions: &mut Vec, + ) { + let mamba_value = params + .mamba_value + .as_ref() + .expect("mamba insert requires a donated mamba_value"); + let slot_len = mamba_value.size()[0] as usize; + + if is_new_leaf { + tree_core + .arena + .set_device_value(node_id, MAMBA, mamba_value.shallow_clone()); + tree_core.device_lru_list_mut(MAMBA).insert_mru(node_id); + tree_core.inc_evictable_size(MAMBA, slot_len); + self.emit_excess_path_states_eviction_(tree_core.arena.node(node_id).id, cache_actions); + return; + } + if !tree_core.arena.has_device_value(node_id, MAMBA) { + // Tombstone refill: the node moves from the host LRU to the device LRU. + tree_core + .arena + .set_device_value(node_id, MAMBA, mamba_value.shallow_clone()); + let host_lru = tree_core.host_lru_list_mut(MAMBA); + if host_lru.in_list(Some(node_id)) { + host_lru.remove_node(node_id); + } + tree_core.device_lru_list_mut(MAMBA).insert_mru(node_id); + tree_core.inc_evictable_size(MAMBA, slot_len); + let tick = tree_core.arena.get_and_bump_access_counter(); + tree_core.arena.node_mut(node_id).last_access_counter = tick; + self.emit_excess_path_states_eviction_(tree_core.arena.node(node_id).id, cache_actions); + return; + } + tree_core.device_lru_list_mut(MAMBA).reset_node_mru(node_id); + let tick = tree_core.arena.get_and_bump_access_counter(); + tree_core.arena.node_mut(node_id).last_access_counter = tick; + result.mamba_exist = true; + } + + /// Mamba data stays on the original leaf; the new prefix node gets none. + /// Evict shallow Mamba device checkpoints beyond the per-path cap on the + /// tail's root path; Full KV, host backups, the tail, forks, locked nodes, + /// and device leaves are preserved (a best-effort soft cap). + fn evict_excess_path_states( + &self, + tree_core: &mut UnifiedTreeCore, + tail_node_id: NodeIdx_, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + ) { + let Some(cap) = self.mamba_max_states_per_path else { + return; + }; + // Mamba-value holders on the root path, tail-first. + let mut holders: Vec = Vec::new(); + let mut cursor = Some(tail_node_id); + while let Some(node_id) = cursor { + let node = tree_core.arena.node(node_id); + if node.is_root() { + break; + } + if node.has_device_value(MAMBA) { + holders.push(node_id); + } + cursor = node.parent; + } + let mut excess = holders.len().saturating_sub(cap); + if excess == 0 { + return; + } + // Cache-level apply: the counts are not reported, only the frees. + let mut tracker: HashMap = HashMap::new(); + for &node_id in holders.iter().rev() { + if excess == 0 || node_id == tail_node_id { + break; + } + let node = tree_core.arena.node(node_id); + if node.device_lock_ref(MAMBA) > 0 || node.children.len() != 1 { + continue; + } + if tree_core.evictable_device_leaves.contains(node_id) { + continue; + } + tree_core.evict_component_and_detach_lru_( + node_id, + MAMBA, + device_frees, + host_frees, + EvictLayer::Device, + Some(&mut tracker), + ); + tree_core.cascade_evict_( + node_id, + MAMBA, + &mut tracker, + device_frees, + host_frees, + EvictLayer::Device, + ); + excess -= 1; + } + } + + fn redistribute_on_node_split( + &self, + tree_core: &mut UnifiedTreeCore, + new_parent_id: NodeIdx_, + _child_id: NodeIdx_, + ) { + let new_parent = tree_core.arena.node_mut(new_parent_id); + if new_parent.has_device_value(MAMBA) { + let _ = new_parent.take_device_value(MAMBA); + } + new_parent.set_lock_ref_(ValueSlotIdx::device(MAMBA), 0); + // HiCache: mamba host_value stays on child (mamba = leaf-only data). + if new_parent.has_host_value(MAMBA) { + let _ = new_parent.take_host_value(MAMBA); + } + new_parent.set_lock_ref_(ValueSlotIdx::host(MAMBA), 0); + } + + /// Free the node's mamba slot on the targeted layer(s). + fn evict_component( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + target: EvictLayer, + ) -> (usize, usize) { + let ct = MAMBA; + let node = tree_core.arena.node_mut(node_id); + let mut freed = 0; + let mut host_freed = 0; + + // Device layer + if target.contains(EvictLayer::Device) && node.has_device_value(MAMBA) { + freed = node.device_value_len(MAMBA); + device_frees + .entry(ct) + .or_default() + .push(node.take_device_value(MAMBA)); + tree_core.dec_evictable_size(MAMBA, freed); + } + + // Host layer + let node = tree_core.arena.node_mut(node_id); + if target.contains(EvictLayer::Host) && node.has_host_value(MAMBA) { + host_freed = node.host_value_len(MAMBA); + host_frees + .entry(ct) + .or_default() + .push(node.take_host_value(MAMBA)); + let host_lru = tree_core.host_lru_list_mut(MAMBA); + if host_lru.in_list(Some(node_id)) { + host_lru.remove_node(node_id); + } + } + + // After device tombstone: if only host_value remains, insert into host LRU + let node = tree_core.arena.node(node_id); + if target == EvictLayer::Device + && !node.has_device_value(MAMBA) + && node.has_host_value(MAMBA) + { + let host_lru = tree_core.host_lru_list_mut(MAMBA); + if !host_lru.in_list(Some(node_id)) { + host_lru.insert_mru(node_id); + } + } + + (freed, host_freed) + } + + /// Begin the device-eviction walk from this component's LRU cursor. + fn evict_device_start(&self, tree_core: &mut UnifiedTreeCore, request_cnt: usize) { + tree_core.set_evict_device_start(MAMBA, request_cnt); + let cursor = tree_core + .device_lru_list(MAMBA) + .get_lru_no_lock(&tree_core.arena); + tree_core.component_state_mut(MAMBA).evict_device_cursor = cursor; + } + + /// Advance one device-eviction step and return a leaf, if selected. + /// + /// An internal tombstone is one complete step so the caller can apply its + /// pending frees and recheck allocator capacity before the next mutation. + fn evict_device_next_node( + &self, + tree_core: &mut UnifiedTreeCore, + tracker: &mut HashMap, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + ) -> Option { + let ct = MAMBA; + assert!( + tree_core.component_state(MAMBA).is_evict_device_ongoing, + "Mamba device eviction not started" + ); + let mut cursor = tree_core.component_state(MAMBA).evict_device_cursor; + // The cursor is re-validated (reset to LRU head) if the previous + // node's eviction removed it. + if cursor.is_some_and(|c| !tree_core.device_lru_list(MAMBA).in_list(Some(c))) { + cursor = tree_core + .device_lru_list(MAMBA) + .get_lru_no_lock(&tree_core.arena); + } + let next = loop { + if tracker[&ct] >= tree_core.component_state(MAMBA).evict_device_request_cnt { + break None; + } + let Some(x) = cursor else { + break None; + }; + if !tree_core.device_lru_list(MAMBA).in_list(Some(x)) { + break None; + } + assert!( + tree_core.arena.has_device_value(x, MAMBA), + "Mamba eviction cursor on a valueless node {x}" + ); + cursor = tree_core + .device_lru_list(MAMBA) + .get_prev_no_lock(x, &tree_core.arena); + // A load-back pin means an in-flight DMA targets this node's slices. + if tree_core.arena.node(x).is_load_back_pending() { + continue; + } + if tree_core.evictable_device_leaves.contains(x) { + break Some(x); + } + // Internal nodes are tombstoned inline (no IO). + tree_core.evict_component_and_detach_lru_( + x, + ct, + device_frees, + host_frees, + EvictLayer::Device, + Some(tracker), + ); + tree_core.cascade_evict_(x, ct, tracker, device_frees, host_frees, EvictLayer::Device); + break None; + }; + tree_core.component_state_mut(MAMBA).evict_device_cursor = cursor; + next + } + + /// Clear the device-eviction walk cursor state. + fn evict_device_end(&self, tree_core: &mut UnifiedTreeCore) { + tree_core.set_evict_device_end(MAMBA); + } + + /// Single-node mamba lock; host locks also detach from the host LRU. + fn acquire_component_lock( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + mut result: IncLockRefResult, + lock_host: bool, + ) -> IncLockRefResult { + let node = tree_core.arena.node(node_id); + if node.is_root() { + return result; + } + // A node in skip_lock_node_ids was a tombstone when this lock was acquired. + if !Self::has_value(node, lock_host) { + result + .skip_lock_node_ids + .entry(MAMBA) + .or_default() + .insert(node.id); + return result; + } + if lock_host { + if node.host_lock_ref(MAMBA) == 0 { + let host_lru = tree_core.host_lru_list_mut(MAMBA); + if host_lru.in_list(Some(node_id)) { + host_lru.remove_node(node_id); + } + } + tree_core.arena.inc_host_lock_ref(node_id, MAMBA); + } else { + let value_len = node.device_value_len(MAMBA); + if node.device_lock_ref(MAMBA) == 0 { + tree_core.dec_evictable_size(MAMBA, value_len); + tree_core.inc_protected_size(MAMBA, value_len); + } + tree_core.arena.inc_device_lock_ref(node_id, MAMBA); + } + result + } + + /// Single-node mamba unlock; host unlocks reinsert into the host LRU. + fn release_component_lock( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + params: Option<&DecLockRefParams>, + lock_host: bool, + ) { + if tree_core.arena.node(node_id).is_root() { + return; + } + if let Some(params) = params + && params + .skip_lock_node_ids + .get(&MAMBA) + .is_some_and(|ids| ids.contains(&tree_core.arena.node(node_id).id)) + { + return; + } + if lock_host { + let node = tree_core.arena.node_mut(node_id); + node.dec_host_lock_ref(MAMBA); + if node.host_lock_ref(MAMBA) == 0 + && !node.has_device_value(MAMBA) + && node.has_host_value(MAMBA) + { + let host_lru = tree_core.host_lru_list_mut(MAMBA); + if !host_lru.in_list(Some(node_id)) { + host_lru.insert_mru(node_id); + } + } + return; + } + let node = tree_core.arena.node(node_id); + let device_lock_ref = node.device_lock_ref(MAMBA); + if device_lock_ref > 0 { + if device_lock_ref == 1 { + let value_len = node.device_value_len(MAMBA); + tree_core.inc_evictable_size(MAMBA, value_len); + tree_core.dec_protected_size(MAMBA, value_len); + } + tree_core.arena.dec_device_lock_ref(node_id, MAMBA); + } + } + + /// Build the mamba transfer descriptors for the given phase. + fn build_hicache_transfers( + &self, + tree_core: &UnifiedTreeCore, + node_id: NodeIdx_, + phase: CacheTransferPhase, + mamba_pool_idx: Option, + host_indices: Option, + _token_ids: Option<&[i64]>, + _prefetch_tokens: usize, + _last_hash: Option<&str>, + ) -> Result>, TreeCoreRuntimeError> { + Ok(match phase { + CacheTransferPhase::BackupHost => { + let node = tree_core.arena.node(node_id); + if node.has_host_value(MAMBA) { + return Ok(None); + } + node.try_device_value(MAMBA).map(|value| { + vec![PoolTransfer { + name: PoolName::Mamba, + device_indices: Some(value.shallow_clone()), + ..Default::default() + }] + }) + } + CacheTransferPhase::LoadBack => { + let node = tree_core.arena.node(node_id); + if node.has_device_value(MAMBA) { + return Ok(None); + } + let mut transfers = Vec::new(); + // restore single node if host_value exists + if let Some(host_value) = node.try_host_value(MAMBA) { + transfers.push(PoolTransfer { + name: PoolName::Mamba, + host_indices: Some(host_value.shallow_clone()), + nodes_to_load: Some(vec![node.id]), + ..Default::default() + }); + } + // Per-request mamba CoW (H->D copy into the request's device slot, + // pre-allocated on caller side). + if let (Some(mamba_pool_idx), Some(host_value)) = + (mamba_pool_idx, node.try_host_value(MAMBA)) + { + transfers.push(PoolTransfer { + name: PoolName::Mamba, + host_indices: Some(host_value.shallow_clone()), + device_indices: Some(mamba_pool_idx.unsqueeze(0)), + ..Default::default() + }); + } + if transfers.is_empty() { + None + } else { + Some(transfers) + } + } + CacheTransferPhase::BackupStorage => { + let node = tree_core.arena.node(node_id); + let Some(host_value) = node.try_host_value(MAMBA) else { + return Ok(None); + }; + let Some(hash_value) = node.hash_value.as_ref().filter(|h| !h.is_empty()) else { + return Ok(None); + }; + Some(vec![PoolTransfer { + name: PoolName::Mamba, + host_indices: Some(host_value.shallow_clone()), + keys: Some(vec![hash_value[hash_value.len() - 1].clone()]), + hit_policy: PoolHitPolicy::TrailingPages, + ..Default::default() + }]) + } + CacheTransferPhase::Prefetch => { + let host_indices = + host_indices.expect("Mamba PREFETCH build requires host indices"); + Some(vec![PoolTransfer { + name: PoolName::Mamba, + host_indices: Some(host_indices), + keys: Some(vec!["__placeholder__".to_string()]), + hit_policy: PoolHitPolicy::TrailingPages, + ..Default::default() + }]) + } + }) + } + + /// Post-transfer mamba bookkeeping for the given phase. + fn commit_hicache_transfer( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + phase: CacheTransferPhase, + transfers: Vec, + cache_actions: &mut Vec, + insert_result: Option<&mut InsertResult>, + pool_storage_result: Option<&PoolTransferResult>, + ) { + match phase { + CacheTransferPhase::BackupHost => { + if let Some(transfer) = transfers.first() + && let Some(host_indices) = &transfer.host_indices + { + let node = tree_core.arena.node_mut(node_id); + if !node.has_host_value(MAMBA) { + node.set_host_value(MAMBA, host_indices.copy()); + } + } + } + CacheTransferPhase::LoadBack => { + let Some(transfer) = transfers.first() else { + return; + }; + if let Some(device_indices) = &transfer.device_indices { + let node = tree_core.arena.node_mut(node_id); + node.set_device_value(MAMBA, device_indices.copy()); + let count = node.device_value_len(MAMBA); + // Move from host LRU to device LRU + let host_lru = tree_core.host_lru_list_mut(MAMBA); + if host_lru.in_list(Some(node_id)) { + host_lru.remove_node(node_id); + } + tree_core.device_lru_list_mut(MAMBA).insert_mru(node_id); + tree_core.inc_evictable_size(MAMBA, count); + } + } + // The python elif chain has no BACKUP_STORAGE arm. + CacheTransferPhase::BackupStorage => {} + CacheTransferPhase::Prefetch => { + let Some(transfer) = transfers.first() else { + return; + }; + let host_indices = transfer.host_indices.as_ref(); + let loaded = pool_storage_result.is_some_and(|result| { + result + .extra_pool_hit_pages + .get(&PoolName::Mamba) + .copied() + .unwrap_or(0) + >= 1 + }); + let target_node_id = insert_result + .as_deref() + .and_then(|result| result.inserted_host_node) + .map(|id| tree_core.arena.resolve(id)); + let attach_target = match (host_indices, target_node_id) { + (Some(_), Some(target)) + if loaded && !tree_core.arena.has_host_value(target, MAMBA) => + { + Some(target) + } + _ => None, + }; + let Some(target) = attach_target else { + // The buffer cannot attach: free it and let the caller keep + // its own donated slot bookkeeping. + cache_actions.push(CacheAction::FreeComponentHostSlot { + component_type: MAMBA, + host_indices: host_indices + .map(|host| vec![host.shallow_clone()]) + .unwrap_or_default(), + }); + if let Some(insert_result) = insert_result { + insert_result.mamba_exist = true; + } + return; + }; + let host_indices = host_indices.expect("an attach target implies host indices"); + tree_core + .arena + .set_host_value(target, MAMBA, host_indices.copy()); + if !tree_core.arena.has_device_value(target, MAMBA) { + let host_lru = tree_core.host_lru_list_mut(MAMBA); + if !host_lru.in_list(Some(target)) { + host_lru.insert_mru(target); + } + } + if let Some(insert_result) = insert_result { + insert_result.mamba_exist = false; + } + } + } + } + + /// Evict mamba host resources: internal nodes tombstone privately, host + /// leaves evict atomically. + fn drive_host_eviction( + &self, + tree_core: &mut UnifiedTreeCore, + num_tokens: usize, + tracker: &mut HashMap, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + ) { + let ct = MAMBA; + let mut x = tree_core + .host_lru_list(MAMBA) + .get_lru_no_lock(&tree_core.arena); + loop { + if tracker[&ct] >= num_tokens { + break; + } + let Some(cur) = x else { + break; + }; + if !tree_core.host_lru_list(MAMBA).in_list(Some(cur)) { + break; + } + let x_next = tree_core + .host_lru_list(MAMBA) + .get_prev_no_lock(cur, &tree_core.arena); + // A load-back pin means an in-flight DMA reads this node's host slices. + if tree_core.arena.node(cur).is_load_back_pending() { + x = x_next; + continue; + } + if tree_core.evictable_host_leaves.contains(cur) { + // Host leaf: atomic eviction (all components host + delete) + tree_core.evict_host_leaf_(cur, tracker, device_frees, host_frees); + } else { + // Internal: tombstone Mamba + cascade + assert!( + tree_core.arena.has_host_value(cur, MAMBA), + "Mamba host LRU member {cur} has no host value" + ); + tree_core.evict_component_and_detach_lru_( + cur, + ct, + device_frees, + host_frees, + EvictLayer::Host, + Some(tracker), + ); + tree_core.cascade_evict_( + cur, + ct, + tracker, + device_frees, + host_frees, + EvictLayer::Host, + ); + } + x = x_next; + } + } +} + +#[cfg(test)] +#[path = "../tests/components/mamba.rs"] +mod tests; diff --git a/rust/mem-cache/src/components/mod.rs b/rust/mem-cache/src/components/mod.rs new file mode 100644 index 000000000..b768b0917 --- /dev/null +++ b/rust/mem-cache/src/components/mod.rs @@ -0,0 +1,491 @@ +//! Per-component drivers; each receives the whole `UnifiedTreeCore` for backward access. +#![allow(unused_variables)] + +use std::collections::HashMap; + +use tch::Tensor; + +use crate::node::{ChildKeyType, NodeArena, NodeIdx_, TreeCoreRuntimeError}; +use crate::unified_tree_core::{ + CacheAction, CacheTransferPhase, DecLockRefParams, EvictLayer, IncLockRefResult, InsertParams, + InsertResult, LRURefreshPhase, MatchPrefixParams, MatchResult, PoolTransfer, + PoolTransferResult, UnifiedTreeCore, +}; + +mod full; +mod mamba; +mod swa; + +pub use full::FullComponent; +pub use mamba::MambaComponent; +pub use swa::SwaComponent; + +/// Whether `node_id` holds the component's data on `target`, checking its +/// device or host slot. +pub(crate) fn node_has_component_data( + arena: &NodeArena, + node_id: NodeIdx_, + component_type: ComponentType, + target: EvictLayer, +) -> bool { + match target { + EvictLayer::Device => arena.has_device_value(node_id, component_type), + EvictLayer::Host => arena.has_host_value(node_id, component_type), + EvictLayer::All => panic!("node_has_component_data: EvictLayer::All is not a single layer"), + } +} + +/// Every device value of the component across all roots, concatenated. +pub(crate) fn all_values_flatten( + tree_core: &UnifiedTreeCore, + component_type: ComponentType, +) -> Tensor { + let mut values: Vec = Vec::new(); + let mut stack: Vec = vec![tree_core.arena.root()]; + while let Some(node_id) = stack.pop() { + let node = tree_core.arena.node(node_id); + if let Some(value) = node.try_device_value(component_type) { + values.push(value.shallow_clone()); + } + stack.extend(node.children.values().copied()); + } + if values.is_empty() { + return tree_core.empty_device_indices.shallow_clone(); + } + Tensor::cat(&values, 0) +} + +/// A per-component lock/value/eviction driver over the shared `UnifiedTreeCore`. +pub trait TreeComponent { + /// The component this driver serves. + fn component_type(&self) -> ComponentType; + + /// Whether this component has device data that still needs a host backup. + fn needs_incremental_backup( + &self, + _tree_core: &UnifiedTreeCore, + _node_id: NodeIdx_, + ) -> bool { + false + } + + /// Refresh this component's LRU position for `node_id` at the given walk phase. + fn refresh_lru( + &self, + tree_core: &mut UnifiedTreeCore, + phase: LRURefreshPhase, + node_id: NodeIdx_, + ) { + // Python reference — tree_component.py::TreeComponent.refresh_lru: + // def refresh_lru( + // self, + // phase: LRURefreshPhase, + // node: UnifiedTreeNode, + // root_node: UnifiedTreeNode, + // ) -> None: + // ct = self.component_type + // match phase: + // case LRURefreshPhase.WALKDOWN: + // if node.component_data[ct].value is None: + // return + // self.tree_core.lru_lists[ct].reset_node_mru(node) + // case LRURefreshPhase.MATCH_END: + // self.tree_core.lru_lists[ct].reset_node_and_parents_mru( + // node, root_node, self.node_has_component_data + // ) + // case LRURefreshPhase.INSERT_END: + // # WALKDOWN already refreshed every node on the insert path + // # (including the new leaf), so there is nothing more to do. + // return + // case _: + // raise ValueError(f"Unknown LRURefreshPhase: {phase}") + unimplemented!("TreeComponent.refresh_lru") + } + + /// Return a per-match stateful predicate deciding whether a node is a valid + /// match boundary for this component. + // Python reference — tree_component.py::TreeComponent.create_match_validator: + // @abstractmethod + // def create_match_validator( + // self, match_device_only: bool = False + // ) -> Callable[[UnifiedTreeNode], bool]: + // """Return a per-match stateful predicate that decides whether a node + // is a valid match boundary for this component. + // Called once per match_prefix; the returned closure may carry state. + // When match_device_only is true, host-backed nodes must not be accepted + // as valid match boundaries. + // - Full: returns True if the node has full component data. + // - SWA: tracks accumulated length since last gap; returns True only + // when the contiguous window reaches swa_sliding_window_size. + // - Mamba: returns True iff the node has mamba component data.""" + // ... + fn create_match_validator( + &self, + tree_core: &UnifiedTreeCore, + match_device_only: bool, + ) -> Box, NodeIdx_) -> bool>; + + /// Tree-side post-processing inside the match walk (no cache access). + fn finalize_match_result_in_tree_core( + &self, + tree_core: &UnifiedTreeCore, + result: MatchResult, + params: &MatchPrefixParams<'_, K>, + value_chunks: &[Tensor], + best_value_len: usize, + ) -> MatchResult { + result + } + + /// Called per-node when an insert's key overlaps an existing node. + /// Returns the index within `value_slice` from which this component + /// consumed (took ownership of) the underlying KV pool slots. + /// Returns `prefix_len` if nothing was consumed (default). + /// The insert walk uses this to free only the non-consumed duplicate + /// portion: `value_slice[dup_start..consumed_from]`. + fn update_component_on_insert_overlap( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + prefix_len: usize, + total_prefix_len: usize, + value_slice: Tensor, + params: &InsertParams<'_, K>, + result: &mut InsertResult, + cache_actions: &mut Vec, + ) -> usize { + prefix_len + } + + /// Called after `unevict_node_on_insert_` restores the base (Full) value + /// on an evicted node. Aux components (e.g. SWA) override this to rebuild + /// their own data from the freshly assigned base value when their entry + /// is still tombstoned. Default no-op. + fn recover_after_unevict( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + prefix_len: usize, + total_prefix_len: usize, + params: &InsertParams<'_, K>, + result: &mut InsertResult, + cache_actions: &mut Vec, + ) { + } + + /// Finalize component data on the target (leaf) node after the insert + /// walk completes. Called once per insert. + /// - Full: no-op (full data is handled by `add_new_node_`). + /// - SWA: for new leaves, checks whether the node straddles the SWA + /// eviction boundary (`swa_evicted_seqlen`). If so, splits the node + /// via `split_node_` — the parent becomes a tombstone (no SWA) and the + /// child (the deeper portion) receives SWA data. If the entire node + /// is within the window, sets SWA directly. If entirely outside, + /// leaves SWA as None (tombstone). + /// - Mamba: sets the mamba component value from params, inserts into the + /// mamba LRU list, and increments evictable size. If the node already + /// has mamba data, resets its LRU position instead. + fn commit_insert_component_data( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + is_new_leaf: bool, + params: &InsertParams<'_, K>, + result: &mut InsertResult, + cache_actions: &mut Vec, + ) { + } + + /// Evict shallow device checkpoints beyond the per-path state cap on the + /// tail's root path; only the Mamba component caps its states. + fn evict_excess_path_states( + &self, + tree_core: &mut UnifiedTreeCore, + tail_node_id: NodeIdx_, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + ) { + } + + /// Redistribute component data between `new_parent` and `child` when a node is + /// split; `new_parent` is the newly created prefix node. + // Python reference — tree_component.py::TreeComponent.redistribute_on_node_split: + // @abstractmethod + // def redistribute_on_node_split( + // self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode + // ): + // """Redistribute component data between new_parent and child when a + // node is split. new_parent is the newly created prefix node. + // - Full: copies child's lock_ref to new_parent. + // - SWA: slices (or clones) the swa value for new_parent, copies + // lock_ref and component_uuid metadata, then syncs child's swa + // value with its (now-trimmed) full_value. + // - Mamba: sets new_parent's mamba value to None and lock_ref to 0 + // (mamba data stays on the original leaf, not on prefix nodes).""" + // ... + fn redistribute_on_node_split( + &self, + tree_core: &mut UnifiedTreeCore, + new_parent_id: NodeIdx_, + child_id: NodeIdx_, + ); + + /// Free this component's KV resources on a node being evicted; returns + /// (device_freed, host_freed) token counts. + // Python reference — tree_component.py::TreeComponent.evict_component: + // @abstractmethod + // def evict_component( + // self, + // node: UnifiedTreeNode, + // device_frees: dict[ComponentType, list[torch.Tensor]], + // host_frees: dict[ComponentType, list[torch.Tensor]], + // target: EvictLayer = EvictLayer.DEVICE, + // ) -> tuple[int, int]: + // """Free this component's KV resources on a node being evicted. + // + // *target* controls which layer(s) to evict: + // - DEVICE: free device memory and tombstone (value = None). + // Host data is untouched. + // - HOST: free host memory (host_value = None). + // Device data is untouched. + // - ALL: free both device and host memory. + // No tombstone — caller will delete the node. + // + // Returns (device_freed, host_freed) token counts.""" + // ... + fn evict_component( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + target: EvictLayer, + ) -> (usize, usize); + + /// Eviction priority on this node type; higher = evicted later, and evicting a + /// component cascade-evicts every component of equal or lower priority. + fn eviction_priority(&self, is_leaf: bool) -> i64 { + 0 + } + + /// Begin this component's device-eviction walk (build its cursor/heap). + fn evict_device_start(&self, tree_core: &mut UnifiedTreeCore, request_cnt: usize); + + /// Advance one eviction step and return a device leaf, if selected. + /// + /// Implementations must return after one allocator-relevant internal + /// mutation so the caller can drain pending frees before continuing. + fn evict_device_next_node( + &self, + tree_core: &mut UnifiedTreeCore, + tracker: &mut HashMap, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + ) -> Option; + + /// Clear this component's device-eviction walk state. + fn evict_device_end(&self, tree_core: &mut UnifiedTreeCore); + + /// Increment component lock refs, protecting nodes from eviction. + // Python reference — tree_component.py::TreeComponent.acquire_component_lock: + // @abstractmethod + // def acquire_component_lock( + // self, + // node: UnifiedTreeNode, + // result: IncLockRefResult, + // lock_host: bool = False, + // ) -> IncLockRefResult: + // """Increment component lock refs, protecting nodes from + // eviction. Updates evictable → protected size on first lock. + // - Full: path-lock — walks from node up to root, incrementing + // lock_ref on every ancestor. + // - SWA: path-lock — walks upward collecting swa values until the + // sliding window is filled; records a component_uuid at the + // boundary for release_component_lock to know where to stop. + // - Mamba: single-node lock — only increments lock_ref on the + // node itself (mamba state is per-leaf, not per-path). + // + // When ``lock_host`` is True, the lock applies to host-side state: + // - Full: single-node host lock. + // - SWA: host window-lock with a dedicated host UUID boundary. + // - Mamba: single-node host lock with host LRU detach.""" + // ... + fn acquire_component_lock( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + result: IncLockRefResult, + lock_host: bool, + ) -> IncLockRefResult; + + /// Decrement component lock refs, un-protecting nodes. + // Python reference — tree_component.py::TreeComponent.release_component_lock: + // @abstractmethod + // def release_component_lock( + // self, + // node: UnifiedTreeNode, + // params: Optional[DecLockRefParams], + // lock_host: bool = False, + // ) -> None: + // """Decrement component lock refs, un-protecting nodes. + // Updates protected → evictable size when lock_ref drops to 0. + // - Full: path-unlock — walks from node up to root, decrementing + // lock_ref on every ancestor. + // - SWA: path-unlock — walks upward, stopping at the node whose + // component_uuid matches the one recorded during acquire. + // - Mamba: single-node unlock — only decrements lock_ref on the + // node itself. + // + // When ``lock_host`` is True, the inverse host-side semantics apply.""" + // ... + fn release_component_lock( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + params: Option<&DecLockRefParams>, + lock_host: bool, + ); + + /// Early-release the SWA lock along [node, swa_uuid_for_lock] while leaving + /// the other components' locks intact; only the SWA component supports it. + fn release_window_lock( + &self, + _tree_core: &mut UnifiedTreeCore, + _node_id: NodeIdx_, + _swa_uuid_for_lock: Option, + _device_frees: &mut HashMap>, + _host_frees: &mut HashMap>, + ) { + unimplemented!("release_window_lock is SWA-only") + } + + /// Build transfer descriptors for this component in the given phase; None when + /// the component has nothing to transfer. + fn build_hicache_transfers( + &self, + tree_core: &UnifiedTreeCore, + node_id: NodeIdx_, + phase: CacheTransferPhase, + mamba_pool_idx: Option, + host_indices: Option, + token_ids: Option<&[i64]>, + prefetch_tokens: usize, + last_hash: Option<&str>, + ) -> Result>, TreeCoreRuntimeError> { + // Python reference — tree_component.py::TreeComponent.build_hicache_transfers: + // def build_hicache_transfers( + // self, + // node: UnifiedTreeNode, + // phase: CacheTransferPhase, + // *, + // mamba_pool_idx: Optional[torch.Tensor] = None, + // host_indices: Optional[torch.Tensor] = None, + // token_ids: Optional[Sequence[int]] = None, + // prefetch_tokens: int = 0, + // last_hash: Optional[str] = None, + // ) -> Optional[list[PoolTransfer]]: + // """Build transfer descriptors for this component in the given phase. + // Returns None if the component has nothing to transfer.""" + // return None + unimplemented!("TreeComponent.build_hicache_transfers") + } + + /// Post-transfer bookkeeping: store host indices, update LRU, etc. + fn commit_hicache_transfer( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + phase: CacheTransferPhase, + transfers: Vec, + cache_actions: &mut Vec, + insert_result: Option<&mut InsertResult>, + pool_storage_result: Option<&PoolTransferResult>, + ) { + // Python reference — tree_component.py::TreeComponent.commit_hicache_transfer: + // def commit_hicache_transfer( + // self, + // node: UnifiedTreeNode, + // phase: CacheTransferPhase, + // transfers: list[PoolTransfer] = (), + // *, + // cache_actions: list[CacheAction | ComponentAction], + // insert_result: Optional[InsertResult] = None, + // pool_storage_result: Optional[PoolTransferResult] = None, + // ) -> None: + // """Post-transfer bookkeeping: store host indices, update LRU, etc.""" + // pass + unimplemented!("TreeComponent.commit_hicache_transfer") + } + + /// Reclaim host values that coexist with device values before ordinary + /// host eviction. Called only under the write-back policy. + fn reclaim_coexisting_host_values( + &self, + _tree_core: &mut UnifiedTreeCore, + _num_tokens: usize, + _tracker: &mut HashMap, + _device_frees: &mut HashMap>, + _host_frees: &mut HashMap>, + ) { + } + + /// Evict from this component's host-side resources. + /// Called by HostPoolGroup when the host pool is full. + /// Default no-op for components without host storage. + fn drive_host_eviction( + &self, + _tree_core: &mut UnifiedTreeCore, + _num_tokens: usize, + _tracker: &mut HashMap, + _device_frees: &mut HashMap>, + _host_frees: &mut HashMap>, + ) { + } +} + +// Tree component types. + +/// The tree components; discriminants define the per-component array indexes. +#[derive(Copy, Clone, PartialEq, Eq, Debug, Hash)] +pub enum ComponentType { + Full = 0, + Swa = 1, + Mamba = 2, +} + +/// Short call-site aliases for the component types. +pub const FULL: ComponentType = ComponentType::Full; +pub const SWA: ComponentType = ComponentType::Swa; +pub const MAMBA: ComponentType = ComponentType::Mamba; + +/// The base component every tree runs; the others are auxiliary. +pub const BASE_COMPONENT_TYPE: ComponentType = ComponentType::Full; + +/// Slots per tier — the arrays are sized to this, not the enabled subset. +pub const NUM_COMPONENT_TYPES: usize = ComponentType::Mamba as usize + 1; + +impl ComponentType { + /// Index into a per-component array. + pub const fn idx(self) -> usize { + self as usize + } + + /// Whether the component stores one state slot per node (Mamba) instead of + /// one row per key atom. + pub fn single_value_per_node(self) -> bool { + matches!(self, ComponentType::Mamba) + } + + /// The component at a per-component array index; panics out of range. + pub fn from_idx(idx: usize) -> ComponentType { + match idx { + 0 => ComponentType::Full, + 1 => ComponentType::Swa, + 2 => ComponentType::Mamba, + _ => panic!("from_idx: {idx} is not a component index"), + } + } +} +#[cfg(test)] +#[path = "../tests/components/base.rs"] +mod tests; diff --git a/rust/mem-cache/src/components/swa.rs b/rust/mem-cache/src/components/swa.rs new file mode 100644 index 000000000..27906e8a3 --- /dev/null +++ b/rust/mem-cache/src/components/swa.rs @@ -0,0 +1,1197 @@ +//! SWA (sliding-window attention) component driver: overrides the methods SWA +//! customizes and inherits the rest from the `TreeComponent` defaults. +//! SWA values arrive pool-resolved; the full->SWA index translation happens at +//! the cache boundary. + +use std::collections::{HashMap, HashSet}; + +use tch::{Kind, Tensor}; + +use crate::components::TreeComponent; +use crate::components::{ComponentType, FULL, SWA}; +use crate::node::ChildKeyType; +use crate::node::Node; +use crate::node::{NodeId, NodeIdx_, TreeCoreRuntimeError, ValueSlotIdx}; +use crate::unified_tree_core::{ + CacheAction, CacheInitParams, CacheTransferPhase, DecLockRefParams, EvictLayer, + IncLockRefResult, InsertParams, InsertResult, LRURefreshPhase, MatchPrefixParams, MatchResult, + PoolHitPolicy, PoolName, PoolTransfer, PoolTransferResult, UnifiedTreeCore, +}; + +/// SWA component driver; owns the SWA device/host value slots. +pub struct SwaComponent { + /// Sliding window size in tokens. + sliding_window_size: usize, +} + +impl SwaComponent { + /// The component's device value slot. + pub const DEVICE: ValueSlotIdx = ValueSlotIdx::device(SWA); + /// The component's host value slot. + pub const HOST: ValueSlotIdx = ValueSlotIdx::host(SWA); +} + +impl SwaComponent { + /// Build the driver from the tree's init params. + pub fn new(params: &CacheInitParams) -> Self { + SwaComponent { + sliding_window_size: params + .swa_sliding_window_size + .expect("the SWA component requires swa_sliding_window_size"), + } + } + + /// Cap a fresh in-window SWA leaf at one page-aligned window so locking it pins + /// only one window of SWA pool, not the whole (long chunked-prefill) leaf; return + /// the split-off parent (older window) or None. The SWA value is stamped later, so + /// this runs on the tombstone leaf. + fn maybe_split_leaf_for_swa_lock_( + &self, + tree_core: &mut UnifiedTreeCore, + leaf_id: NodeIdx_, + ) -> Option { + let leaf = tree_core.arena.node(leaf_id); + if leaf.is_root() || leaf.device_lock_ref(SWA) > 0 { + return None; + } + + let page_size = tree_core.page_size; + // Smallest page-aligned size that still covers the sliding window. + let tail_size = self.sliding_window_size.div_ceil(page_size) * page_size; + let leaf_len = leaf.key.atom_len(); + if leaf_len <= tail_size { + return None; + } + let split_at = leaf_len - tail_size; + if page_size > 1 + && (!split_at.is_multiple_of(page_size) || !leaf_len.is_multiple_of(page_size)) + { + return None; + } + + let (new_parent, action) = tree_core.split_node_(leaf_id, split_at); + assert!( + action.is_none(), + "fresh SWA leaf cannot be write-through-pending" + ); + Some(new_parent) + } + + // Tier-selected SWA slot reads for the lock walks; `host` picks the host slot. + fn has_value(node: &Node, host: bool) -> bool { + if host { + node.has_host_value(SWA) + } else { + node.has_device_value(SWA) + } + } + + fn lock_ref(node: &Node, host: bool) -> u32 { + if host { + node.host_lock_ref(SWA) + } else { + node.device_lock_ref(SWA) + } + } + + fn inc_lock_ref(node: &mut Node, host: bool) { + if host { + node.inc_host_lock_ref(SWA); + } else { + node.inc_device_lock_ref(SWA); + } + } + + fn dec_lock_ref(node: &mut Node, host: bool) { + if host { + node.dec_host_lock_ref(SWA); + } else { + node.dec_device_lock_ref(SWA); + } + } + + fn value_len(node: &Node, host: bool) -> usize { + if host { + node.host_value_len(SWA) + } else { + node.device_value_len(SWA) + } + } + + fn swa_uuid(node: &Node, host: bool) -> Option { + if host { + node.swa_host_uuid + } else { + node.swa_uuid + } + } + + /// The node's SWA lock-window uuid for the tier, stamping a fresh one if absent. + fn ensure_swa_uuid( + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + host: bool, + ) -> i64 { + match Self::swa_uuid(tree_core.arena.node(node_id), host) { + Some(uuid) => uuid, + None => { + let minted = tree_core.next_swa_uuid_(); + let node = tree_core.arena.node_mut(node_id); + if host { + node.swa_host_uuid = Some(minted); + } else { + node.swa_uuid = Some(minted); + } + minted + } + } + } + + fn next_host_unlocked_device_lru_node( + tree_core: &UnifiedTreeCore, + from: Option, + ) -> Option { + let lru = tree_core.device_lru_list(SWA); + let unlocked = |id: NodeIdx_| tree_core.arena.node(id).host_lock_ref(SWA) == 0; + match from { + Some(node_id) => lru.get_prev_where(node_id, unlocked), + None => lru.get_lru_where(unlocked), + } + } +} + +impl SwaComponent { + /// Queue a free of the given SWA host slots; empty tensors are dropped. + fn release_swa_host_(&self, host_indices: Tensor, cache_actions: &mut Vec) { + if host_indices.numel() > 0 { + cache_actions.push(CacheAction::FreeComponentHostSlot { + component_type: SWA, + host_indices: vec![host_indices], + }); + } + } + + /// Write host_indices into node's SWA host_value and refresh tree state. + fn attach_swa_host_value_( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + host_indices: Tensor, + ) { + let node = tree_core.arena.node_mut(node_id); + let device_on = node.has_device_value(SWA); + node.set_host_value(SWA, host_indices.copy()); + let host_lru = tree_core.host_lru_list_mut(SWA); + if !device_on && !host_lru.in_list(Some(node_id)) { + host_lru.insert_mru(node_id); + } + tree_core.update_evictable_leaf_sets_(node_id); + if let Some(parent) = tree_core.arena.node(node_id).try_parent() { + tree_core.update_evictable_leaf_sets_(parent); + } + } + + /// Fill the prefetched SWA window onto the leaf→anchor path. + /// + /// All-or-nothing over one full window: `loaded_pages` is the cross-rank + /// MIN, so `loaded_pages < window_pages` drops the whole window (keeps the + /// tree identical across TP ranks). Otherwise map the buffer to token range + /// `[loaded_start, total_len)` and walk leaf→anchor, filling SWA + /// tombstones and releasing slices that already have host_value. + fn commit_prefetch_( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + transfers: Vec, + cache_actions: &mut Vec, + insert_result: Option<&InsertResult>, + pool_storage_result: Option<&PoolTransferResult>, + ) { + if transfers.is_empty() { + return; + } + let page_size = tree_core.page_size; + let transfer = &transfers[0]; + let window_require_pages = transfer + .host_indices + .as_ref() + .map_or(0, |host| host.numel() / page_size); + let loaded_pages = pool_storage_result.map_or(0, |result| { + result + .extra_pool_hit_pages + .get(&PoolName::Swa) + .copied() + .unwrap_or(0) + }); + let target = insert_result + .and_then(|result| result.inserted_host_node) + .map(|id| tree_core.arena.resolve(id)); + + let (Some(target), Some(host_indices)) = (target, transfer.host_indices.as_ref()) else { + if let Some(host_indices) = &transfer.host_indices { + self.release_swa_host_(host_indices.shallow_clone(), cache_actions); + } + return; + }; + if window_require_pages == 0 || loaded_pages < window_require_pages { + self.release_swa_host_(host_indices.shallow_clone(), cache_actions); + return; + } + let insert_result = insert_result.expect("target implies an insert result"); + + // The buffer covers token range [loaded_start, total_len). + let loaded_start = insert_result.total_len - window_require_pages * page_size; + + // Walk leaf -> anchor; pos is the right edge of cur in tokens. + let mut pos = insert_result.total_len; + let mut cur = target; + while cur != node_id && pos > loaded_start { + let cur_node = tree_core.arena.node(cur); + let node_start = pos - cur_node.key.atom_len(); + // Intersection of cur's range and the buffer. + let fill_start = node_start.max(loaded_start); + let fill_len = pos - fill_start; + let buf_off = fill_start - loaded_start; + let slice = host_indices.narrow(0, buf_off as i64, fill_len as i64); + let parent = cur_node + .try_parent() + .expect("prefetch walk reached a root before the anchor"); + + if !cur_node.has_host_value(SWA) && fill_len > 0 { + // Tombstone: split off the in-buffer tail if needed, then fill. + if fill_start > node_start { + let (_, action) = tree_core.split_node_(cur, fill_start - node_start); + if let Some(action) = action { + cache_actions.push(action); + } + } + self.attach_swa_host_value_(tree_core, cur, slice); + } else { + // Already has SWA (or empty overlap): drop this slice. + self.release_swa_host_(slice, cache_actions); + } + + pos = node_start; + cur = parent; + } + + // Buffer prefix that fell outside the anchor->leaf path. + if pos > loaded_start { + self.release_swa_host_( + host_indices.narrow(0, 0, (pos - loaded_start) as i64), + cache_actions, + ); + } + } +} + +impl TreeComponent for SwaComponent { + fn component_type(&self) -> ComponentType { + SWA + } + + fn refresh_lru( + &self, + tree_core: &mut UnifiedTreeCore, + phase: LRURefreshPhase, + node_id: NodeIdx_, + ) { + match phase { + // Walk-down would refresh every visited ancestor to MRU, but most + // are outside the active sliding window and must stay evictable. + // Window-bounded refresh runs at MATCH_END / INSERT_END instead. + LRURefreshPhase::Walkdown => {} + LRURefreshPhase::MatchEnd | LRURefreshPhase::InsertEnd => { + let window = self.sliding_window_size + tree_core.page_size; + let (lru, arena) = tree_core.device_lru_list_mut_with_arena(SWA); + lru.reset_node_and_window_ancestors_mru(node_id, window, arena, |node| { + node.has_device_value(SWA) + }); + } + } + } + + fn create_match_validator( + &self, + tree_core: &UnifiedTreeCore, + match_device_only: bool, + ) -> Box, NodeIdx_) -> bool> { + let sliding_window_size = self.sliding_window_size; + // unified_kv never caches the SWA ring (per-request, not + // content-stable), so SWA bookkeeping must not gate the match here. + let swa_device_only_hicache = !tree_core.has_swa_host_pool && tree_core.enable_hicache; + let mut contiguous_len = usize::MAX; + Box::new(move |tree_core: &UnifiedTreeCore, node_id: NodeIdx_| { + let node = tree_core.arena.node(node_id); + // HiCache: a host-only tombstone is a valid match boundary too + // — load_back will restore SWA from host before use. + if !node.has_device_value(SWA) && (match_device_only || !node.has_host_value(SWA)) { + contiguous_len = 0; + return swa_device_only_hicache && (node.backuped() || !node.evicted()); + } + contiguous_len = contiguous_len.saturating_add(node.key.atom_len()); + contiguous_len >= sliding_window_size + }) + } + + fn finalize_match_result_in_tree_core( + &self, + tree_core: &UnifiedTreeCore, + mut result: MatchResult, + params: &MatchPrefixParams<'_, K>, + value_chunks: &[Tensor], + best_value_len: usize, + ) -> MatchResult { + // Sum the SWA tokens backing the match, walking up from the best match + // until one sliding window is covered; host-resident chunks count + // toward the SWA host hit. + let mut n_swa = 0; + let mut swa_host_hit = 0; + let mut node = tree_core + .arena + .node(tree_core.arena.resolve(result.best_match_node_id)); + while !node.is_root() && n_swa < self.sliding_window_size { + if node.has_device_value(SWA) { + n_swa += node.device_value_len(SWA); + } else if node.has_host_value(SWA) { + // TODO(hzh): once load_back is constrained to fetch only one + // sliding window worth of pages, cap swa_host_hit at + // sliding_window_size so the scheduler budget matches the + // actual device-pool consumption. + let host_len = node.host_value_len(SWA); + swa_host_hit += host_len; + n_swa += host_len; + } else { + break; + } + node = tree_core.arena.node(node.parent()); + } + if swa_host_hit > 0 { + result.swa_host_hit_length = result.swa_host_hit_length.max(swa_host_hit); + } + result + } + + fn update_component_on_insert_overlap( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + prefix_len: usize, + total_prefix_len: usize, + value_slice: Tensor, + params: &InsertParams<'_, K>, + result: &mut InsertResult, + cache_actions: &mut Vec, + ) -> usize { + if params.prev_prefix_len >= total_prefix_len + prefix_len { + return prefix_len; + } + + let node = tree_core.arena.node_mut(node_id); + let is_tombstone = !node.has_device_value(SWA); + if !is_tombstone { + return prefix_len; + } + + let swa_evicted_seqlen = params.swa_evicted_seqlen; + assert_eq!( + node.device_lock_ref(SWA), + 0, + "tombstone Swa lock_ref should be 0, node {node_id}" + ); + assert_eq!( + swa_evicted_seqlen % tree_core.page_size, + 0, + "Swa: swa_evicted_seqlen must be page-aligned, swa_evicted_seqlen={swa_evicted_seqlen}" + ); + + if swa_evicted_seqlen <= total_prefix_len { + // Branch 1: entire value_slice is within SWA window — recover + result.record_adopted_range(SWA, total_prefix_len, total_prefix_len + prefix_len); + if node.device_lock_ref(FULL) > 0 { + cache_actions.push(CacheAction::RecoverSwaWithLockedFull { + node_id: node.id, + kept_full: node.device_value(FULL).shallow_clone(), + incoming_full: value_slice, + }); + return 0; + } + result.record_adopted_range(FULL, total_prefix_len, total_prefix_len + prefix_len); + let old_full = node.take_device_value(FULL); + node.set_device_value(FULL, value_slice.copy()); + cache_actions.push(CacheAction::FreeDeviceKVFullOnly(vec![old_full])); + cache_actions.push(CacheAction::SwaRebuild { + node_id: node.id, + source_value: value_slice, + }); + 0 + } else if swa_evicted_seqlen < total_prefix_len + prefix_len { + // Branch 2: value_slice[start_idx:] is within SWA window — partial recover + let start_idx = swa_evicted_seqlen - total_prefix_len; + result.record_adopted_range(SWA, swa_evicted_seqlen, total_prefix_len + prefix_len); + let node_ext_id = node.id; + let is_locked = node.device_lock_ref(FULL) > 0; + let full_len = node.device_value_len(FULL); + let old_full = + node.device_value(FULL) + .narrow(0, start_idx as i64, (full_len - start_idx) as i64); + let (_, action) = tree_core.split_node_(node_id, start_idx); + if let Some(action) = action { + cache_actions.push(action); + } + let new_full = value_slice.narrow(0, start_idx as i64, (prefix_len - start_idx) as i64); + if is_locked { + cache_actions.push(CacheAction::RecoverSwaWithLockedFull { + node_id: node_ext_id, + kept_full: old_full, + incoming_full: new_full, + }); + return start_idx; + } + result.record_adopted_range(FULL, swa_evicted_seqlen, total_prefix_len + prefix_len); + let node = tree_core.arena.node_mut(node_id); + let _ = node.take_device_value(FULL); + node.set_device_value(FULL, new_full.copy()); + cache_actions.push(CacheAction::FreeDeviceKVFullOnly(vec![old_full])); + cache_actions.push(CacheAction::SwaRebuild { + node_id: node_ext_id, + source_value: new_full, + }); + start_idx + } else { + // Branch 3: entire value_slice is outside SWA window — not consumed + prefix_len + } + } + + fn recover_after_unevict( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + prefix_len: usize, + total_prefix_len: usize, + params: &InsertParams<'_, K>, + result: &mut InsertResult, + cache_actions: &mut Vec, + ) { + // _unevict_node_on_insert already wrote the request's fresh KV slice + // into the base value. We just need to rebuild SWA from that slice for + // the in-window portion. There is no old SWA slot to free here. + let node = tree_core.arena.node(node_id); + if node.has_device_value(SWA) { + return; + } + assert_eq!( + node.device_lock_ref(SWA), + 0, + "tombstone Swa lock_ref should be 0 on unevict, node {node_id}" + ); + let swa_evicted_seqlen = params.swa_evicted_seqlen; + assert_eq!( + swa_evicted_seqlen % tree_core.page_size, + 0, + "Swa: swa_evicted_seqlen must be page-aligned, swa_evicted_seqlen={swa_evicted_seqlen}" + ); + + if swa_evicted_seqlen <= total_prefix_len { + // entire node is within the SWA window + } else if swa_evicted_seqlen < total_prefix_len + prefix_len { + let start_idx = swa_evicted_seqlen - total_prefix_len; + let (_, action) = tree_core.split_node_(node_id, start_idx); + if let Some(action) = action { + cache_actions.push(action); + } + } else { + return; + } + result.record_adopted_range( + SWA, + total_prefix_len.max(swa_evicted_seqlen), + total_prefix_len + prefix_len, + ); + cache_actions.push(CacheAction::SwaRebuild { + node_id: tree_core.arena.node(node_id).id, + source_value: tree_core.arena.device_value(node_id, FULL).shallow_clone(), + }); + } + + fn commit_insert_component_data( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + is_new_leaf: bool, + params: &InsertParams<'_, K>, + result: &mut InsertResult, + cache_actions: &mut Vec, + ) { + if !is_new_leaf { + return; + } + + let node_start = result.prefix_len; + let node_end = node_start + tree_core.arena.node(node_id).key.atom_len(); + // A boundary above the leaf skips the split (Python's negative split_pos). + if params.swa_evicted_seqlen >= node_start { + let split_pos = params.swa_evicted_seqlen - node_start; + if split_pos >= tree_core.arena.node(node_id).key.atom_len() { + // Entire leaf is outside the SWA window — left as a tombstone. + return; + } + if split_pos > 0 { + // Node straddles the boundary: split into an out-of-window parent + // (tombstone) and an in-window child; `node` becomes the child. + let (_, action) = tree_core.split_node_(node_id, split_pos); + assert!(action.is_none(), "new leaf cannot be write-through-pending"); + } + } + result.record_adopted_range(SWA, node_start.max(params.swa_evicted_seqlen), node_end); + // Cap the in-window leaf at one window for lock granularity, then rebuild SWA + // onto the in-window node(s) at apply time; rebuild the older prefix first so + // the in-window tail lands more-MRU. + let capped_parent = self.maybe_split_leaf_for_swa_lock_(tree_core, node_id); + if let Some(capped_parent) = capped_parent { + cache_actions.push(CacheAction::SwaRebuild { + node_id: tree_core.arena.node(capped_parent).id, + source_value: tree_core + .arena + .device_value(capped_parent, FULL) + .shallow_clone(), + }); + } + cache_actions.push(CacheAction::SwaRebuild { + node_id: tree_core.arena.node(node_id).id, + source_value: tree_core.arena.device_value(node_id, FULL).shallow_clone(), + }); + } + + fn redistribute_on_node_split( + &self, + tree_core: &mut UnifiedTreeCore, + new_parent_id: NodeIdx_, + child_id: NodeIdx_, + ) { + let (new_parent, child) = tree_core.arena.node_pair_mut(new_parent_id, child_id); + let split_len = new_parent.key.atom_len() as i64; + new_parent.copy_device_lock_ref(SWA, child); + if child.has_device_value(SWA) { + Node::redistribute_child_device_value(new_parent, child, SWA, split_len); + } + if child.has_host_value(SWA) { + Node::redistribute_child_host_value(new_parent, child, SWA, split_len); + // Device-tombstoned sides park in the host LRU. + let parent_is_tombstone = !new_parent.has_device_value(SWA); + let child_is_tombstone = !child.has_device_value(SWA); + let host_lru = tree_core.host_lru_list_mut(SWA); + if parent_is_tombstone { + host_lru.insert_mru(new_parent_id); + } + if child_is_tombstone && !host_lru.in_list(Some(child_id)) { + host_lru.insert_mru(child_id); + } + } + + // parent inherits the swa_uuid from child for swa lock ref + let swa_uuid = tree_core.arena.node_mut(child_id).swa_uuid.take(); + tree_core.arena.node_mut(new_parent_id).swa_uuid = swa_uuid; + } + + fn evict_component( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + target: EvictLayer, + ) -> (usize, usize) { + let ct = SWA; + let node = tree_core.arena.node_mut(node_id); + let mut freed = 0; + let mut host_freed = 0; + + // Device layer + if target.contains(EvictLayer::Device) && node.has_device_value(SWA) { + // Pass full indices to free_swa so slots with no SWA pair are + // skipped. Freeing swa_value directly would double free those + // entries since they all map to the same sentinel slot. + device_frees + .entry(ct) + .or_default() + .push(node.device_value(FULL).shallow_clone()); + freed = node.device_value_len(SWA); + let _ = node.take_device_value(SWA); + tree_core.dec_evictable_size(SWA, freed); + } + + // Host layer + let node = tree_core.arena.node_mut(node_id); + if target.contains(EvictLayer::Host) && node.has_host_value(SWA) { + host_freed = node.host_value_len(SWA); + host_frees + .entry(ct) + .or_default() + .push(node.take_host_value(SWA)); + let host_lru = tree_core.host_lru_list_mut(SWA); + if host_lru.in_list(Some(node_id)) { + host_lru.remove_node(node_id); + } + } + + // After device tombstone: if host_value remains, move into host LRU + let node = tree_core.arena.node(node_id); + if target == EvictLayer::Device && !node.has_device_value(SWA) && node.has_host_value(SWA) { + let host_lru = tree_core.host_lru_list_mut(SWA); + if !host_lru.in_list(Some(node_id)) { + host_lru.insert_mru(node_id); + } + } + + (freed, host_freed) + } + + fn eviction_priority(&self, is_leaf: bool) -> i64 { + if is_leaf { 0 } else { 1 } + } + + fn evict_device_start(&self, tree_core: &mut UnifiedTreeCore, request_cnt: usize) { + tree_core.set_evict_device_start(SWA, request_cnt); + let cursor = tree_core + .device_lru_list(SWA) + .get_lru_no_lock(&tree_core.arena); + tree_core.component_state_mut(SWA).evict_device_cursor = cursor; + } + + /// Advance one device-eviction step and return a leaf, if selected. + /// + /// An internal tombstone is one complete step so the caller can apply its + /// pending frees and recheck allocator capacity before the next mutation. + fn evict_device_next_node( + &self, + tree_core: &mut UnifiedTreeCore, + tracker: &mut HashMap, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + ) -> Option { + let ct = SWA; + assert!( + tree_core.component_state(SWA).is_evict_device_ongoing, + "Swa device eviction not started" + ); + let mut cursor = tree_core.component_state(SWA).evict_device_cursor; + // The cursor is re-validated (reset to LRU head) if the previous + // node's eviction removed it. + if cursor.is_some_and(|c| !tree_core.device_lru_list(SWA).in_list(Some(c))) { + cursor = tree_core + .device_lru_list(SWA) + .get_lru_no_lock(&tree_core.arena); + } + let next = loop { + if tracker[&ct] >= tree_core.component_state(SWA).evict_device_request_cnt { + break None; + } + let Some(x) = cursor else { + break None; + }; + if !tree_core.device_lru_list(SWA).in_list(Some(x)) { + break None; + } + assert!( + tree_core.arena.has_device_value(x, SWA), + "Swa eviction cursor on a valueless node {x}" + ); + cursor = tree_core + .device_lru_list(SWA) + .get_prev_no_lock(x, &tree_core.arena); + // A load-back pin means an in-flight DMA targets this node's slices. + if tree_core.arena.node(x).is_load_back_pending() { + continue; + } + if tree_core.evictable_device_leaves.contains(x) { + break Some(x); + } + // Internal nodes are tombstoned inline (no IO). + tree_core.evict_component_and_detach_lru_( + x, + ct, + device_frees, + host_frees, + EvictLayer::Device, + Some(tracker), + ); + tree_core.cascade_evict_(x, ct, tracker, device_frees, host_frees, EvictLayer::Device); + break None; + }; + tree_core.component_state_mut(SWA).evict_device_cursor = cursor; + next + } + + fn evict_device_end(&self, tree_core: &mut UnifiedTreeCore) { + tree_core.set_evict_device_end(SWA); + } + + fn reclaim_coexisting_host_values( + &self, + tree_core: &mut UnifiedTreeCore, + num_tokens: usize, + tracker: &mut HashMap, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + ) { + for spare_imminent_demotes in [true, false] { + if tracker[&SWA] >= num_tokens { + break; + } + let mut next = Self::next_host_unlocked_device_lru_node(tree_core, None); + while let Some(node_id) = next { + if tracker[&SWA] >= num_tokens { + break; + } + next = Self::next_host_unlocked_device_lru_node(tree_core, Some(node_id)); + if spare_imminent_demotes && tree_core.evictable_device_leaves.contains(node_id) { + continue; + } + if !tree_core.can_reclaim_coexisting_host_value_(node_id, SWA) { + continue; + } + tree_core.release_coexisting_host_value_( + node_id, + SWA, + tracker, + device_frees, + host_frees, + ); + } + } + } + + /// Evict SWA host resources. + /// Internal nodes: private tombstone (free SWA host only). + /// Host leaves: atomic eviction via _evict_host_leaf. + fn drive_host_eviction( + &self, + tree_core: &mut UnifiedTreeCore, + num_tokens: usize, + tracker: &mut HashMap, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + ) { + let ct = SWA; + let mut x = tree_core + .host_lru_list(SWA) + .get_lru_no_lock(&tree_core.arena); + loop { + if tracker[&ct] >= num_tokens { + break; + } + let Some(cur) = x else { + break; + }; + if !tree_core.host_lru_list(SWA).in_list(Some(cur)) { + break; + } + let x_next = tree_core + .host_lru_list(SWA) + .get_prev_no_lock(cur, &tree_core.arena); + // A load-back pin means an in-flight DMA reads this node's host slices. + if tree_core.arena.node(cur).is_load_back_pending() { + x = x_next; + continue; + } + if tree_core.evictable_host_leaves.contains(cur) { + tree_core.evict_host_leaf_(cur, tracker, device_frees, host_frees); + } else { + assert!( + tree_core.arena.has_host_value(cur, SWA), + "SWA host LRU member {cur} has no host value" + ); + tree_core.evict_component_and_detach_lru_( + cur, + ct, + device_frees, + host_frees, + EvictLayer::Host, + Some(tracker), + ); + tree_core.cascade_evict_( + cur, + ct, + tracker, + device_frees, + host_frees, + EvictLayer::Host, + ); + } + x = x_next; + } + } + + fn build_hicache_transfers( + &self, + tree_core: &UnifiedTreeCore, + node_id: NodeIdx_, + phase: CacheTransferPhase, + _mamba_pool_idx: Option, + host_indices: Option, + _token_ids: Option<&[i64]>, + _prefetch_tokens: usize, + _last_hash: Option<&str>, + ) -> Result>, TreeCoreRuntimeError> { + // unified_kv keeps SWA as a device-only ring. + if !tree_core.has_swa_host_pool && tree_core.enable_hicache { + return Ok(None); + } + Ok(match phase { + CacheTransferPhase::BackupHost => { + let node = tree_core.arena.node(node_id); + if node.has_host_value(SWA) { + return Ok(None); + } + // cd.value already holds SWA-pool indices (translated at insert time). + // Host pool indexing wants int64. + node.try_device_value(SWA).map(|value| { + vec![PoolTransfer { + name: PoolName::Swa, + device_indices: Some(value.to_kind(Kind::Int64)), + ..Default::default() + }] + }) + } + CacheTransferPhase::LoadBack => { + // `node` is best_match_node; the SWA validator guarantees every + // ancestor within `sliding_window_size` has value or host_value. + let mut n_swa = 0; + let mut backed_up: Vec = Vec::new(); + let mut nodes_to_load: Vec = Vec::new(); + let mut cur = tree_core.arena.node(node_id); + while !cur.is_root() && n_swa < self.sliding_window_size { + if let Some(value) = cur.try_device_value(SWA) { + // Device exists, skip it. + n_swa += value.size()[0] as usize; + } else if let Some(host_value) = cur.try_host_value(SWA) { + // Host only, collect it. + backed_up.push(host_value.shallow_clone()); + nodes_to_load.push(cur.id); + n_swa += host_value.size()[0] as usize; + } else { + return Err(TreeCoreRuntimeError::SwaLoadBackMissingValue { + node_id: cur.id, + }); + } + cur = tree_core.arena.node(cur.parent()); + } + if backed_up.is_empty() { + return Ok(None); + } + backed_up.reverse(); + nodes_to_load.reverse(); + Some(vec![PoolTransfer { + name: PoolName::Swa, + host_indices: Some(Tensor::cat(&backed_up, 0)), + nodes_to_load: Some(nodes_to_load), + ..Default::default() + }]) + } + CacheTransferPhase::BackupStorage => { + let node = tree_core.arena.node(node_id); + let Some(host_value) = node.try_host_value(SWA) else { + return Ok(None); + }; + let Some(hash_value) = node.hash_value.as_ref().filter(|h| !h.is_empty()) else { + return Ok(None); + }; + let page_size = tree_core.page_size as i64; + let num_pages = host_value.size()[0] / page_size; + if num_pages == 0 { + return Ok(None); + } + let host_len = host_value.size()[0]; + Some(vec![PoolTransfer { + name: PoolName::Swa, + host_indices: Some(host_value.narrow( + 0, + host_len - num_pages * page_size, + num_pages * page_size, + )), + keys: Some( + hash_value[hash_value.len().saturating_sub(num_pages as usize)..].to_vec(), + ), + hit_policy: PoolHitPolicy::TrailingPages, + ..Default::default() + }]) + } + CacheTransferPhase::Prefetch => { + let host_indices = host_indices.expect("SWA PREFETCH build requires host indices"); + let sw_pages = host_indices.numel() / tree_core.page_size; + Some(vec![PoolTransfer { + name: PoolName::Swa, + host_indices: Some(host_indices), + keys: Some(vec!["__placeholder__".to_string(); sw_pages]), + hit_policy: PoolHitPolicy::TrailingPages, + ..Default::default() + }]) + } + }) + } + + fn commit_hicache_transfer( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + phase: CacheTransferPhase, + transfers: Vec, + cache_actions: &mut Vec, + insert_result: Option<&mut InsertResult>, + pool_storage_result: Option<&PoolTransferResult>, + ) { + match phase { + CacheTransferPhase::BackupHost => { + if let Some(transfer) = transfers.first() + && let Some(host_indices) = &transfer.host_indices + { + let node = tree_core.arena.node_mut(node_id); + if !node.has_host_value(SWA) { + node.set_host_value(SWA, host_indices.copy()); + } + } + } + CacheTransferPhase::LoadBack => { + let transfer = transfers + .first() + .expect("SWA LOAD_BACK commit requires a transfer"); + let device_indices = transfer + .device_indices + .as_ref() + .expect("SWA LOAD_BACK commit requires device indices"); + let mut full_chunks: Vec = Vec::new(); + let mut swa_chunks: Vec = Vec::new(); + let mut offset = 0i64; + for &loaded_id in transfer.nodes_to_load.iter().flatten() { + let loaded_idx = tree_core.arena.resolve(loaded_id); + let n_tokens = tree_core.arena.host_value_len(loaded_idx, SWA) as i64; + let swa_chunk = device_indices.narrow(0, offset, n_tokens).copy(); + tree_core.set_component_device_value_( + loaded_idx, + SWA, + swa_chunk.shallow_clone(), + ); + let full_value = tree_core + .arena + .device_value(loaded_idx, FULL) + .shallow_clone(); + assert_eq!(full_value.size()[0], n_tokens); + full_chunks.push(full_value); + swa_chunks.push(swa_chunk); + offset += n_tokens; + } + let host_indices = transfer + .host_indices + .as_ref() + .expect("SWA LOAD_BACK commit requires host indices"); + assert_eq!(offset, host_indices.size()[0]); + // Rebuild the full->swa mapping for the loaded chunks at the orchestration layer. + if !full_chunks.is_empty() { + cache_actions.push(CacheAction::RebuildFullToSwaMapping { + full_indices: full_chunks, + swa_indices: swa_chunks, + }); + } + } + CacheTransferPhase::Prefetch => { + self.commit_prefetch_( + tree_core, + node_id, + transfers, + cache_actions, + insert_result.as_deref(), + pool_storage_result, + ); + } + CacheTransferPhase::BackupStorage => {} + } + } + + fn acquire_component_lock( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + mut result: IncLockRefResult, + lock_host: bool, + ) -> IncLockRefResult { + let ct = SWA; + let sliding_window_size = self.sliding_window_size; + let mut swa_lock_size = 0; + let mut swa_uuid = None; + + // Tombstoned nodes (cd.value is None) have no SWA chunk to protect + // skip them and keep walking up. This path is hit when HiCache + // backs up a FULL present internal node whose SWA was already evicted. + let mut cur = node_id; + loop { + let node = tree_core.arena.node_mut(cur); + if node.is_root() || swa_lock_size >= sliding_window_size { + break; + } + let parent = node.parent(); + if !Self::has_value(node, lock_host) { + result + .skip_lock_node_ids + .entry(ct) + .or_default() + .insert(node.id); + cur = parent; + continue; + } + let key_len = node.key.atom_len(); + let newly_locked = Self::lock_ref(node, lock_host) == 0; + Self::inc_lock_ref(node, lock_host); + swa_lock_size += Self::value_len(node, lock_host); + if newly_locked { + if lock_host { + let host_lru = tree_core.host_lru_list_mut(SWA); + if host_lru.in_list(Some(cur)) { + host_lru.remove_node(cur); + } + } else { + tree_core.dec_evictable_size(SWA, key_len); + tree_core.inc_protected_size(SWA, key_len); + } + } + if swa_lock_size >= sliding_window_size { + swa_uuid = Some(Self::ensure_swa_uuid(tree_core, cur, lock_host)); + } + cur = parent; + } + + if lock_host { + result.swa_uuid_for_host_lock = swa_uuid; + } else { + result.swa_uuid_for_lock = swa_uuid; + } + result + } + + fn release_component_lock( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + params: Option<&DecLockRefParams>, + lock_host: bool, + ) { + let ct = SWA; + let swa_uuid_for_lock = params.and_then(|p| { + if lock_host { + p.swa_uuid_for_host_lock + } else { + p.swa_uuid_for_lock + } + }); + let empty = HashSet::new(); + let skip_lock_node_ids = params + .and_then(|p| p.skip_lock_node_ids.get(&ct)) + .unwrap_or(&empty); + + // A node in skip_lock_node_ids was a tombstone when this lock was acquired. + let mut cur = node_id; + loop { + let node = tree_core.arena.node_mut(cur); + if node.is_root() { + break; + } + let parent = node.parent(); + if skip_lock_node_ids.contains(&node.id) { + cur = parent; + continue; + } + let lock_ref = Self::lock_ref(node, lock_host); + if lock_ref == 0 { + cur = parent; + continue; + } + if lock_ref == 1 { + if lock_host { + if !node.has_device_value(SWA) && node.has_host_value(SWA) { + let host_lru = tree_core.host_lru_list_mut(SWA); + if !host_lru.in_list(Some(cur)) { + host_lru.insert_mru(cur); + } + } + } else { + let key_len = node.device_value_len(SWA); + tree_core.inc_evictable_size(SWA, key_len); + tree_core.dec_protected_size(SWA, key_len); + } + } + Self::dec_lock_ref(tree_core.arena.node_mut(cur), lock_host); + if swa_uuid_for_lock.is_some() + && Self::swa_uuid(tree_core.arena.node(cur), lock_host) == swa_uuid_for_lock + { + break; + } + cur = parent; + } + } + + /// Early-release the SWA lock along [node, swa_uuid_for_lock] while + /// leaving Full and Mamba locks intact. + /// + /// Called when a request's decode position has advanced past the sliding + /// window — the SWA portion of the tree lock is no longer needed but the + /// Full lock must stay so the request's prefix is protected. + /// + /// Caller (UnifiedRadixCache.dec_swa_lock_only) must ensure this is + /// invoked at most once per (node, swa_uuid_for_lock) pair. + fn release_window_lock( + &self, + tree_core: &mut UnifiedTreeCore, + node_id: NodeIdx_, + swa_uuid_for_lock: Option, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + ) { + let ct = SWA; + let mut cur = node_id; + loop { + let node = tree_core.arena.node_mut(cur); + if node.is_root() { + break; + } + let parent = node.parent(); + // Acquire skips tombstoned nodes; release must skip them too. Same + // for nodes with lock_ref == 0 — acquire never credited them. + if !node.has_device_value(SWA) || node.device_lock_ref(SWA) == 0 { + if swa_uuid_for_lock.is_some() && node.swa_uuid == swa_uuid_for_lock { + break; + } + cur = parent; + continue; + } + + node.dec_device_lock_ref(SWA); + if node.device_lock_ref(SWA) == 0 { + let key_len = node.key.atom_len(); + tree_core.dec_protected_size(SWA, key_len); + tree_core.inc_evictable_size(SWA, key_len); + if tree_core.is_evictable_device_leaf_(tree_core.arena.node(cur)) { + tree_core.evict_component_and_detach_lru_( + cur, + ct, + device_frees, + host_frees, + EvictLayer::Device, + /* tracker = */ None, + ); + } + } + + if swa_uuid_for_lock.is_some() + && tree_core.arena.node(cur).swa_uuid == swa_uuid_for_lock + { + break; + } + cur = parent; + } + } +} + +#[cfg(test)] +#[path = "../tests/components/swa.rs"] +mod tests; diff --git a/rust/mem-cache/src/lib.rs b/rust/mem-cache/src/lib.rs new file mode 100644 index 000000000..3f67a7d5e --- /dev/null +++ b/rust/mem-cache/src/lib.rs @@ -0,0 +1,27 @@ +//! Rust radix tree core for SGLang's KV cache, exposed to Python as `mem_cache`. +// TODO(Jialin): Replace recoverable panics with explicit Rust errors and map +// them to PyErr at the Python boundary. +#![allow( + dead_code, + unsafe_op_in_unsafe_fn, + clippy::unwrap_used, + clippy::expect_used, + clippy::owned_cow, + clippy::panic, + clippy::print_stdout, + clippy::too_many_arguments, + clippy::type_complexity, + clippy::unimplemented, + clippy::unreachable, + clippy::useless_conversion +)] + +mod components; +mod node; +#[cfg(feature = "python-extension")] +mod python_bindings; +#[cfg(test)] +#[path = "tests/test_utils.rs"] +pub(crate) mod test_utils; +mod unified_lru_list; +mod unified_tree_core; diff --git a/rust/mem-cache/src/node.rs b/rust/mem-cache/src/node.rs new file mode 100644 index 000000000..1e3758edb --- /dev/null +++ b/rust/mem-cache/src/node.rs @@ -0,0 +1,1571 @@ +//! A radix-tree node, owned by the `NodeArena` and referenced by `NodeIdx_`. + +use std::borrow::{Borrow, Cow}; +use std::collections::HashMap; +use std::collections::hash_map::{DefaultHasher, RandomState}; +use std::fmt::Debug; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use hashbrown::hash_map::Entry; +use hashbrown::{Equivalent, HashMap as HashBrownMap}; +use sha2::{Digest, Sha256}; +use tch::Tensor; + +use crate::components::{ComponentType, FULL, NUM_COMPONENT_TYPES}; + +/// The two independent dimensions that partition a radix tree. +#[derive(Debug)] +struct KeyNamespaceData { + extra_key: Option>, + cache_salt: Option>, + hash: u64, +} + +/// Compact, shared namespace stored on nodes and child edges. +/// +/// The default namespace is represented without an allocation. Non-default +/// namespaces share one immutable allocation down a radix path. +#[derive(Clone, Debug, Default)] +pub struct KeyNamespace(Option>); + +/// Borrowed namespace used by match/insert lookups without allocating. +#[derive(Clone, Copy, Debug, Default)] +pub struct KeyNamespaceRef<'a> { + pub extra_key: Option<&'a str>, + pub cache_salt: Option<&'a str>, + hash: u64, +} + +fn key_namespace_hash(extra_key: Option<&str>, cache_salt: Option<&str>) -> u64 { + if extra_key.is_none() && cache_salt.is_none() { + return 0; + } + let mut hasher = DefaultHasher::new(); + extra_key.hash(&mut hasher); + cache_salt.hash(&mut hasher); + hasher.finish() +} + +impl<'a> KeyNamespaceRef<'a> { + pub fn new(extra_key: Option<&'a str>, cache_salt: Option<&'a str>) -> Self { + let cache_salt = cache_salt.filter(|salt| !salt.is_empty()); + Self { + extra_key, + cache_salt, + hash: key_namespace_hash(extra_key, cache_salt), + } + } + + pub fn to_owned(self) -> KeyNamespace { + if self.extra_key.is_none() && self.cache_salt.is_none() { + return KeyNamespace::default(); + } + KeyNamespace(Some(Arc::new(KeyNamespaceData { + extra_key: self.extra_key.map(Into::into), + cache_salt: self.cache_salt.map(Into::into), + hash: self.hash, + }))) + } +} + +impl PartialEq for KeyNamespaceRef<'_> { + fn eq(&self, other: &Self) -> bool { + self.extra_key == other.extra_key && self.cache_salt == other.cache_salt + } +} + +impl Eq for KeyNamespaceRef<'_> {} + +impl Hash for KeyNamespaceRef<'_> { + fn hash(&self, state: &mut H) { + state.write_u64(self.hash); + } +} + +impl KeyNamespace { + pub fn new(extra_key: Option<&str>, cache_salt: Option<&str>) -> Self { + KeyNamespaceRef::new(extra_key, cache_salt).to_owned() + } + + pub fn as_ref(&self) -> KeyNamespaceRef<'_> { + match self.0.as_deref() { + Some(namespace) => KeyNamespaceRef { + extra_key: namespace.extra_key.as_deref(), + cache_salt: namespace.cache_salt.as_deref(), + hash: namespace.hash, + }, + None => KeyNamespaceRef::default(), + } + } + + pub fn extra_key(&self) -> Option<&str> { + self.as_ref().extra_key + } + + pub fn cache_salt(&self) -> Option<&str> { + self.as_ref().cache_salt + } + + pub fn cache_salt_arc(&self) -> Option> { + self.0 + .as_ref() + .and_then(|namespace| namespace.cache_salt.clone()) + } +} + +impl PartialEq for KeyNamespace { + fn eq(&self, other: &Self) -> bool { + self.as_ref() == other.as_ref() + } +} + +impl Eq for KeyNamespace {} + +impl Hash for KeyNamespace { + fn hash(&self, state: &mut H) { + self.as_ref().hash(state); + } +} + +type ChildMap = HashBrownMap<(KeyNamespace, K), NodeIdx_, RandomState>; + +/// Borrowed view of a namespaced child edge used for allocation-free lookup. +struct ChildEdgeRef<'a, K: ChildKeyType> { + namespace: KeyNamespaceRef<'a>, + page: &'a [K::Atom], +} + +impl Hash for ChildEdgeRef<'_, K> { + fn hash(&self, state: &mut H) { + self.namespace.hash(state); + self.page.hash(state); + } +} + +impl Equivalent<(KeyNamespace, K)> for ChildEdgeRef<'_, K> { + fn equivalent(&self, edge: &(KeyNamespace, K)) -> bool { + self.namespace == edge.0.as_ref() && self.page == edge.1.as_ref() + } +} + +/// A radix-tree node, generic over the child-key type `K` (single-token or bigram). +pub struct Node { + /// Parent handle; `None` for the root or a not-yet-attached child. + pub parent: Option, + /// Own arena slot; stamped by the arena on allocation (hand-built nodes + /// default it to the id value). + pub(crate) idx: NodeIdx_, + /// The namespace this node belongs to. + pub namespace: KeyNamespace, + /// Child edges keyed by (namespace, the child's page key); the namespace + /// component mirrors the child's namespace at every level. + pub children: ChildMap, + /// The page key labelling the edge from the parent (also this node's key in the + /// parent's `children`); empty for the root. + pub key: K, + /// Per-(component × tier) value state, indexed by `ValueSlotIdx` (device + /// slots first, host after); device states also sit at plain component + /// indices. + pub values: [ValueState; NUM_VALUE_SLOTS], + /// SWA lock-window uuid; stamped where a device lock walk fills the window. + pub swa_uuid: Option, + /// SWA lock-window uuid for host locks; stamped where a host lock walk fills the window. + pub swa_host_uuid: Option, + /// Per-page hash chain; None when the node was never hashed. + /// TODO: Store raw digests and hex-encode only at the Python or storage boundary. + pub hash_value: Option>, + /// The in-flight write-through backup's ack id. + pub write_through_pending_id: Option, + /// Load-back anchor currently reading this node's host slots. + pub load_back_pending_id: Option, + /// Monotonic access tick for LRU ordering (exact; not wall-clock). + pub last_access_counter: i64, + /// Tick stamped at construction. + pub creation_counter: i64, + /// Match hits accumulated for write-through and LFU decisions. + pub hit_count: i64, + /// Eviction priority; the root uses `i64::MIN` and is never a leaf. + pub priority: i64, + /// This node's external handle; minted once, never recycled. + pub id: NodeId, +} + +impl Node { + /// Whether this is the root (no parent). + pub fn is_root(&self) -> bool { + self.parent.is_none() + } + + /// The parent's id; panics on a root. + #[track_caller] + pub fn parent(&self) -> NodeIdx_ { + self.parent + .unwrap_or_else(|| panic!("node {} is a root and has no parent", self.id)) + } + + /// The parent's id, or `None` on a root. + pub fn try_parent(&self) -> Option { + self.parent + } + + /// Whether this node has no children (a tree leaf). + pub fn is_leaf(&self) -> bool { + self.children.is_empty() + } + + /// Tree-level: Full KV not on device (non-root with value=None). + pub fn evicted(&self) -> bool { + self.parent.is_some() && !self.has_device_value(FULL) + } + + /// Tree-level: Full KV present on host. + pub fn backuped(&self) -> bool { + self.has_host_value(FULL) + } + + /// The last page's hash value, or None when the node was never hashed. + pub fn get_last_hash_value(&self) -> Option<&str> { + self.hash_value + .as_ref() + .and_then(|h| h.last()) + .map(String::as_str) + } + + /// The component's device value; panics if unset. + pub fn device_value(&self, component_type: ComponentType) -> &Tensor { + self.value_(ValueSlotIdx::device(component_type)) + } + + /// The component's device value, or None when unset. + pub fn try_device_value(&self, component_type: ComponentType) -> Option<&Tensor> { + self.try_value_(ValueSlotIdx::device(component_type)) + } + + /// Whether the component's device value is present. + pub fn has_device_value(&self, component_type: ComponentType) -> bool { + self.has_value_(ValueSlotIdx::device(component_type)) + } + + /// The component's device value length, or 0 when value-less. + pub fn device_value_len(&self, component_type: ComponentType) -> usize { + self.value_len_(ValueSlotIdx::device(component_type)) + } + + /// Set the component's device value; panics if already set. + pub fn set_device_value(&mut self, component_type: ComponentType, value: Tensor) { + self.set_value_(ValueSlotIdx::device(component_type), value); + } + + /// Take the component's device value; panics if unset. + pub fn take_device_value(&mut self, component_type: ComponentType) -> Tensor { + self.take_value_(ValueSlotIdx::device(component_type)) + } + + /// The component's device lock refcount. + pub fn device_lock_ref(&self, component_type: ComponentType) -> u32 { + self.lock_ref_(ValueSlotIdx::device(component_type)) + } + + /// Bump the component's device lock refcount by one. + pub fn inc_device_lock_ref(&mut self, component_type: ComponentType) { + self.inc_lock_ref_(ValueSlotIdx::device(component_type)); + } + + /// Drop the component's device lock refcount by one; panics when unlocked. + pub fn dec_device_lock_ref(&mut self, component_type: ComponentType) { + self.dec_lock_ref_(ValueSlotIdx::device(component_type)); + } + + /// Copy the component's device lock refcount from `src_node`. + pub fn copy_device_lock_ref(&mut self, component_type: ComponentType, src_node: &Node) { + let slot = ValueSlotIdx::device(component_type); + self.set_lock_ref_(slot, src_node.lock_ref_(slot)); + } + + /// Split the component's device value between a new parent and the child. + pub fn redistribute_child_device_value( + parent_node: &mut Node, + child_node: &mut Node, + component_type: ComponentType, + split_len: i64, + ) { + Node::redistribute_child_value_( + parent_node, + child_node, + ValueSlotIdx::device(component_type), + split_len, + ); + } + + /// The component's host value; panics if unset. + pub fn host_value(&self, component_type: ComponentType) -> &Tensor { + self.value_(ValueSlotIdx::host(component_type)) + } + + /// The component's host value, or None when unset. + pub fn try_host_value(&self, component_type: ComponentType) -> Option<&Tensor> { + self.try_value_(ValueSlotIdx::host(component_type)) + } + + /// Whether the component's host value is present. + pub fn has_host_value(&self, component_type: ComponentType) -> bool { + self.has_value_(ValueSlotIdx::host(component_type)) + } + + /// The component's host value length, or 0 when value-less. + pub fn host_value_len(&self, component_type: ComponentType) -> usize { + self.value_len_(ValueSlotIdx::host(component_type)) + } + + /// Set the component's host value; panics if already set. + pub fn set_host_value(&mut self, component_type: ComponentType, value: Tensor) { + self.set_value_(ValueSlotIdx::host(component_type), value); + } + + /// Take the component's host value; panics if unset. + pub fn take_host_value(&mut self, component_type: ComponentType) -> Tensor { + self.take_value_(ValueSlotIdx::host(component_type)) + } + + /// The component's host lock refcount (the host lock). + pub fn host_lock_ref(&self, component_type: ComponentType) -> u32 { + self.lock_ref_(ValueSlotIdx::host(component_type)) + } + + /// Bump the component's host lock refcount by one. + pub fn inc_host_lock_ref(&mut self, component_type: ComponentType) { + self.inc_lock_ref_(ValueSlotIdx::host(component_type)); + } + + /// Drop the component's host lock refcount by one; panics when unlocked. + pub fn dec_host_lock_ref(&mut self, component_type: ComponentType) { + self.dec_lock_ref_(ValueSlotIdx::host(component_type)); + } + + /// Split the component's host value between a new parent and the child. + pub fn redistribute_child_host_value( + parent_node: &mut Node, + child_node: &mut Node, + component_type: ComponentType, + split_len: i64, + ) { + Node::redistribute_child_value_( + parent_node, + child_node, + ValueSlotIdx::host(component_type), + split_len, + ); + } + + /// Whether any component holds a device lock on this node. + pub fn is_device_locked(&self) -> bool { + self.values[..NUM_COMPONENT_TYPES] + .iter() + .any(|state| state.lock_ref > 0) + } + + /// Whether any component holds a host lock on this node. + pub fn is_host_locked(&self) -> bool { + self.values[NUM_COMPONENT_TYPES..] + .iter() + .any(|state| state.lock_ref > 0) + } + + /// Whether an in-flight load-back currently pins this node. + pub fn is_load_back_pending(&self) -> bool { + self.load_back_pending_id.is_some() + } + + // ==== Crate-internal tree wiring ==== + + /// A fresh root: no parent, no value, lowest eviction priority. + pub(crate) fn new_root(id: NodeId) -> Self { + Node { + parent: None, + namespace: KeyNamespace::default(), + children: ChildMap::with_hasher(RandomState::new()), + key: K::default(), + values: Default::default(), + swa_uuid: None, + swa_host_uuid: None, + hash_value: Some(Vec::new()), + write_through_pending_id: None, + load_back_pending_id: None, + last_access_counter: 0, + creation_counter: 0, + hit_count: 0, + priority: i64::MIN, + id, + idx: NodeIdx_(id), + } + } + + /// A detached child reached by edge `key`; `attach_child` sets the parent link. + pub(crate) fn new_child(id: NodeId, key: K, priority: i64) -> Self { + Node { + parent: None, + namespace: KeyNamespace::default(), + children: ChildMap::with_hasher(RandomState::new()), + key, + values: Default::default(), + swa_uuid: None, + swa_host_uuid: None, + hash_value: None, + write_through_pending_id: None, + load_back_pending_id: None, + last_access_counter: 0, + creation_counter: 0, + hit_count: 0, + priority, + id, + idx: NodeIdx_(id), + } + } + + /// The namespaced edge key for this node's own edge from its parent. + pub(crate) fn edge_key(&self, page_size: usize) -> (KeyNamespace, K) { + (self.namespace.clone(), self.key.child_key(page_size)) + } + + /// Link `child` under this node, keyed by its namespaced page key; errors on + /// a duplicate key. Panics if `child` is already attached (an internal invariant). + /// The caller sets `child.namespace` first; the edge key mirrors it. + pub(crate) fn attach_child( + &mut self, + child: &mut Node, + page_size: usize, + ) -> Result<(), TreeCoreRuntimeError> { + // Only fresh, detached nodes are ever attached. + assert!( + child.parent.is_none(), + "attach_child: node {} is already attached to parent {:?}", + child.id, + child.parent + ); + let parent_idx = self.idx; + match self.children.entry(child.edge_key(page_size)) { + Entry::Occupied(_) => Err(TreeCoreRuntimeError::DuplicateChildKey { + parent: self.id, + key: format!("{:?}", child.key), + }), + Entry::Vacant(slot) => { + slot.insert(child.idx); + child.parent = Some(parent_idx); + Ok(()) + } + } + } + + /// Unlink this node from `parent`: drop it from `parent.children` and clear its own + /// parent link; panics on a broken parent<->child link (an internal invariant). + pub(crate) fn detach_from_parent(&mut self, parent: &mut Node, page_size: usize) { + // A live child is always registered under its namespaced page key. + match parent.children.remove(&self.edge_key(page_size)) { + Some(idx) if idx == self.idx => self.parent = None, + found => panic!( + "detach_from_parent: parent {} has no child {} under key {:?} (found {:?})", + parent.id, self.id, self.key, found + ), + } + } + + // ==== Internal slot-keyed lookups ==== + + /// The slot's value state (internal slot-keyed lookup). + pub fn state_(&self, slot: ValueSlotIdx) -> &ValueState { + &self.values[slot.idx()] + } + + /// The slot's mutable value state (internal slot-keyed lookup). + pub fn state_mut_(&mut self, slot: ValueSlotIdx) -> &mut ValueState { + &mut self.values[slot.idx()] + } + + /// The slot's value, or None when unset (internal slot-keyed lookup). + pub fn try_value_(&self, slot: ValueSlotIdx) -> Option<&Tensor> { + self.state_(slot).value.as_ref() + } + + /// The slot's value; panics if no value is set (internal slot-keyed lookup). + pub fn value_(&self, slot: ValueSlotIdx) -> &Tensor { + self.try_value_(slot).unwrap_or_else(|| { + panic!( + "value: {:?}/{} slot has no value on node {}", + slot.component_type(), + slot.tier(), + self.id + ) + }) + } + + /// Whether the slot's value is present (internal slot-keyed lookup). + pub fn has_value_(&self, slot: ValueSlotIdx) -> bool { + self.state_(slot).value.is_some() + } + + /// The slot's value length, or 0 when value-less (internal slot-keyed lookup). + pub fn value_len_(&self, slot: ValueSlotIdx) -> usize { + self.state_(slot) + .value + .as_ref() + .map_or(0, |v| v.size()[0] as usize) + } + + /// Set the slot's value; panics if a value is already set (internal slot-keyed lookup). + pub fn set_value_(&mut self, slot: ValueSlotIdx, value: Tensor) { + if slot.component_type().single_value_per_node() { + assert_eq!( + value.size()[0], + 1, + "set_value: {:?}/{} expects a single state slot on node {}", + slot.component_type(), + slot.tier(), + self.id + ); + } else { + assert_eq!( + value.size()[0] as usize, + self.key.atom_len(), + "set_value: {:?}/{} value length differs from the key on node {}", + slot.component_type(), + slot.tier(), + self.id + ); + } + let node_id = self.id; + let state = self.state_mut_(slot); + assert!( + state.value.is_none(), + "set_value: {:?}/{} slot already set on node {node_id}", + slot.component_type(), + slot.tier() + ); + state.value = Some(value); + } + + /// Take the slot's value, leaving it value-less; panics if no value is set + /// (internal slot-keyed lookup). + pub fn take_value_(&mut self, slot: ValueSlotIdx) -> Tensor { + let node_id = self.id; + self.state_mut_(slot).value.take().unwrap_or_else(|| { + panic!( + "take_value: {:?}/{} slot has no value on node {node_id}", + slot.component_type(), + slot.tier() + ) + }) + } + + /// Split the slot's value on `child_node` at `split_len`: the deep-copied head + /// rows land on `parent_node`, the tail rows replace the value on `child_node` + /// (internal slot-keyed lookup). + pub fn redistribute_child_value_( + parent_node: &mut Node, + child_node: &mut Node, + slot: ValueSlotIdx, + split_len: i64, + ) { + let child_node_id = child_node.id; + let value = child_node.take_value_(slot); + let len = value.size()[0]; + // A boundary split would leave one side with a present-but-empty value. + assert!( + 0 < split_len && split_len < len, + "redistribute_child_value: split_len {split_len} out of range (0, {len}) on node {child_node_id}" + ); + let head = value.narrow(0, 0, split_len).copy(); + let tail = value.narrow(0, split_len, len - split_len).copy(); + child_node.set_value_(slot, tail); + parent_node.set_value_(slot, head); + } + + /// The slot's lock refcount (internal slot-keyed lookup). + pub fn lock_ref_(&self, slot: ValueSlotIdx) -> u32 { + self.state_(slot).lock_ref + } + + /// Set the slot's lock refcount (internal slot-keyed lookup). + pub fn set_lock_ref_(&mut self, slot: ValueSlotIdx, lock_ref: u32) { + self.state_mut_(slot).lock_ref = lock_ref; + } + + /// Bump the slot's lock refcount by one (internal slot-keyed lookup). + pub fn inc_lock_ref_(&mut self, slot: ValueSlotIdx) { + self.state_mut_(slot).lock_ref += 1; + } + + /// Drop the slot's lock refcount by one; panics on an unlocked node + /// (internal slot-keyed lookup). + pub fn dec_lock_ref_(&mut self, slot: ValueSlotIdx) { + let node_id = self.id; + let state = self.state_mut_(slot); + state.lock_ref = state.lock_ref.checked_sub(1).unwrap_or_else(|| { + panic!( + "dec_lock_ref: {:?}/{} lock_ref underflow on node {node_id}", + slot.component_type(), + slot.tier() + ) + }); + } +} + +// Node handles and per-slot value state. + +/// External node handle — the only node identity that crosses the FFI. +/// Minted monotonically and never recycled, so a freed node's id can never +/// alias a later allocation (the arena's id map is the ABA guard). +pub type NodeId = usize; + +/// Internal arena slot index; recycled by the freelist and never crosses the FFI. +#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)] +pub(crate) struct NodeIdx_(pub(crate) usize); + +impl std::fmt::Display for NodeIdx_ { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.0) + } +} + +/// Number of (component × tier) value slots on a node: device tier first, host after. +pub const NUM_VALUE_SLOTS: usize = 2 * NUM_COMPONENT_TYPES; + +/// Flat index of one (component × tier) value slot in a node's `values` array. +/// All slot arithmetic lives here; nothing else computes raw offsets. +#[derive(Copy, Clone, PartialEq, Eq, Debug)] +pub struct ValueSlotIdx(usize); + +impl ValueSlotIdx { + /// The component's device-tier slot. + pub const fn device(component_type: ComponentType) -> Self { + Self(component_type.idx()) + } + + /// The component's host-tier slot. + pub const fn host(component_type: ComponentType) -> Self { + Self(NUM_COMPONENT_TYPES + component_type.idx()) + } + + /// Index into a node's `values` array. + pub const fn idx(self) -> usize { + self.0 + } + + /// The slot at a flat index; panics out of range. + pub fn from_idx(idx: usize) -> Self { + assert!( + idx < NUM_VALUE_SLOTS, + "from_idx: {idx} is not a value-slot index" + ); + Self(idx) + } + + /// Whether this is a host-tier slot. + pub const fn is_host(self) -> bool { + self.0 >= NUM_COMPONENT_TYPES + } + + /// The component this slot belongs to. + pub fn component_type(self) -> ComponentType { + ComponentType::from_idx(self.0 % NUM_COMPONENT_TYPES) + } + + /// The slot's tier name, for diagnostics. + pub fn tier(self) -> &'static str { + if self.is_host() { "host" } else { "device" } + } +} + +/// Per-(component × tier) node state: the KV-index `value` (held opaquely) +/// and the in-flight `lock_ref`. +#[derive(Default)] +pub struct ValueState { + /// KV pool indices; `None` = value-less (root) or tombstone (evicted / out-of-window). + pub value: Option, + /// In-flight request refcount; a host slot's own `lock_ref` IS the host lock. + pub lock_ref: u32, +} + +// Tree-core runtime errors. + +/// Errors surfaced from the tree-core runtime API when a caller violates a documented +/// contract (freeing an unallocated node, allocating under a freed parent). +#[allow(clippy::enum_variant_names)] +#[derive(Debug, thiserror::Error)] +pub enum TreeCoreRuntimeError { + /// A public NodeId no longer names a live arena node. + #[error("node {node_id} is not allocated")] + NodeNotAllocated { node_id: NodeId }, + /// `begin_insert`/`insert` called while a resumable insert is suspended. + #[error("concurrent insert walks")] + ConcurrentInsertWalk, + /// `resume_insert` called without a suspended insert. + #[error("no in-flight insert")] + NoInFlightInsert, + /// A `NodeIdx_` beyond the arena's bounds — never allocated. `size` is the + /// arena's current slot count, so valid ids are `0..size`. + #[error("node access out of bounds: id {id} not in [0, {size})")] + NodeAccessOutOfBound { id: NodeIdx_, size: usize }, + /// `free` called on an in-range slot that is already free — a double free. + #[error("double free: node {id} is already free")] + NodeDoubleFree { id: NodeIdx_ }, + /// `free` called on a root (no parent) — roots are protected. + #[error("cannot free node {id}: it is a root (protected)")] + RootNotFreeable { id: NodeIdx_ }, + /// `free` called on a non-leaf node — only leaves are freeable. + #[error("cannot free non-leaf node {id}: it has {num_children} children")] + FreeNonLeafNode { id: NodeIdx_, num_children: usize }, + /// `alloc_child` under an in-range parent slot that holds no live node. + #[error("alloc_child: parent {id} is not allocated")] + ParentNotAllocated { id: NodeIdx_ }, + /// `alloc_child` under a parent that already has a child at the same key. + #[error("cannot add a child under parent {parent}: key {key} already has a child")] + DuplicateChildKey { parent: NodeId, key: String }, + /// `demote` requires a device-resident Full value with a completed host backup. + #[error( + "cannot demote node {node_id}: expected device-resident Full value with host backup (evicted={evicted}, backuped={backuped})" + )] + InvalidDemoteState { + node_id: NodeId, + evicted: bool, + backuped: bool, + }, + /// SWA load-back cannot cross a node that has no value on either tier. + #[error("SWA load-back traversal reached node {node_id} without a device or host value")] + SwaLoadBackMissingValue { node_id: NodeId }, + /// A host insert below a non-root anchor must remain in that anchor's namespace. + #[error("insert_host namespace does not match non-root anchor {node_id}")] + InsertHostNamespaceMismatch { node_id: NodeId }, +} + +// Unigram and bigram child keys. + +/// An owned child key (one radix page); `Atom` is the per-position token (`i64` single, +/// `(i64, i64)` bigram/EAGLE). `Borrow<[Atom]>` lets a borrowed page slice +/// query a `HashMap` keyed by owned keys. +pub trait ChildKeyType: + Clone + + Eq + + Hash + + Default + + Debug + + AsRef<[Self::Atom]> + + Borrow<[Self::Atom]> + + From> +{ + type Atom: Copy + Eq + Hash + Send + Sync; + + /// Whether this key represents overlapping token bigrams rather than + /// individual tokens. This is type metadata, not per-node state. + const IS_BIGRAM: bool; + + /// The key over the boundary's raw token ids; ownership passes straight + /// through, so the unigram key never copies. + fn key_from(token_ids: Cow<'_, Vec>) -> Cow<'_, Self>; + + /// The atom's token ids as u32 storage-hash words. + fn hash_words(atom: &Self::Atom) -> impl Iterator; + + /// The raw token ids spanned by `atoms`; the unigram view borrows, bigram + /// atoms (overlapping by one) materialize. + fn raw_token_ids(atoms: &[Self::Atom]) -> Cow<'_, [i64]>; + + /// Number of atoms in this key; an atom is one radix position — a token normally, + /// a token pair for bigram/EAGLE keys. + fn atom_len(&self) -> usize { + self.as_ref().len() + } + + /// The first radix page (`page_size` atoms) as an owned key — a node's + /// child-map key under its parent; panics on a key shorter than a page. + fn child_key(&self, page_size: usize) -> Self { + let atom_len = self.atom_len(); + assert!( + atom_len >= page_size, + "child_key: key of {atom_len} atoms is shorter than a page ({page_size})" + ); + self.as_ref()[..page_size].to_vec().into() + } + + /// The single page starting at `start`, zero-copy. + fn page_at(&self, start: usize, page_size: usize) -> &[Self::Atom] { + let end = start + page_size; + let atom_len = self.atom_len(); + assert!( + end <= atom_len, + "page_at: page [{start}, {end}) reaches beyond the key length {atom_len}" + ); + &self.as_ref()[start..end] + } + + /// Page-quantized common-prefix length of the tail from `start` with `other`. + fn match_len(&self, start: usize, other: &Self, page_size: usize) -> usize { + let atom_len = self.atom_len(); + assert!( + start <= atom_len, + "match_len: start {start} beyond the key length {atom_len}" + ); + let common = self.as_ref()[start..] + .iter() + .zip(other.as_ref()) + .take_while(|(a, b)| a == b) + .count(); + common / page_size * page_size + } + + /// The owned suffix from `start`; empty when `start` equals the length. + fn suffix(&self, start: usize) -> Self { + let atom_len = self.atom_len(); + assert!( + start <= atom_len, + "suffix: start {start} beyond the key length {atom_len}" + ); + self.as_ref()[start..].to_vec().into() + } + + /// The key truncated to a whole number of pages. + fn page_aligned(&self, page_size: usize) -> Self { + let aligned_len = self.atom_len() / page_size * page_size; + self.as_ref()[..aligned_len].to_vec().into() + } + + /// Split into (head, tail) owned keys at `split_idx`; panics on a boundary + /// split, which would leave one side empty. + fn split_at(&self, split_idx: usize) -> (Self, Self) { + let atom_len = self.atom_len(); + assert!( + 0 < split_idx && split_idx < atom_len, + "split_at: split_idx {split_idx} out of range (0, {atom_len})" + ); + let (head, tail) = self.as_ref().split_at(split_idx); + (head.to_vec().into(), tail.to_vec().into()) + } +} + +/// A token id as a u32 storage-hash word; token ids beyond u32 are rejected. +fn hash_word(token_id: i64) -> u32 { + u32::try_from(token_id).expect("token id does not fit in uint32") +} + +impl ChildKeyType for Vec { + type Atom = i64; + const IS_BIGRAM: bool = false; + + fn key_from(token_ids: Cow<'_, Vec>) -> Cow<'_, Self> { + token_ids + } + + fn hash_words(atom: &i64) -> impl Iterator { + std::iter::once(hash_word(*atom)) + } + + fn raw_token_ids(atoms: &[i64]) -> Cow<'_, [i64]> { + Cow::Borrowed(atoms) + } +} + +impl ChildKeyType for Vec<(i64, i64)> { + type Atom = (i64, i64); + const IS_BIGRAM: bool = true; + + /// N+1 raw token ids become N overlapping (t_i, t_{i+1}) bigram atoms. + fn key_from(token_ids: Cow<'_, Vec>) -> Cow<'_, Self> { + Cow::Owned(token_ids.windows(2).map(|w| (w[0], w[1])).collect()) + } + + fn hash_words(atom: &(i64, i64)) -> impl Iterator { + [hash_word(atom.0), hash_word(atom.1)].into_iter() + } + + fn raw_token_ids(atoms: &[(i64, i64)]) -> Cow<'_, [i64]> { + let Some(first) = atoms.first() else { + return Cow::Owned(Vec::new()); + }; + Cow::Owned( + std::iter::once(first.0) + .chain(atoms.iter().map(|atom| atom.1)) + .collect(), + ) + } +} + +// Per-page hash chains. + +pub(crate) const DIGEST_LEN: usize = 32; +pub(crate) type HashDigest = [u8; DIGEST_LEN]; + +/// SHA256(prior_digest || page atom words as little-endian u32 bytes). +pub(crate) fn hash_page( + page: &[K::Atom], + prior: Option<&HashDigest>, +) -> HashDigest { + let mut hasher = Sha256::new(); + if let Some(prior) = prior { + hasher.update(prior); + } + for atom in page { + for word in K::hash_words(atom) { + hasher.update(word.to_le_bytes()); + } + } + hasher.finalize().into() +} + +/// Lowercase-hex encoding of a digest. +fn digest_to_hex(digest: &HashDigest) -> String { + const HEX: &[u8; 16] = b"0123456789abcdef"; + let mut out = String::with_capacity(DIGEST_LEN * 2); + for byte in digest { + out.push(HEX[(byte >> 4) as usize] as char); + out.push(HEX[(byte & 0x0f) as usize] as char); + } + out +} + +/// Decode a chained-in hex hash back to its raw digest. +fn parse_prior_hash(prior_hash: &str) -> HashDigest { + let bytes = prior_hash.as_bytes(); + assert_eq!( + bytes.len(), + DIGEST_LEN * 2, + "prior hash must be a 64-char hex digest" + ); + let nibble = |b: u8| -> u8 { + match b { + b'0'..=b'9' => b - b'0', + b'a'..=b'f' => b - b'a' + 10, + b'A'..=b'F' => b - b'A' + 10, + _ => panic!("prior hash contains a non-hex character"), + } + }; + let mut digest = [0u8; DIGEST_LEN]; + for (i, out) in digest.iter_mut().enumerate() { + *out = (nibble(bytes[i * 2]) << 4) | nibble(bytes[i * 2 + 1]); + } + digest +} + +/// Per-page chained hashes over key atoms, seeded from an optional prior hex hash. +pub fn get_hash_str( + atoms: &[K::Atom], + prior_hash: Option<&str>, + page_size: usize, +) -> Vec { + assert!(page_size > 0, "page_size must be positive"); + // An empty prior chains nothing. + let mut prior = prior_hash + .filter(|prior| !prior.is_empty()) + .map(parse_prior_hash); + let mut hash_values = Vec::with_capacity(atoms.len().div_ceil(page_size)); + for page in atoms.chunks(page_size) { + let digest = hash_page::(page, prior.as_ref()); + hash_values.push(digest_to_hex(&digest)); + prior = Some(digest); + } + hash_values +} + +/// Per-page chained raw digests, seeded from an optional prior digest. +pub(crate) fn get_hash_digests( + atoms: &[K::Atom], + prior: Option<&HashDigest>, + page_size: usize, +) -> Vec { + assert!(page_size > 0, "page_size must be positive"); + let mut prior = prior.copied(); + let mut digests = Vec::with_capacity(atoms.len().div_ceil(page_size)); + for page in atoms.chunks(page_size) { + let digest = hash_page::(page, prior.as_ref()); + digests.push(digest); + prior = Some(digest); + } + digests +} + +/// The hash's first 16 hex chars as a signed 64-bit block id for events. +pub fn hash_str_to_int64(hash_str: &str) -> i64 { + u64::from_str_radix(&hash_str[..16], 16).expect("hash must be a hex digest") as i64 +} + +/// The raw digest's first eight bytes as the event protocol's signed i64. +pub(crate) fn hash_digest_to_int64(digest: &HashDigest) -> i64 { + i64::from_be_bytes(digest[..8].try_into().expect("digest has eight bytes")) +} + +/// Split a node's hash list at a page boundary; None-safe when never hashed. +pub fn split_node_hash_value( + hash_values: Option>, + split_idx: usize, + page_size: usize, +) -> (Option>, Option>) { + let Some(mut new_node_hash) = hash_values else { + return (None, None); + }; + let child_hash = new_node_hash.split_off(split_idx / page_size); + // Progressive splits must not retain the pre-split capacity on the head. + new_node_hash.shrink_to_fit(); + (Some(new_node_hash), Some(child_hash)) +} + +// Node arena storage. + +/// Owns every `Node`; parent/children/LRU hold `NodeIdx_`s into it, with a freelist +/// and a single root; child edges are keyed by (namespace, page key). +pub struct NodeArena { + /// Node store indexed by `NodeIdx_`; `None` marks a freed slot. + nodes: Vec>>, + /// Freed slot ids available for reuse. + free: Vec, + /// External handle -> live slot; a freed `NodeId` leaves the map, so a + /// stale handle can never alias a recycled slot. + id_map: HashMap, + /// Next external handle; monotonic, never recycled (survives `reset`). + next_id: NodeId, + /// The tree's single root; every namespace hangs off it. + root: NodeIdx_, + /// Monotonic counter stamped into `Node::last_access_counter`. + access_counter: i64, + /// The component types this tree runs (every root is locked for each). + component_types: Vec, + /// Atoms per radix page; children are keyed by their key's first page. + page_size: usize, +} + +impl NodeArena { + /// Build an arena for the given component types and install a fresh root. + pub fn new(component_types: Vec, page_size: usize) -> Self { + let mut arena = NodeArena { + nodes: Vec::new(), + free: Vec::new(), + id_map: HashMap::new(), + next_id: 0, + root: NodeIdx_(0), + access_counter: 0, + component_types, + page_size, + }; + arena.reset(); + arena + } + + /// Drop all nodes, then reinstall the root. + pub fn reset(&mut self) { + self.nodes.clear(); + self.free.clear(); + // next_id is NOT reset: pre-reset handles must miss, never alias. + self.id_map.clear(); + self.access_counter = 0; + self.root = self.alloc_root(); + } + + /// The live slot for an external handle; panics on a freed or unknown id. + #[track_caller] + pub fn resolve(&self, id: NodeId) -> NodeIdx_ { + *self + .id_map + .get(&id) + .unwrap_or_else(|| panic!("node {id} is not allocated")) + } + + /// The live slot for an external handle, or None if freed/unknown. + pub fn try_resolve(&self, id: NodeId) -> Option { + self.id_map.get(&id).copied() + } + + /// Mint the next external handle for the slot and index it. + fn mint_id_(&mut self, idx: NodeIdx_) -> NodeId { + let id = self.next_id; + self.next_id += 1; + self.id_map.insert(id, idx); + id + } + + /// Allocate a protected, value-less root: locked (`lock_ref = 1`) for each + /// component type and never entering a leaf/LRU set. + pub fn alloc_root(&mut self) -> NodeIdx_ { + let idx = self.reserve(); + let id = self.mint_id_(idx); + let tick = self.get_and_bump_access_counter(); + let node = self.nodes[idx.0].insert(Node::new_root(id)); + node.idx = idx; + node.last_access_counter = tick; + node.creation_counter = tick; + for ct in &self.component_types { + node.values[ct.idx()].lock_ref = 1; + } + idx + } + + /// Every live node id, in slot order. + pub fn live_ids(&self) -> impl Iterator + '_ { + self.nodes + .iter() + .enumerate() + .filter_map(|(idx, slot)| slot.as_ref().map(|_| NodeIdx_(idx))) + } + + /// Per-page hash values for a node's key, chained from its parent's last hash. + pub fn compute_node_hash_values(&self, node_id: NodeIdx_, page_size: usize) -> Vec { + let node = self.node(node_id); + let parent_hash = node.parent.and_then(|parent_id| { + let parent = self.node(parent_id); + if parent.key.atom_len() > 0 { + parent.get_last_hash_value() + } else { + None + } + }); + crate::node::get_hash_str::(node.key.as_ref(), parent_hash, page_size) + } + + /// The ancestor chain's hash values ending at `node_id`, in root-to-node + /// order; the walk stops below the nearest never-hashed ancestor. + pub fn prefix_hash_values(&self, node_id: Option) -> Vec { + let mut chunks: Vec<&Vec> = Vec::new(); + let mut cursor = node_id; + while let Some(id) = cursor { + let node = self.node(id); + let Some(hash_value) = node.hash_value.as_ref() else { + break; + }; + chunks.push(hash_value); + cursor = node.parent; + } + chunks + .iter() + .rev() + .flat_map(|chunk| chunk.iter().cloned()) + .collect() + } + + /// The node's caller-defined namespace key; None for the default. + pub fn node_extra_key(&self, node_id: NodeIdx_) -> Option<&str> { + self.node(node_id).namespace.extra_key() + } + + /// The node's cache salt; None for an unsalted namespace. + pub fn node_cache_salt(&self, node_id: NodeIdx_) -> Option<&str> { + self.node(node_id).namespace.cache_salt() + } + + /// The tree's single root. + pub fn root(&self) -> NodeIdx_ { + self.root + } + + /// The node's child on the page within the namespace, if any. + pub fn child_on_page( + &self, + id: NodeIdx_, + extra_key: Option<&str>, + page: &[K::Atom], + ) -> Option { + self.child_on_page_in_namespace( + id, + KeyNamespaceRef::new(extra_key, /* cache_salt = */ None), + page, + ) + } + + /// The node's child on the page within the full radix namespace, if any. + pub fn child_on_page_in_namespace( + &self, + id: NodeIdx_, + namespace: KeyNamespaceRef<'_>, + page: &[K::Atom], + ) -> Option { + self.node(id) + .children + .get(&ChildEdgeRef:: { namespace, page }) + .copied() + } + + /// The root's child on the key's first page within the namespace, if any. + pub fn root_child(&self, extra_key: Option<&str>, page: &[K::Atom]) -> Option { + self.child_on_page(self.root, extra_key, page) + } + + /// The root's child on the key's first page within the full radix namespace. + pub fn root_child_in_namespace( + &self, + namespace: KeyNamespaceRef<'_>, + page: &[K::Atom], + ) -> Option { + self.child_on_page_in_namespace(self.root, namespace, page) + } + + /// Whether any root edge files under the namespace. + pub fn namespace_exists(&self, extra_key: Option<&str>) -> bool { + self.full_namespace_exists(KeyNamespaceRef::new( + extra_key, /* cache_salt = */ None, + )) + } + + /// Whether any root edge files under the full radix namespace. + pub fn full_namespace_exists(&self, namespace: KeyNamespaceRef<'_>) -> bool { + self.node(self.root) + .children + .keys() + .any(|(stored, _)| stored.as_ref() == namespace) + } + + /// Install `child` under `parent` on its namespaced `map_key`; returns the + /// displaced child, if any. The key's namespace mirrors the child's. + pub fn insert_child_edge( + &mut self, + parent: NodeIdx_, + map_key: K, + child: NodeIdx_, + ) -> Option { + let namespace = self.node(child).namespace.clone(); + self.node_mut(parent) + .children + .insert((namespace, map_key), child) + } + + /// Reserve a detached child and attach it under `parent`; `extra_key` names the + /// namespace for root children (deeper nodes inherit the parent's). + pub fn alloc_child( + &mut self, + parent: NodeIdx_, + key: K, + priority: i64, + extra_key: Option<&str>, + ) -> Result { + self.alloc_child_in_namespace( + parent, + key, + priority, + KeyNamespaceRef::new(extra_key, /* cache_salt = */ None), + ) + } + + /// Reserve and attach a child in the full radix namespace. + pub fn alloc_child_in_namespace( + &mut self, + parent: NodeIdx_, + key: K, + priority: i64, + namespace: KeyNamespaceRef<'_>, + ) -> Result { + // Validate the parent and attach the child before committing a slot, so a + // rejected add reserves nothing. + let size = self.nodes.len(); + match self.nodes.get(parent.0) { + None => return Err(TreeCoreRuntimeError::NodeAccessOutOfBound { id: parent, size }), + Some(None) => return Err(TreeCoreRuntimeError::ParentNotAllocated { id: parent }), + Some(Some(_)) => {} + } + let idx = self + .free + .last() + .copied() + .unwrap_or(NodeIdx_(self.nodes.len())); + let mut child_node = Node::new_child(self.next_id, key, priority); + child_node.idx = idx; + let tick = self.get_and_bump_access_counter(); + child_node.last_access_counter = tick; + child_node.creation_counter = tick; + let page_size = self.page_size; + // Root children adopt the op namespace; deeper nodes inherit the parent's. + child_node.namespace = if parent == self.root { + namespace.to_owned() + } else { + self.node(parent).namespace.clone() + }; + self.node_mut(parent) + .attach_child(&mut child_node, page_size)?; + match self.free.pop() { + Some(popped) => { + // The freed slot we peeked is still the freelist head and still free. + assert_eq!(popped, idx, "freelist head changed between peek and pop"); + assert!( + self.nodes[idx.0].is_none(), + "freelist popped a live slot {idx} (freelist corruption)" + ); + self.nodes[idx.0] = Some(child_node); + } + None => self.nodes.push(Some(child_node)), + } + self.mint_id_(idx); + Ok(idx) + } + + /// Allocate a detached node (empty key, no parent) for the tree to wire in. + pub fn alloc_detached(&mut self, priority: i64) -> NodeIdx_ { + let idx = self.reserve(); + let id = self.mint_id_(idx); + let tick = self.get_and_bump_access_counter(); + let node = self.nodes[idx.0].insert(Node::new_child(id, K::default(), priority)); + node.idx = idx; + node.last_access_counter = tick; + node.creation_counter = tick; + idx + } + + /// Reserve an empty slot (reusing a freed one when available), returning its id. + fn reserve(&mut self) -> NodeIdx_ { + match self.free.pop() { + Some(idx) => { + assert!( + self.nodes[idx.0].is_none(), + "freelist popped a live slot {idx} (freelist corruption)" + ); + idx + } + None => { + self.nodes.push(None); + NodeIdx_(self.nodes.len() - 1) + } + } + } + + /// Detach a leaf from its parent and return its slot to the freelist. + pub fn free_leaf(&mut self, id: NodeIdx_) -> Result<(), TreeCoreRuntimeError> { + let size = self.nodes.len(); + // Validate and take the leaf out in one step. + let mut child_node = match self.nodes.get_mut(id.0) { + None => return Err(TreeCoreRuntimeError::NodeAccessOutOfBound { id, size }), + Some(None) => return Err(TreeCoreRuntimeError::NodeDoubleFree { id }), + Some(Some(node)) if node.is_root() => { + return Err(TreeCoreRuntimeError::RootNotFreeable { id }); + } + Some(Some(node)) if !node.is_leaf() => { + return Err(TreeCoreRuntimeError::FreeNonLeafNode { + id, + num_children: node.children.len(), + }); + } + Some(slot) => slot.take().expect("validated non-root leaf"), + }; + // A validated non-root leaf always has a parent to unlink from. + let parent = child_node.parent(); + let page_size = self.page_size; + let parent_node = self.node_mut(parent); + child_node.detach_from_parent(parent_node, page_size); + self.id_map.remove(&child_node.id); + self.free.push(id); + Ok(()) + } + + /// Number of live nodes. + pub fn len(&self) -> usize { + self.nodes.len() - self.free.len() + } + + /// Shared access to a live node; panics on a dead or out-of-range id. + #[track_caller] + pub fn node(&self, id: NodeIdx_) -> &Node { + let size = self.nodes.len(); + self.nodes + .get(id.0) + .unwrap_or_else(|| panic!("node access out of bounds: id {id} not in [0, {size})")) + .as_ref() + .unwrap_or_else(|| panic!("node {id} is not allocated")) + } + + /// Mutable access to a live node; panics on a dead or out-of-range id. + #[track_caller] + pub fn node_mut(&mut self, id: NodeIdx_) -> &mut Node { + let size = self.nodes.len(); + self.nodes + .get_mut(id.0) + .unwrap_or_else(|| panic!("node access out of bounds: id {id} not in [0, {size})")) + .as_mut() + .unwrap_or_else(|| panic!("node {id} is not allocated")) + } + + /// The node's device value for the component; panics if unset. + pub fn device_value(&self, id: NodeIdx_, component_type: ComponentType) -> &Tensor { + self.node(id).device_value(component_type) + } + + /// The node's device value for the component, or None when unset. + pub fn try_device_value(&self, id: NodeIdx_, component_type: ComponentType) -> Option<&Tensor> { + self.node(id).try_device_value(component_type) + } + + /// Whether the node holds the component's device value. + pub fn has_device_value(&self, id: NodeIdx_, component_type: ComponentType) -> bool { + self.node(id).has_device_value(component_type) + } + + /// The node's device value length for the component, or 0 when value-less. + pub fn device_value_len(&self, id: NodeIdx_, component_type: ComponentType) -> usize { + self.node(id).device_value_len(component_type) + } + + /// Set the node's device value for the component; panics if already set. + pub fn set_device_value(&mut self, id: NodeIdx_, component_type: ComponentType, value: Tensor) { + self.node_mut(id).set_device_value(component_type, value); + } + + /// Take the node's device value for the component; panics if unset. + pub fn take_device_value(&mut self, id: NodeIdx_, component_type: ComponentType) -> Tensor { + self.node_mut(id).take_device_value(component_type) + } + + /// The node's device lock refcount for the component. + pub fn device_lock_ref(&self, id: NodeIdx_, component_type: ComponentType) -> u32 { + self.node(id).device_lock_ref(component_type) + } + + /// The node's host value for the component; panics if unset. + pub fn host_value(&self, id: NodeIdx_, component_type: ComponentType) -> &Tensor { + self.node(id).host_value(component_type) + } + + /// Whether the node holds the component's host value. + pub fn has_host_value(&self, id: NodeIdx_, component_type: ComponentType) -> bool { + self.node(id).has_host_value(component_type) + } + + /// The node's host value length for the component, or 0 when value-less. + pub fn host_value_len(&self, id: NodeIdx_, component_type: ComponentType) -> usize { + self.node(id).host_value_len(component_type) + } + + /// Set the node's host value for the component; panics if already set. + pub fn set_host_value(&mut self, id: NodeIdx_, component_type: ComponentType, value: Tensor) { + self.node_mut(id).set_host_value(component_type, value); + } + + /// Take the node's host value for the component; panics if unset. + pub fn take_host_value(&mut self, id: NodeIdx_, component_type: ComponentType) -> Tensor { + self.node_mut(id).take_host_value(component_type) + } + + /// The node's host lock refcount for the component. + pub fn host_lock_ref(&self, id: NodeIdx_, component_type: ComponentType) -> u32 { + self.node(id).host_lock_ref(component_type) + } + + /// Bump the node's device lock refcount for the component. + pub fn inc_device_lock_ref(&mut self, id: NodeIdx_, component_type: ComponentType) { + self.node_mut(id).inc_device_lock_ref(component_type); + } + + /// Drop the node's device lock refcount for the component; panics when unlocked. + pub fn dec_device_lock_ref(&mut self, id: NodeIdx_, component_type: ComponentType) { + self.node_mut(id).dec_device_lock_ref(component_type); + } + + /// Bump the node's host lock refcount for the component. + pub fn inc_host_lock_ref(&mut self, id: NodeIdx_, component_type: ComponentType) { + self.node_mut(id).inc_host_lock_ref(component_type); + } + + /// Mutable access to a live parent/child pair at once. Internal-only accessor: + /// panics on a dead id or when `child_node_id` is not a child of `parent_node_id`. + #[track_caller] + pub fn node_pair_mut( + &mut self, + parent_node_id: NodeIdx_, + child_node_id: NodeIdx_, + ) -> (&mut Node, &mut Node) { + assert_ne!( + parent_node_id, child_node_id, + "node_pair_mut: distinct nodes required, got {parent_node_id} twice" + ); + let size = self.nodes.len(); + assert!( + parent_node_id.0 < size && child_node_id.0 < size, + "node_pair_mut: id out of bounds ({parent_node_id}, {child_node_id}) vs size {size}" + ); + let [parent_slot, child_slot] = self + .nodes + .get_disjoint_mut([parent_node_id.0, child_node_id.0]) + .expect("distinct in-bounds indices"); + let parent_node = parent_slot.as_mut().expect("live node"); + let child_node = child_slot.as_mut().expect("live node"); + assert_eq!( + child_node.try_parent(), + Some(parent_node_id), + "node_pair_mut: node {child_node_id} is not a child of {parent_node_id}" + ); + (parent_node, child_node) + } + + /// Advance the access counter by `delta` ticks and return the newest one, + /// reserving the whole range for the caller to assign. + pub fn get_and_batch_bump_access_counter(&mut self, delta: i64) -> i64 { + assert!( + delta > 0, + "get_and_batch_bump_access_counter: delta {delta} must be positive" + ); + self.access_counter += delta; + self.access_counter + } + + /// Bump the access counter and return the new tick (for stamping `last_access_counter`). + pub fn get_and_bump_access_counter(&mut self) -> i64 { + self.access_counter += 1; + self.access_counter + } +} + +// Eviction-eligible node set. + +/// Set of `NodeIdx_`s with O(1) membership ops and dense iteration. +#[derive(Default)] +pub struct EvictableNodeSet { + /// Dense member list; order is unspecified (swap-remove). + nodes: Vec, + /// Each member's position in `nodes`, indexed by `NodeIdx_`. + slots: Vec>, +} + +impl EvictableNodeSet { + pub fn new() -> Self { + Default::default() + } + + /// Whether `node_id` is a member. + pub fn contains(&self, node_id: NodeIdx_) -> bool { + self.slots.get(node_id.0).copied().flatten().is_some() + } + + /// Insert `node_id`; no-op when already a member. + pub fn add(&mut self, node_id: NodeIdx_) { + if node_id.0 >= self.slots.len() { + self.slots.resize(node_id.0 + 1, None); + } + if self.slots[node_id.0].is_some() { + return; + } + self.slots[node_id.0] = Some(self.nodes.len()); + self.nodes.push(node_id); + } + + /// Remove `node_id`; no-op when not a member. + pub fn discard(&mut self, node_id: NodeIdx_) { + let Some(slot) = self.slots.get(node_id.0).copied().flatten() else { + return; + }; + self.slots[node_id.0] = None; + self.nodes.swap_remove(slot); + // The swapped-in tail member (if any) now lives at `slot`. + if let Some(&moved) = self.nodes.get(slot) { + self.slots[moved.0] = Some(slot); + } + } + + /// The members, in unspecified order. + pub fn iter(&self) -> impl Iterator + '_ { + self.nodes.iter().copied() + } + + // Test-only conveniences: production callers use add/discard/contains/iter. + #[cfg(test)] + pub fn len(&self) -> usize { + self.nodes.len() + } + + #[cfg(test)] + pub fn is_empty(&self) -> bool { + self.nodes.is_empty() + } +} +#[cfg(test)] +#[path = "tests/node.rs"] +mod tests; diff --git a/rust/mem-cache/src/python_bindings.rs b/rust/mem-cache/src/python_bindings.rs new file mode 100644 index 000000000..ce6695c6c --- /dev/null +++ b/rust/mem-cache/src/python_bindings.rs @@ -0,0 +1,3261 @@ +//! Python bindings: the `mem_cache` extension module and its TreeCore adapter. + +use std::borrow::Cow; +use std::collections::{HashMap, HashSet}; +use std::sync::Mutex; + +use pyo3::buffer::PyBuffer; +use pyo3::exceptions::{PyAssertionError, PyKeyError, PyRuntimeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::{PyBytes, PyDict, PyList}; +use tch::{Device, Kind, Tensor}; + +use crate::components::{ComponentType, FULL, MAMBA, SWA}; +use crate::node::ChildKeyType; +use crate::node::{KeyNamespaceRef, NodeId, TreeCoreRuntimeError}; +use crate::unified_tree_core::KvCacheEvent; +use crate::unified_tree_core::{ + BufferBackupSnapshot, BufferBackupState, CacheAction, CacheInitParams, CacheTransferPhase, + DecLockRefParams, EvictLayer, EvictionStepResult, InsertParams, InsertResult, InsertStepResult, + MatchPrefixParams, MatchResult, PoolHitPolicy, PoolName, PoolTransfer, PoolTransferResult, Req, + UnifiedTreeCore, +}; + +/// Parse a torch-style device string (e.g. "cpu", "cuda", "cuda:1"); a bare +/// "cuda" means index 0, so callers must resolve the index themselves. +fn parse_device(device: &str) -> PyResult { + let device = device.to_lowercase(); + if device == "cpu" { + return Ok(Device::Cpu); + } + if device == "cuda" { + return Ok(Device::Cuda(0)); + } + if let Some(index) = device.strip_prefix("cuda:") + && let Ok(index) = index.parse::() + { + return Ok(Device::Cuda(index)); + } + Err(PyValueError::new_err(format!( + "unsupported device string: {device}" + ))) +} + +/// Map a Python ComponentType value onto the Rust enum. +fn parse_component_type(component_type: u8) -> PyResult { + match component_type { + 0 => Ok(FULL), + 1 => Ok(SWA), + 2 => Ok(MAMBA), + other => Err(PyValueError::new_err(format!( + "unknown component type: {other}" + ))), + } +} + +/// Map the Python EvictLayer IntFlag value onto the Rust enum. +fn parse_evict_layer(target: u8) -> PyResult { + match target { + 1 => Ok(EvictLayer::Device), + 2 => Ok(EvictLayer::Host), + 3 => Ok(EvictLayer::All), + other => Err(PyValueError::new_err(format!( + "unknown eviction layer: {other}" + ))), + } +} + +/// Convert an expected tree-core contract failure without unwinding through PyO3. +fn tree_core_runtime_error(error: TreeCoreRuntimeError) -> PyErr { + match error { + TreeCoreRuntimeError::NodeNotAllocated { node_id } => PyKeyError::new_err(node_id), + error => PyRuntimeError::new_err(error.to_string()), + } +} + +fn tree_core_assertion_error(error: TreeCoreRuntimeError) -> PyErr { + match error { + TreeCoreRuntimeError::NodeNotAllocated { node_id } => PyKeyError::new_err(node_id), + error => PyAssertionError::new_err(error.to_string()), + } +} + +/// Map the Rust enum back onto the Python ComponentType value. +fn component_type_to_u8(component_type: ComponentType) -> u8 { + component_type as u8 +} + +/// Newtype bridging `tch::Tensor` and Python `torch.Tensor` over raw THPVariable +/// pointers (inlined from pyo3-tch, MIT/Apache-2.0, by Laurent Mazare). +pub struct PyTensor(pub Tensor); + +impl<'py> FromPyObject<'py> for PyTensor { + fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { + let ptr = ob.as_ptr() as *mut tch::python::CPyObject; + match unsafe { Tensor::pyobject_unpack(ptr) } { + Ok(Some(tensor)) => Ok(PyTensor(tensor)), + Ok(None) => Err(pyo3::exceptions::PyTypeError::new_err(format!( + "expected a torch.Tensor, got {}", + ob.get_type() + ))), + Err(err) => Err(PyValueError::new_err(format!("{err:?}"))), + } + } +} + +impl ToPyObject for PyTensor { + fn to_object(&self, py: Python<'_>) -> PyObject { + PyTensor(self.0.shallow_clone()).into_py(py) + } +} + +impl IntoPy for PyTensor { + fn into_py(self, py: Python<'_>) -> PyObject { + let ptr = self + .0 + .pyobject_wrap() + .expect("failed to wrap a tensor as torch.Tensor"); + unsafe { PyObject::from_owned_ptr(py, ptr as *mut pyo3::ffi::PyObject) } + } +} + +/// Convert a tensor into a Python-held torch.Tensor reference. +fn tensor_to_py(py: Python<'_>, tensor: Tensor) -> PyResult> { + let ptr = tensor + .pyobject_wrap() + .map_err(|err| PyValueError::new_err(format!("{err:?}")))?; + Ok(unsafe { Py::from_owned_ptr(py, ptr as *mut pyo3::ffi::PyObject) }) +} + +/// Convert a Python int64 sequence to an owned `Vec`. +fn py_array_to_vec_i64(py: Python<'_>, key: &Bound<'_, PyAny>) -> PyResult> { + // Special handling for empty keys, as empty pyarray might use + // a random address to represent empty buffer which + // non-deterministically violates alignment check + if key.len().map(|n| n == 0).unwrap_or(false) { + return Ok(Vec::new()); + } + let buffer = key.extract::>()?; + if !buffer.is_c_contiguous() { + return Err(pyo3::exceptions::PyTypeError::new_err( + "Unexpected key received, expected a C-contiguous int64 buffer \ + (e.g. array.array('q'))", + )); + } + buffer.to_vec(py) +} + +/// The tagged-tuple tag of a cache action. +fn cache_action_tag(action: &CacheAction) -> &'static str { + match action { + CacheAction::FreeDeviceKV(_) => "free_device_kv", + CacheAction::FreeDeviceKVFullOnly(_) => "free_device_kv_full_only", + CacheAction::BackupKV(_) => "backup_kv", + CacheAction::ReplaceWriteThroughOnNodeSplit { .. } => "replace_write_through_on_node_split", + CacheAction::MambaEvictExcessPathStates { .. } => "mamba_evict_excess_path_states", + CacheAction::FreeComponentDeviceSlot { .. } => "free_component_device_slot", + CacheAction::FreeComponentHostSlot { .. } => "free_component_host_slot", + CacheAction::RebuildFullToSwaMapping { .. } => "rebuild_full_to_swa_mapping", + CacheAction::RecoverSwaWithLockedFull { .. } => "recover_swa_with_locked_full", + CacheAction::SwaRebuild { .. } => "swa_rebuild", + } +} + +/// Convert a cache action into its Python tagged tuple. +fn cache_action_to_py(py: Python<'_>, action: CacheAction) -> PyResult> { + let tag = cache_action_tag(&action); + match action { + CacheAction::FreeDeviceKV(tensors) => { + let tensors = PyList::new_bound(py, tensors.into_iter().map(PyTensor)); + Ok((tag, tensors).into_py(py)) + } + CacheAction::FreeDeviceKVFullOnly(tensors) => { + let tensors = PyList::new_bound(py, tensors.into_iter().map(PyTensor)); + Ok((tag, tensors).into_py(py)) + } + CacheAction::BackupKV(backup) => Ok((tag, backup.node_ids).into_py(py)), + CacheAction::ReplaceWriteThroughOnNodeSplit { + ack_id, + old_node_id, + new_node_id, + new_child_node_id, + } => Ok((tag, ack_id, old_node_id, new_node_id, new_child_node_id).into_py(py)), + CacheAction::MambaEvictExcessPathStates { tail_node_id } => { + Ok((tag, tail_node_id).into_py(py)) + } + CacheAction::FreeComponentDeviceSlot { + component_type, + indices, + } => { + let indices = PyList::new_bound(py, indices.into_iter().map(PyTensor)); + Ok((tag, component_type_to_u8(component_type), indices).into_py(py)) + } + CacheAction::FreeComponentHostSlot { + component_type, + host_indices, + } => { + let host_indices = PyList::new_bound(py, host_indices.into_iter().map(PyTensor)); + Ok((tag, component_type_to_u8(component_type), host_indices).into_py(py)) + } + CacheAction::RebuildFullToSwaMapping { + full_indices, + swa_indices, + } => { + let full_indices = PyList::new_bound(py, full_indices.into_iter().map(PyTensor)); + let swa_indices = PyList::new_bound(py, swa_indices.into_iter().map(PyTensor)); + Ok((tag, full_indices, swa_indices).into_py(py)) + } + CacheAction::RecoverSwaWithLockedFull { + node_id, + kept_full, + incoming_full, + } => Ok((tag, node_id, PyTensor(kept_full), PyTensor(incoming_full)).into_py(py)), + CacheAction::SwaRebuild { + node_id, + source_value, + } => Ok((tag, node_id, PyTensor(source_value)).into_py(py)), + } +} + +/// Convert cache actions into a Python list of tagged tuples. +fn cache_actions_to_py(py: Python<'_>, actions: Vec) -> PyResult> { + let list = PyList::empty_bound(py); + for action in actions { + list.append(cache_action_to_py(py, action)?)?; + } + Ok(list.unbind()) +} + +/// The python PoolName string for a pool. +fn pool_name_str(name: PoolName) -> &'static str { + match name { + PoolName::Kv => "kv", + PoolName::Mamba => "mamba", + PoolName::Swa => "swa", + PoolName::Indexer => "indexer", + PoolName::DeepseekV4C4 => "deepseek_v4_c4", + PoolName::DeepseekV4C4Indexer => "deepseek_v4_c4_indexer", + PoolName::DeepseekV4C128 => "deepseek_v4_c128", + PoolName::DeepseekV4C4State => "deepseek_v4_c4_state", + PoolName::DeepseekV4C4IndexerState => "deepseek_v4_c4_indexer_state", + PoolName::DeepseekV4C128State => "deepseek_v4_c128_state", + PoolName::Draft => "draft", + PoolName::DraftIndexer => "draft_indexer", + PoolName::DraftSwa => "draft_swa", + } +} + +/// Map a python PoolName string onto the Rust enum. +fn parse_pool_name(name: &str) -> PyResult { + match name { + "kv" => Ok(PoolName::Kv), + "mamba" => Ok(PoolName::Mamba), + "swa" => Ok(PoolName::Swa), + "indexer" => Ok(PoolName::Indexer), + "deepseek_v4_c4" => Ok(PoolName::DeepseekV4C4), + "deepseek_v4_c4_indexer" => Ok(PoolName::DeepseekV4C4Indexer), + "deepseek_v4_c128" => Ok(PoolName::DeepseekV4C128), + "deepseek_v4_c4_state" => Ok(PoolName::DeepseekV4C4State), + "deepseek_v4_c4_indexer_state" => Ok(PoolName::DeepseekV4C4IndexerState), + "deepseek_v4_c128_state" => Ok(PoolName::DeepseekV4C128State), + "draft" => Ok(PoolName::Draft), + "draft_indexer" => Ok(PoolName::DraftIndexer), + "draft_swa" => Ok(PoolName::DraftSwa), + other => Err(PyValueError::new_err(format!("unknown pool name: {other}"))), + } +} + +/// A pool transfer's boundary form: +/// (name, host_indices, device_indices, nodes_to_load, keys, hit_policy). +type TransferArgs = ( + String, + Option, + Option, + Option>, + Option>, + String, +); + +/// Strongly typed, attribute-based input view of a Python MatchResult. +/// Cache actions are intentionally omitted: component finalizers only update +/// match metadata, while the test adapter preserves the original actions. +#[cfg(feature = "inspection")] +#[derive(FromPyObject)] +struct InspectionMatchResultInput { + #[pyo3(attribute)] + device_indices: PyTensor, + #[pyo3(attribute)] + last_device_node: NodeId, + #[pyo3(attribute)] + last_host_node: NodeId, + #[pyo3(attribute)] + best_match_node: NodeId, + #[pyo3(attribute)] + host_hit_length: usize, + #[pyo3(attribute)] + swa_host_hit_length: usize, + #[pyo3(attribute)] + mamba_host_hit_length: usize, + #[pyo3(attribute)] + mamba_branching_seqlen: Option, + #[pyo3(attribute)] + full_kv_hit_length: usize, +} + +/// Map a python CacheTransferPhase value onto the Rust enum. +fn parse_transfer_phase(phase: &str) -> PyResult { + match phase { + "backup_host" => Ok(CacheTransferPhase::BackupHost), + "load_back" => Ok(CacheTransferPhase::LoadBack), + "backup_storage" => Ok(CacheTransferPhase::BackupStorage), + "prefetch" => Ok(CacheTransferPhase::Prefetch), + other => Err(PyValueError::new_err(format!( + "unknown transfer phase: {other}" + ))), + } +} + +/// Map a python PoolHitPolicy value onto the Rust enum. +fn parse_hit_policy(hit_policy: &str) -> PyResult { + match hit_policy { + "all_pages" => Ok(PoolHitPolicy::AllPages), + "trailing_pages" => Ok(PoolHitPolicy::TrailingPages), + other => Err(PyValueError::new_err(format!( + "unknown hit policy: {other}" + ))), + } +} + +/// Convert a pool transfer into its boundary tuple. +fn transfer_to_py(py: Python<'_>, transfer: PoolTransfer) -> PyResult> { + Ok(( + pool_name_str(transfer.name), + transfer.host_indices.map(PyTensor), + transfer.device_indices.map(PyTensor), + transfer.nodes_to_load, + transfer.keys, + transfer.hit_policy.as_str(), + ) + .into_py(py)) +} + +/// Build a pool transfer from its boundary tuple. +fn transfer_from_args(args: TransferArgs) -> PyResult { + let (name, host_indices, device_indices, nodes_to_load, keys, hit_policy) = args; + Ok(PoolTransfer { + name: parse_pool_name(&name)?, + host_indices: host_indices.map(|t| t.0), + device_indices: device_indices.map(|t| t.0), + keys, + hit_policy: parse_hit_policy(&hit_policy)?, + nodes_to_load, + }) +} + +/// Convert per-component transfers into a Python dict of boundary tuples. +fn comp_xfers_to_py( + py: Python<'_>, + comp_xfers: HashMap>, +) -> PyResult> { + let dict = PyDict::new_bound(py); + for (ct, transfers) in comp_xfers { + let list = PyList::empty_bound(py); + for transfer in transfers { + list.append(transfer_to_py(py, transfer)?)?; + } + dict.set_item(component_type_to_u8(ct), list)?; + } + Ok(dict.unbind()) +} + +/// Build per-component transfers from a Python dict of boundary tuples. +fn comp_xfers_from_args( + comp_xfers: HashMap>, +) -> PyResult>> { + comp_xfers + .into_iter() + .map(|(ct, transfers)| { + Ok(( + parse_component_type(ct)?, + transfers + .into_iter() + .map(transfer_from_args) + .collect::>>()?, + )) + }) + .collect() +} + +/// Convert per-component freed tensors into a Python dict keyed by component value. +fn frees_to_py(py: Python<'_>, frees: HashMap>) -> PyResult> { + let frees: HashMap> = frees + .into_iter() + .map(|(ct, tensors)| { + ( + component_type_to_u8(ct), + tensors.into_iter().map(PyTensor).collect(), + ) + }) + .collect(); + let dict = PyDict::new_bound(py); + for (ct, tensors) in frees { + dict.set_item(ct, tensors)?; + } + Ok(dict.unbind()) +} + +/// Python-visible tree-core init params; converts into CacheInitParams. +#[pyclass(get_all, set_all)] +#[derive(Clone)] +pub struct TreeCoreInitParamsBinding { + pub eviction_policy: String, + pub page_size: usize, + pub is_write_back: bool, + pub enable_hicache: bool, + pub write_through_threshold: i64, + pub device: String, + pub swa_sliding_window_size: Option, + pub enable_kv_cache_events: bool, + pub mamba_cache_chunk_size: Option, + pub mamba_max_states_per_path: Option, +} + +impl TreeCoreInitParamsBinding { + /// Convert into the tree core's construction params. + fn to_cache_init_params(&self) -> PyResult { + Ok(CacheInitParams { + eviction_policy: self.eviction_policy.clone(), + page_size: self.page_size, + is_write_back: self.is_write_back, + enable_hicache: self.enable_hicache, + write_through_threshold: self.write_through_threshold, + device: parse_device(&self.device)?, + swa_sliding_window_size: self.swa_sliding_window_size, + // Wired post-construction via set_has_swa_host_pool. + has_swa_host_pool: false, + enable_kv_cache_events: self.enable_kv_cache_events, + mamba_cache_chunk_size: self.mamba_cache_chunk_size, + mamba_max_states_per_path: self.mamba_max_states_per_path, + }) + } +} + +#[pymethods] +impl TreeCoreInitParamsBinding { + #[new] + #[pyo3(signature = (eviction_policy = "lru".to_string(), page_size = 1, is_write_back = false, enable_hicache = false, write_through_threshold = 256, device = "cpu".to_string(), swa_sliding_window_size = None, enable_kv_cache_events = false, mamba_cache_chunk_size = None, mamba_max_states_per_path = None))] + fn new( + eviction_policy: String, + page_size: usize, + is_write_back: bool, + enable_hicache: bool, + write_through_threshold: i64, + device: String, + swa_sliding_window_size: Option, + enable_kv_cache_events: bool, + mamba_cache_chunk_size: Option, + mamba_max_states_per_path: Option, + ) -> Self { + TreeCoreInitParamsBinding { + eviction_policy, + page_size, + is_write_back, + enable_hicache, + write_through_threshold, + device, + swa_sliding_window_size, + enable_kv_cache_events, + mamba_cache_chunk_size, + mamba_max_states_per_path, + } + } +} + +/// Python-visible match params; converts into MatchPrefixParams. +#[pyclass(get_all, set_all)] +#[derive(Clone)] +pub struct MatchParamsBinding { + pub key: Vec, + pub extra_key: Option, + pub cache_salt: Option, +} + +#[pymethods] +impl MatchParamsBinding { + #[new] + #[pyo3(signature = (key, extra_key = None, cache_salt = None))] + fn new( + py: Python<'_>, + key: &Bound<'_, PyAny>, + extra_key: Option, + cache_salt: Option, + ) -> PyResult { + Ok(MatchParamsBinding { + key: py_array_to_vec_i64(py, key)?, + extra_key, + cache_salt, + }) + } +} + +/// Python-visible insert params; converts into InsertParams. The value tensor +/// stays a Python-held reference until the insert call unwraps it. +#[pyclass(get_all, set_all)] +pub struct InsertParamsBinding { + pub key: Vec, + pub value: Py, + pub extra_key: Option, + pub cache_salt: Option, + pub mamba_value: Option>, + pub prev_prefix_len: usize, + pub swa_evicted_seqlen: usize, + pub chunked: bool, + pub priority: i64, + pub track_adopted_ranges: bool, +} + +#[pymethods] +impl InsertParamsBinding { + #[new] + #[pyo3(signature = (key, value, extra_key = None, cache_salt = None, prev_prefix_len = 0, swa_evicted_seqlen = 0, chunked = false, priority = 0, mamba_value = None, track_adopted_ranges = false))] + fn new( + py: Python<'_>, + key: &Bound<'_, PyAny>, + value: Py, + extra_key: Option, + cache_salt: Option, + prev_prefix_len: usize, + swa_evicted_seqlen: usize, + chunked: bool, + priority: i64, + mamba_value: Option>, + track_adopted_ranges: bool, + ) -> PyResult { + Ok(InsertParamsBinding { + key: py_array_to_vec_i64(py, key)?, + value, + extra_key, + cache_salt, + mamba_value, + prev_prefix_len, + swa_evicted_seqlen, + chunked, + priority, + track_adopted_ranges, + }) + } +} + +/// Python-visible match result; tensors and actions are Python-held. +#[pyclass(get_all)] +pub struct MatchResultBinding { + device_indices: Py, + last_device_node_id: NodeId, + last_host_node_id: NodeId, + best_match_node_id: NodeId, + host_hit_length: usize, + swa_host_hit_length: usize, + mamba_host_hit_length: usize, + mamba_branching_seqlen: Option, + full_kv_hit_length: usize, + cache_actions: Py, +} + +impl MatchResultBinding { + /// Move a core match result across the boundary. + fn from_match_result(py: Python<'_>, result: MatchResult) -> PyResult { + Ok(MatchResultBinding { + device_indices: tensor_to_py(py, result.device_indices)?, + last_device_node_id: result.last_device_node_id, + last_host_node_id: result.last_host_node_id, + best_match_node_id: result.best_match_node_id, + host_hit_length: result.host_hit_length, + swa_host_hit_length: result.swa_host_hit_length, + mamba_host_hit_length: result.mamba_host_hit_length, + mamba_branching_seqlen: result.mamba_branching_seqlen, + full_kv_hit_length: result.full_kv_hit_length, + cache_actions: cache_actions_to_py(py, result.cache_actions)?, + }) + } +} + +/// Python-visible insert result; actions are Python-held. +#[pyclass(get_all)] +pub struct InsertResultBinding { + prefix_len: usize, + total_len: usize, + last_device_node: Option, + inserted_host_node: Option, + host_insert_dropped: bool, + mamba_exist: bool, + adopted_ranges: Option>>, + cache_actions: Py, +} + +/// One step of the resumable insert: the actions to apply at this barrier +/// and the final result once the walk completes. +#[pyclass(get_all)] +pub struct InsertStepResultBinding { + actions: Py, + result: Option>, +} + +impl InsertStepResultBinding { + /// Move a core insert step across the boundary. + fn from_insert_step(py: Python<'_>, step: InsertStepResult) -> PyResult { + let result = match step.result { + Some(result) => Some(Py::new( + py, + InsertResultBinding::from_insert_result(py, result)?, + )?), + None => None, + }; + Ok(InsertStepResultBinding { + actions: cache_actions_to_py(py, step.actions)?, + result, + }) + } +} + +impl InsertResultBinding { + /// Move a core insert result across the boundary. + fn from_insert_result(py: Python<'_>, result: InsertResult) -> PyResult { + Ok(InsertResultBinding { + prefix_len: result.prefix_len, + total_len: result.total_len, + last_device_node: result.last_device_node_id, + inserted_host_node: result.inserted_host_node, + host_insert_dropped: result.host_insert_dropped, + mamba_exist: result.mamba_exist, + adopted_ranges: result.adopted_ranges.map(|ranges| { + ranges + .into_iter() + .map(|(component_type, ranges)| (component_type_to_u8(component_type), ranges)) + .collect() + }), + cache_actions: cache_actions_to_py(py, result.cache_actions)?, + }) + } +} + +/// Python-visible dec-lock params; converts into DecLockRefParams. +#[pyclass(get_all, set_all)] +#[derive(Clone, Default)] +pub struct DecLockRefParamsBinding { + pub swa_uuid_for_lock: Option, + pub swa_uuid_for_host_lock: Option, + pub skip_lock_node_ids: HashMap>, +} + +#[pymethods] +impl DecLockRefParamsBinding { + #[new] + #[pyo3(signature = (swa_uuid_for_lock = None, swa_uuid_for_host_lock = None, skip_lock_node_ids = None))] + fn new( + swa_uuid_for_lock: Option, + swa_uuid_for_host_lock: Option, + skip_lock_node_ids: Option>>, + ) -> Self { + DecLockRefParamsBinding { + swa_uuid_for_lock, + swa_uuid_for_host_lock, + skip_lock_node_ids: skip_lock_node_ids.unwrap_or_default(), + } + } +} + +impl DecLockRefParamsBinding { + /// Convert into the tree core's dec-lock params. + fn to_dec_lock_ref_params(&self) -> PyResult { + Ok(DecLockRefParams { + swa_uuid_for_lock: self.swa_uuid_for_lock, + swa_uuid_for_host_lock: self.swa_uuid_for_host_lock, + skip_lock_node_ids: self + .skip_lock_node_ids + .iter() + .map(|(ct, node_ids)| { + Ok::<_, PyErr>((parse_component_type(*ct)?, node_ids.clone())) + }) + .collect::>()?, + }) + } +} + +/// Python-visible inc_lock_ref result; hand skip_lock_node_ids back to the +/// matching dec_lock_ref. +#[pyclass(get_all)] +pub struct IncLockRefResultBinding { + delta: Option, + swa_uuid_for_lock: Option, + swa_uuid_for_host_lock: Option, + skip_lock_node_ids: HashMap>, +} + +impl IncLockRefResultBinding { + fn from_result(result: crate::unified_tree_core::IncLockRefResult) -> Self { + Self { + delta: result.delta, + swa_uuid_for_lock: result.swa_uuid_for_lock, + swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, + skip_lock_node_ids: result + .skip_lock_node_ids + .into_iter() + .map(|(ct, node_ids)| (component_type_to_u8(ct), node_ids)) + .collect(), + } + } +} + +/// Convert a Python component-keyed tracker into the core's counts. +fn tracker_from_py(tracker: HashMap) -> PyResult> { + tracker + .into_iter() + .map(|(ct, freed)| Ok((parse_component_type(ct)?, freed))) + .collect() +} + +/// Convert the core's tracker back into Python component-keyed counts. +fn tracker_to_py(tracker: HashMap) -> HashMap { + tracker + .into_iter() + .map(|(ct, freed)| (component_type_to_u8(ct), freed)) + .collect() +} + +/// Next-eviction-node step result: the node to evict, whether the walk made +/// progress, this step's per-component evicted counts, and this step's newly +/// freed tensors. +#[pyclass(get_all)] +pub struct EvictDeviceNextNodeResultBinding { + node_id: Option, + made_progress: bool, + tracker: HashMap, + new_device_frees: Py, + new_host_frees: Py, +} + +/// Leaf-eviction step result: the backup action for an unbacked +/// write-back leaf (else None), this step's per-component evicted counts, +/// and this step's newly freed tensors. +#[pyclass(get_all)] +pub struct EvictDeviceLeafResultBinding { + backup_kv: Option>, + tracker: HashMap, + new_device_frees: Py, + new_host_frees: Py, +} + +/// Drop-subtree result: whether the drop happened, this step's +/// per-component evicted counts, and the subtree's newly freed tensors. +#[pyclass(get_all)] +pub struct DropSubtreeResultBinding { + dropped: bool, + tracker: HashMap, + new_device_frees: Py, + new_host_frees: Py, +} + +/// Demote result: this step's per-component evicted counts and the +/// demoted node's newly freed tensors. +#[pyclass(get_all)] +pub struct DemoteResultBinding { + tracker: HashMap, + new_device_frees: Py, + new_host_frees: Py, +} + +/// Python-visible device->storage backup spec. +#[pyclass(get_all)] +pub struct StorageBackupSpecBinding { + host_value: Py, + /// Raw native-endian int64 bytes; python rebuilds them via array("q").frombytes. + token_ids: Py, + hash_value: Option>, + prefix_keys: Option>, + comp_xfers: Py, +} + +#[pyclass(get_all)] +pub struct BufferBackupSnapshotBinding { + node_id: NodeId, + parent_node_id: NodeId, + parent_is_root: bool, + parent_last_hash: Option, + key_token_ids: Py, + extra_key: Option, + cache_salt: Option, + is_bigram: bool, + hash_values: Vec, + prefix_keys: Option>, +} + +impl BufferBackupSnapshotBinding { + fn from_snapshot(py: Python<'_>, snapshot: BufferBackupSnapshot) -> Self { + let mut token_bytes = Vec::with_capacity(snapshot.token_ids.len() * 8); + for token in snapshot.token_ids { + token_bytes.extend_from_slice(&token.to_ne_bytes()); + } + Self { + node_id: snapshot.node_id, + parent_node_id: snapshot.parent_node_id, + parent_is_root: snapshot.parent_is_root, + parent_last_hash: snapshot.parent_last_hash, + key_token_ids: PyBytes::new_bound(py, &token_bytes).unbind(), + extra_key: snapshot.extra_key, + cache_salt: snapshot.cache_salt, + is_bigram: snapshot.is_bigram, + hash_values: snapshot.hash_values, + prefix_keys: snapshot.prefix_keys, + } + } +} + +#[pyclass(get_all)] +pub struct BufferBackupStateBinding { + parent_node_id: NodeId, + parent_is_root: bool, + parent_last_hash: Option, +} + +impl From for BufferBackupStateBinding { + fn from(state: BufferBackupState) -> Self { + Self { + parent_node_id: state.parent_node_id, + parent_is_root: state.parent_is_root, + parent_last_hash: state.parent_last_hash, + } + } +} + +/// Python-visible KV-canary walk rows: parallel int64 tensors of slots, +/// positions, and preceding slots. +#[pyclass(get_all)] +pub struct KvCanaryWalkResultBinding { + slot_indices: Py, + positions: Py, + prev_slot_indices: Py, +} + +/// Host-eviction drive result: this step's per-component evicted counts +/// and the drive's newly freed tensors. +#[pyclass(get_all)] +pub struct HostEvictionResultBinding { + tracker: HashMap, + new_device_frees: Py, + new_host_frees: Py, +} + +impl HostEvictionResultBinding { + fn from_eviction_step(py: Python<'_>, result: EvictionStepResult) -> PyResult { + Ok(Self { + tracker: tracker_to_py(result.tracker), + new_device_frees: frees_to_py(py, result.device_frees)?, + new_host_frees: frees_to_py(py, result.host_frees)?, + }) + } +} + +/// The generic UnifiedTreeCore adapter the per-key-type pyclasses delegate to; +/// the Mutex makes the Send-only core satisfy pyclass's Sync bound. +struct TreeCoreBinding { + core: Mutex>, + /// The core's construction device, kept outside the Mutex for pre-lock validation. + device: Device, + /// The core's page size, kept outside the Mutex for pre-lock validation. + page_size: usize, +} + +// Send + Sync lets allow_threads release the GIL around core calls. +impl TreeCoreBinding { + /// Build a tree core for the given component types from the cache's + /// init params. + fn new(init_params: &TreeCoreInitParamsBinding, component_types: Vec) -> PyResult { + let component_types = component_types + .into_iter() + .map(parse_component_type) + .collect::>>()?; + if component_types != [FULL] + && component_types != [FULL, SWA] + && component_types != [FULL, MAMBA] + && component_types != [FULL, SWA, MAMBA] + { + return Err(PyValueError::new_err( + "only the [Full], [Full, Swa], [Full, Mamba], and [Full, Swa, Mamba] component sets are supported", + )); + } + if component_types.contains(&SWA) && init_params.swa_sliding_window_size.is_none() { + return Err(PyValueError::new_err( + "the Swa component requires swa_sliding_window_size", + )); + } + if component_types.contains(&MAMBA) && init_params.mamba_cache_chunk_size.is_none() { + return Err(PyValueError::new_err( + "the Mamba component requires mamba_cache_chunk_size", + )); + } + if init_params.page_size == 0 { + return Err(PyValueError::new_err("page_size must be at least 1")); + } + let eviction_policy = init_params.eviction_policy.to_lowercase(); + if !matches!( + eviction_policy.as_str(), + "lru" | "lfu" | "fifo" | "mru" | "filo" | "priority" | "slru" + ) { + return Err(PyValueError::new_err(format!( + "Unknown eviction policy: {eviction_policy}. Supported policies: \ + 'lru', 'lfu', 'fifo', 'mru', 'filo', 'priority', 'slru'." + ))); + } + let params = init_params.to_cache_init_params()?; + let device = params.device; + let page_size = params.page_size; + Ok(TreeCoreBinding { + core: Mutex::new(UnifiedTreeCore::new(params, component_types)), + device, + page_size, + }) + } + + /// Lock the core for one adapter call. A panic can leave a mutation half-applied, + /// so a poisoned core is never reused. + fn core(&self) -> std::sync::MutexGuard<'_, UnifiedTreeCore> { + self.core.lock().unwrap_or_else(|_| { + panic!("Rust TreeCore mutex poisoned; refusing to reuse state after an earlier panic") + }) + } + + /// Reject an insert value whose dtype, device, or length cannot cover the key. + fn validate_insert_value(&self, value: &Tensor, key_atom_len: usize) -> PyResult<()> { + if value.kind() != Kind::Int64 { + return Err(PyValueError::new_err(format!( + "insert value must be an int64 tensor, got {:?}", + value.kind() + ))); + } + if value.device() != self.device { + return Err(PyValueError::new_err(format!( + "insert value device {:?} does not match the tree core device {:?}", + value.device(), + self.device + ))); + } + let aligned_key_len = key_atom_len / self.page_size * self.page_size; + let value_len = value.size().first().copied().unwrap_or(0); + if value_len < aligned_key_len as i64 { + return Err(PyValueError::new_err(format!( + "insert value length {value_len} is shorter than the aligned key length {aligned_key_len}" + ))); + } + Ok(()) + } + + /// Drop the entire tree and reinitialize empty state. + fn reset(&self, py: Python<'_>) { + py.allow_threads(|| self.core().reset()); + } + + /// Match a key against the tree. + fn match_prefix( + &self, + py: Python<'_>, + params: &MatchParamsBinding, + ) -> PyResult { + let key = K::key_from(Cow::Borrowed(¶ms.key)); + let key = key.as_ref(); + let params = MatchPrefixParams { + key, + namespace: KeyNamespaceRef::new( + params.extra_key.as_deref(), + params.cache_salt.as_deref(), + ), + }; + let result = py.allow_threads(|| self.core().match_prefix(¶ms)); + MatchResultBinding::from_match_result(py, result) + } + + /// The empty match result anchored at the root. + fn empty_match_result(&self, py: Python<'_>) -> PyResult { + let result = py.allow_threads(|| self.core().empty_match_result()); + MatchResultBinding::from_match_result(py, result) + } + + /// Insert device values into the tree per the provided key. + fn insert( + &self, + py: Python<'_>, + params: &InsertParamsBinding, + ) -> PyResult { + let key = K::key_from(Cow::Borrowed(¶ms.key)); + let key = key.as_ref(); + let value: PyTensor = params.value.bind(py).extract()?; + // The value covers key atoms (bigram: raw len - 1), so validate the converted key. + self.validate_insert_value(&value.0, key.atom_len())?; + let mamba_value = match ¶ms.mamba_value { + Some(mamba_value) => Some(mamba_value.bind(py).extract::()?.0), + None => None, + }; + let params = InsertParams { + key, + namespace: KeyNamespaceRef::new( + params.extra_key.as_deref(), + params.cache_salt.as_deref(), + ), + value: value.0, + mamba_value, + prev_prefix_len: params.prev_prefix_len, + swa_evicted_seqlen: params.swa_evicted_seqlen, + chunked: params.chunked, + priority: params.priority, + track_adopted_ranges: params.track_adopted_ranges, + }; + let result = py + .allow_threads(move || self.core().try_insert(¶ms)) + .map_err(tree_core_runtime_error)?; + InsertResultBinding::from_insert_result(py, result) + } + + /// Start the resumable insert, running to its first barrier or completion. + fn begin_insert( + &self, + py: Python<'_>, + params: &InsertParamsBinding, + ) -> PyResult { + let key = K::key_from(Cow::Borrowed(¶ms.key)); + let key = key.as_ref(); + let value: PyTensor = params.value.bind(py).extract()?; + // The value covers key atoms (bigram: raw len - 1), so validate the converted key. + self.validate_insert_value(&value.0, key.atom_len())?; + let mamba_value = match ¶ms.mamba_value { + Some(mamba_value) => Some(mamba_value.bind(py).extract::()?.0), + None => None, + }; + let params = InsertParams { + key, + namespace: KeyNamespaceRef::new( + params.extra_key.as_deref(), + params.cache_salt.as_deref(), + ), + value: value.0, + mamba_value, + prev_prefix_len: params.prev_prefix_len, + swa_evicted_seqlen: params.swa_evicted_seqlen, + chunked: params.chunked, + priority: params.priority, + track_adopted_ranges: params.track_adopted_ranges, + }; + let step = py + .allow_threads(move || self.core().try_begin_insert(¶ms)) + .map_err(tree_core_runtime_error)?; + InsertStepResultBinding::from_insert_step(py, step) + } + + /// Continue the suspended insert after its step actions were applied. + fn resume_insert(&self, py: Python<'_>) -> PyResult { + let step = py + .allow_threads(|| self.core().try_resume_insert()) + .map_err(tree_core_runtime_error)?; + InsertStepResultBinding::from_insert_step(py, step) + } + + /// Whether an insert walk is suspended at a barrier. + fn has_ongoing_insert(&self, py: Python<'_>) -> bool { + py.allow_threads(|| self.core().has_ongoing_insert()) + } + + /// Finish the insert (idempotent); returns still-pending actions to drain. + fn end_insert(&self, py: Python<'_>) -> PyResult> { + let actions = py.allow_threads(|| self.core().end_insert()); + cache_actions_to_py(py, actions) + } + + /// Bump the reference count on a node's component locks. + fn inc_lock_ref( + &self, + py: Python<'_>, + node_id: NodeId, + skip_lock_components: Option>, + ) -> PyResult { + let skip_lock_components = skip_lock_components + .unwrap_or_default() + .into_iter() + .map(parse_component_type) + .collect::>>()?; + let result = py.allow_threads(|| { + self.core() + .inc_lock_ref_with_skip(node_id, &skip_lock_components) + }); + Ok(IncLockRefResultBinding::from_result(result)) + } + + /// Decrease the reference count on a node's component locks. + fn dec_lock_ref( + &self, + py: Python<'_>, + node_id: NodeId, + params: Option<&DecLockRefParamsBinding>, + skip_swa: bool, + ) -> PyResult<()> { + let params = params.map(|p| p.to_dec_lock_ref_params()).transpose()?; + py.allow_threads(|| self.core().dec_lock_ref(node_id, params.as_ref(), skip_swa)); + Ok(()) + } + + /// Early-release the SWA portion of a request's tree lock; returns this + /// release's per-component (device_frees, host_frees). + fn dec_swa_lock_only( + &self, + py: Python<'_>, + node_id: NodeId, + swa_uuid_for_lock: Option, + skip_lock_node_ids: Option>>, + ) -> PyResult<(Py, Py)> { + let skip_lock_node_ids = skip_lock_node_ids + .unwrap_or_default() + .into_iter() + .map(|(ct, node_ids)| Ok((parse_component_type(ct)?, node_ids))) + .collect::>>()?; + let (device_frees, host_frees) = py.allow_threads(|| { + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + self.core().dec_swa_lock_only_with_skip( + node_id, + swa_uuid_for_lock, + Some(&skip_lock_node_ids), + &mut device_frees, + &mut host_frees, + ); + (device_frees, host_frees) + }); + Ok((frees_to_py(py, device_frees)?, frees_to_py(py, host_frees)?)) + } + + /// Store a component's device value on a node (the SWA rebuild write-back). + fn set_component_device_value( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + value: PyTensor, + ) -> PyResult<()> { + let component_type = parse_component_type(component_type)?; + let value = value.0; + if value.kind() != Kind::Int64 { + return Err(PyValueError::new_err(format!( + "component device value must be an int64 tensor, got {:?}", + value.kind() + ))); + } + if value.device() != self.device { + return Err(PyValueError::new_err(format!( + "component device value device {:?} does not match the tree core device {:?}", + value.device(), + self.device + ))); + } + py.allow_threads(|| { + self.core() + .set_component_device_value(node_id, component_type, value) + }); + Ok(()) + } + + /// A component's device value on a node, if set. + fn get_component_device_value( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + ) -> PyResult> { + let component_type = parse_component_type(component_type)?; + let value = py.allow_threads(|| { + self.core() + .get_component_device_value(node_id, component_type) + .map(|tensor| tensor.shallow_clone()) + }); + Ok(value.map(PyTensor)) + } + + // TODO(jialino): batch a full no-backup eviction round in Rust (one crossing + // instead of 2N+2); step-wise stays for the write-back interleave. + /// Begin a component's device-eviction walk for up to request_cnt tokens. + fn evict_device_start( + &self, + py: Python<'_>, + component_type: u8, + request_cnt: usize, + ) -> PyResult<()> { + let ct = parse_component_type(component_type)?; + py.allow_threads(|| self.core().evict_device_start(ct, request_cnt)); + Ok(()) + } + + /// Advance one component eviction step. A missing node with + /// `made_progress` set means an internal tombstone completed the step; + /// otherwise a missing node means the walk is exhausted. + fn evict_device_next_node( + &self, + py: Python<'_>, + component_type: u8, + tracker: HashMap, + ) -> PyResult { + let ct = parse_component_type(component_type)?; + let baseline = tracker_from_py(tracker)?; + let (node_id, result) = + py.allow_threads(move || self.core().evict_device_next_node(ct, &baseline)); + let made_progress = node_id.is_some() || !result.tracker.is_empty(); + Ok(EvictDeviceNextNodeResultBinding { + node_id, + made_progress, + tracker: tracker_to_py(result.tracker), + new_device_frees: frees_to_py(py, result.device_frees)?, + new_host_frees: frees_to_py(py, result.host_frees)?, + }) + } + + /// Evict one device leaf; an unbacked write-back leaf returns its backup + /// action for the caller to execute before demoting. + fn evict_device_leaf( + &self, + py: Python<'_>, + node_id: NodeId, + ) -> PyResult { + let (backup, result) = py.allow_threads(move || { + let mut core = self.core(); + let is_write_back = core.is_write_back; + core.evict_device_leaf(node_id, is_write_back) + }); + Ok(EvictDeviceLeafResultBinding { + backup_kv: backup + .map(|backup| cache_action_to_py(py, CacheAction::BackupKV(backup))) + .transpose()?, + tracker: tracker_to_py(result.tracker), + new_device_frees: frees_to_py(py, result.device_frees)?, + new_host_frees: frees_to_py(py, result.host_frees)?, + }) + } + + /// Finish a component's device-eviction walk. + fn evict_device_end(&self, py: Python<'_>, component_type: u8) -> PyResult<()> { + let ct = parse_component_type(component_type)?; + py.allow_threads(|| self.core().evict_device_end(ct)); + Ok(()) + } + + /// Verify tree-structure, leaf-set, LRU, size, and ongoing-op invariants; + /// ongoing_* args are (id, node_id) pairs. + fn sanity_check( + &self, + py: Python<'_>, + ongoing_write_through: Vec<(i64, NodeId)>, + ongoing_load_back: Vec<(i64, NodeId)>, + ) -> PyResult<()> { + py.allow_threads(|| { + self.core() + .try_sanity_check(&ongoing_write_through, &ongoing_load_back) + }) + .map_err(PyAssertionError::new_err) + } + + /// Concatenated FULL device values from from_node up to (exclusive) until_node. + fn collect_full_device_indices( + &self, + py: Python<'_>, + from_node_id: NodeId, + until_node_id: NodeId, + ) -> PyTensor { + PyTensor(py.allow_threads(|| { + self.core() + .collect_full_device_indices(from_node_id, until_node_id) + })) + } + + /// Every FULL device value in the tree, concatenated. + fn all_values_flatten(&self, py: Python<'_>) -> PyTensor { + PyTensor(py.allow_threads(|| self.core().all_values_flatten())) + } + + /// Every Mamba device value in the tree, concatenated. + fn all_mamba_values_flatten(&self, py: Python<'_>) -> PyTensor { + PyTensor(py.allow_threads(|| self.core().all_mamba_values_flatten())) + } + + /// Flatten every FULL device slot into (slot, position, prev-slot) rows for the KV-canary sweep. + fn walk_for_kv_canary( + &self, + py: Python<'_>, + unlocked_only: bool, + swa_resident_only: bool, + ) -> PyResult { + let result = py.allow_threads(|| { + self.core() + .walk_for_kv_canary(unlocked_only, swa_resident_only) + }); + Ok(KvCanaryWalkResultBinding { + slot_indices: tensor_to_py(py, Tensor::from_slice(&result.slot_indices))?, + positions: tensor_to_py(py, Tensor::from_slice(&result.positions))?, + prev_slot_indices: tensor_to_py(py, Tensor::from_slice(&result.prev_slot_indices))?, + }) + } + + /// Evictable token count of the FULL (base) component. + fn evictable_size(&self, py: Python<'_>) -> usize { + py.allow_threads(|| self.core().evictable_size()) + } + + /// Protected (locked) token count of the FULL (base) component. + fn protected_size(&self, py: Python<'_>) -> usize { + py.allow_threads(|| self.core().protected_size()) + } + + /// FULL component evictable token count. + fn full_evictable_size(&self, py: Python<'_>) -> usize { + py.allow_threads(|| self.core().full_evictable_size()) + } + + /// FULL component protected token count. + fn full_protected_size(&self, py: Python<'_>) -> usize { + py.allow_threads(|| self.core().full_protected_size()) + } + + /// Evictable token count for one component (0 if the component is absent). + fn component_evictable_size(&self, py: Python<'_>, component_type: u8) -> PyResult { + let ct = parse_component_type(component_type)?; + Ok(py.allow_threads(|| self.core().component_evictable_size(ct))) + } + + /// Protected token count for one component (0 if the component is absent). + fn component_protected_size(&self, py: Python<'_>, component_type: u8) -> PyResult { + let ct = parse_component_type(component_type)?; + Ok(py.allow_threads(|| self.core().component_protected_size(ct))) + } + + /// (full_tokens, aux_tokens) summed across the whole tree. + fn total_size(&self, py: Python<'_>) -> (usize, usize) { + py.allow_threads(|| self.core().total_size()) + } + + /// Whether the node's FULL device value has been evicted. + fn is_full_device_evicted(&self, py: Python<'_>, node_id: NodeId) -> bool { + py.allow_threads(|| self.core().is_full_device_evicted(node_id)) + } + + /// Mark the host tier (HiCache) as wired. + fn set_hicache_enabled(&self, py: Python<'_>) { + py.allow_threads(|| self.core().set_hicache_enabled()); + } + + /// Whether the host tier (HiCache) is wired. + fn enable_hicache(&self, py: Python<'_>) -> bool { + py.allow_threads(|| self.core().enable_hicache) + } + + /// Mark the SWA host pool as wired (HiCache). + fn set_has_swa_host_pool(&self, py: Python<'_>) { + py.allow_threads(|| self.core().set_has_swa_host_pool()); + } + + /// Whether the SWA host pool is wired. + fn has_swa_host_pool(&self, py: Python<'_>) -> bool { + py.allow_threads(|| self.core().has_swa_host_pool) + } + + /// Insert a host-side (backuped) tree path descending from the given node. + fn insert_host( + &self, + py: Python<'_>, + node_id: NodeId, + extra_key: Option, + key: &Bound<'_, PyAny>, + host_value: PyTensor, + hash_value: Vec, + cache_salt: Option, + ) -> PyResult { + let key = K::key_from(Cow::Owned(py_array_to_vec_i64(py, key)?)).into_owned(); + let host_value = host_value.0; + if host_value.kind() != Kind::Int64 { + return Err(PyValueError::new_err(format!( + "insert_host host_value must be an int64 tensor, got {:?}", + host_value.kind() + ))); + } + let result = py + .allow_threads(move || { + self.core().try_insert_host_in_namespace( + node_id, + KeyNamespaceRef::new(extra_key.as_deref(), cache_salt.as_deref()), + key, + host_value, + hash_value, + ) + }) + .map_err(tree_core_runtime_error)?; + InsertResultBinding::from_insert_result(py, result) + } + + /// Gather a node's device value plus per-component BACKUP_HOST transfers. + fn build_backup_spec( + &self, + py: Python<'_>, + node_id: NodeId, + ) -> PyResult<(PyTensor, Py)> { + let (device_value, comp_xfers) = + py.allow_threads(|| self.core().build_backup_spec(node_id)); + Ok((PyTensor(device_value), comp_xfers_to_py(py, comp_xfers)?)) + } + + /// Gather a node's device->storage backup spec; None if the node is not backuped. + fn build_storage_backup_spec( + &self, + py: Python<'_>, + node_id: NodeId, + pass_prefix_keys: bool, + ) -> PyResult> { + let spec = py.allow_threads(|| { + self.core() + .build_storage_backup_spec(node_id, pass_prefix_keys) + }); + let Some(spec) = spec else { + return Ok(None); + }; + let mut token_bytes = Vec::with_capacity(spec.token_ids.len() * 8); + for token in &spec.token_ids { + token_bytes.extend_from_slice(&token.to_ne_bytes()); + } + Ok(Some(StorageBackupSpecBinding { + host_value: tensor_to_py(py, spec.host_value)?, + token_ids: PyBytes::new_bound(py, &token_bytes).unbind(), + hash_value: spec.hash_value, + prefix_keys: spec.prefix_keys, + comp_xfers: comp_xfers_to_py(py, spec.comp_xfers)?, + })) + } + + /// Route a build_hicache_transfers call to the component for the given type. + #[allow(clippy::too_many_arguments)] + fn build_hicache_transfers( + &self, + py: Python<'_>, + component_type: u8, + node_id: NodeId, + phase: &str, + host_indices: Option, + token_ids: Option>, + prefetch_tokens: usize, + last_hash: Option, + ) -> PyResult>>> { + let component_type = parse_component_type(component_type)?; + let phase = parse_transfer_phase(phase)?; + let host_indices = host_indices.map(|t| t.0); + let transfers = py + .allow_threads(|| { + self.core().try_build_hicache_transfers( + component_type, + node_id, + phase, + host_indices, + token_ids.as_deref(), + prefetch_tokens, + last_hash.as_deref(), + ) + }) + .map_err(tree_core_assertion_error)?; + transfers + .map(|transfers| { + transfers + .into_iter() + .map(|transfer| transfer_to_py(py, transfer)) + .collect::>>() + }) + .transpose() + } + + /// The anchor node's namespace; None for root-like anchors. + fn prefetch_anchor_info( + &self, + py: Python<'_>, + node_id: NodeId, + ) -> PyResult<(Option, Option)> { + py.allow_threads(|| self.core().try_prefetch_anchor_info(node_id)) + .map_err(tree_core_runtime_error) + } + + /// Whether the node's Full KV is present on host. + fn node_backuped(&self, py: Python<'_>, node_id: NodeId) -> PyResult { + py.allow_threads(|| self.core().try_node_backuped(node_id)) + .map_err(tree_core_runtime_error) + } + + /// Whether the node is a (default or named) root. + fn is_root(&self, py: Python<'_>, node_id: NodeId) -> PyResult { + py.allow_threads(|| self.core().try_is_root(node_id)) + .map_err(tree_core_runtime_error) + } + + /// The node's last page hash, or None when it was never hashed. + fn get_last_hash_value(&self, py: Python<'_>, node_id: NodeId) -> PyResult> { + py.allow_threads(|| self.core().try_get_last_hash_value(node_id)) + .map_err(tree_core_runtime_error) + } + + /// The hash chain of the node's ancestors, in root-to-parent order. + fn get_prefix_hash_values(&self, py: Python<'_>, node_id: NodeId) -> PyResult> { + py.allow_threads(|| self.core().try_get_prefix_hash_values(node_id)) + .map_err(tree_core_runtime_error) + } + + fn get_hash_values(&self, py: Python<'_>, node_id: NodeId) -> PyResult> { + py.allow_threads(|| self.core().try_get_hash_values(node_id)) + .map_err(tree_core_runtime_error) + } + + fn snapshot_buffer_backup( + &self, + py: Python<'_>, + node_id: NodeId, + pass_prefix_keys: bool, + ) -> Option { + py.allow_threads(|| { + self.core() + .snapshot_buffer_backup(node_id, pass_prefix_keys) + }) + .map(|snapshot| BufferBackupSnapshotBinding::from_snapshot(py, snapshot)) + } + + fn validate_buffer_backup( + &self, + py: Python<'_>, + node_id: NodeId, + expected_key_length: usize, + ) -> Option { + py.allow_threads(|| { + self.core() + .validate_buffer_backup(node_id, expected_key_length) + }) + .map(BufferBackupStateBinding::from) + } + + /// Hash every node built while storage was disabled. + fn backfill_missing_hash_values(&self, py: Python<'_>) -> usize { + py.allow_threads(|| self.core().backfill_missing_hash_values()) + } + + fn root_node_handle(&self, py: Python<'_>, extra_key: Option) -> NodeId { + py.allow_threads(|| self.core().root_node_handle(extra_key.as_deref())) + } + + fn dfs_weight_order(&self, py: Python<'_>, node_ids: Vec) -> PyResult> { + py.allow_threads(|| self.core().try_dfs_weight_order(&node_ids)) + .map_err(tree_core_runtime_error) + } + + /// Commit each component's HiCache transfers; returns the new cache actions. + fn commit_hicache_transfers( + &self, + py: Python<'_>, + node_id: NodeId, + phase: &str, + comp_xfers: HashMap>, + insert_result: Option<(usize, Option, bool)>, + pool_storage_result: Option<(usize, HashMap)>, + ) -> PyResult<(Py, Option)> { + let phase = parse_transfer_phase(phase)?; + let comp_xfers = comp_xfers_from_args(comp_xfers)?; + let insert_result = + insert_result.map( + |(total_len, inserted_host_node, mamba_exist)| InsertResult { + total_len, + inserted_host_node, + mamba_exist, + ..InsertResult::default() + }, + ); + let pool_storage_result = pool_storage_result + .map(|(kv_hit_pages, extra_pool_hit_pages)| { + Ok::<_, PyErr>(PoolTransferResult { + kv_hit_pages, + extra_pool_hit_pages: extra_pool_hit_pages + .into_iter() + .map(|(name, pages)| Ok((parse_pool_name(&name)?, pages))) + .collect::>>()?, + }) + }) + .transpose()?; + let (cache_actions, mamba_exist) = py.allow_threads(move || { + let mut cache_actions = Vec::new(); + let mut insert_result = insert_result; + self.core().commit_hicache_transfers( + node_id, + phase, + comp_xfers, + &mut cache_actions, + insert_result.as_mut(), + pool_storage_result.as_ref(), + ); + ( + cache_actions, + insert_result.map(|result| result.mamba_exist), + ) + }); + Ok((cache_actions_to_py(py, cache_actions)?, mamba_exist)) + } + + /// Commit a successful backup to the node. + fn commit_backup( + &self, + py: Python<'_>, + node_id: NodeId, + host_indices: PyTensor, + comp_xfers: HashMap>, + ) -> PyResult<()> { + let comp_xfers = comp_xfers_from_args(comp_xfers)?; + let host_indices = host_indices.0; + py.allow_threads(move || self.core().commit_backup(node_id, host_indices, comp_xfers)); + Ok(()) + } + + /// Build the H->D load-back KV transfer plus per-component aux transfers. + fn build_load_back_spec( + &self, + py: Python<'_>, + node_id: NodeId, + mamba_pool_idx: Option, + ) -> PyResult<(Py, Py)> { + let req = Req { + mamba_pool_idx: mamba_pool_idx.map(|t| t.0), + }; + let (kv_xfer, comp_xfers) = py + .allow_threads(move || self.core().try_build_load_back_spec(node_id, Some(&req))) + .map_err(tree_core_assertion_error)?; + Ok(( + transfer_to_py(py, kv_xfer)?, + comp_xfers_to_py(py, comp_xfers)?, + )) + } + + /// Commit a successful H->D load-back onto the node; returns its actions. + fn commit_load_back( + &self, + py: Python<'_>, + node_id: NodeId, + device_indices: PyTensor, + kv_xfer: TransferArgs, + comp_xfers: HashMap>, + ) -> PyResult> { + let kv_xfer = transfer_from_args(kv_xfer)?; + let comp_xfers = comp_xfers_from_args(comp_xfers)?; + let device_indices = device_indices.0; + let actions = py.allow_threads(move || { + self.core() + .commit_load_back(node_id, device_indices, kv_xfer, comp_xfers) + }); + cache_actions_to_py(py, actions) + } + + /// Release a node's device KV once its host copy exists. + fn demote(&self, py: Python<'_>, node_id: NodeId) -> PyResult { + let result = py + .allow_threads(move || self.core().try_demote(node_id)) + .map_err(tree_core_assertion_error)?; + Ok(DemoteResultBinding { + tracker: tracker_to_py(result.tracker), + new_device_frees: frees_to_py(py, result.device_frees)?, + new_host_frees: frees_to_py(py, result.host_frees)?, + }) + } + + /// Evict up to num_tokens of one component's host resources. + fn drive_host_eviction( + &self, + py: Python<'_>, + component_type: u8, + num_tokens: usize, + ) -> PyResult { + let ct = parse_component_type(component_type)?; + let result = py.allow_threads(move || self.core().drive_host_eviction(ct, num_tokens)); + Ok(HostEvictionResultBinding { + tracker: tracker_to_py(result.tracker), + new_device_frees: frees_to_py(py, result.device_frees)?, + new_host_frees: frees_to_py(py, result.host_frees)?, + }) + } + + /// Evict shallow Mamba device checkpoints beyond the per-path cap on the + /// tail's root path; returns the step's freed tensors. + fn evict_excess_path_states( + &self, + py: Python<'_>, + tail_node_id: NodeId, + ) -> PyResult { + let result = py.allow_threads(move || self.core().evict_excess_path_states(tail_node_id)); + Ok(HostEvictionResultBinding { + tracker: tracker_to_py(result.tracker), + new_device_frees: frees_to_py(py, result.device_frees)?, + new_host_frees: frees_to_py(py, result.host_frees)?, + }) + } + + /// Bump the reference count on a node's host-side component locks. + fn inc_host_lock_ref(&self, py: Python<'_>, node_id: NodeId) -> IncLockRefResultBinding { + let result = py.allow_threads(|| self.core().inc_host_lock_ref(node_id)); + IncLockRefResultBinding { + delta: result.delta, + swa_uuid_for_lock: result.swa_uuid_for_lock, + swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, + skip_lock_node_ids: result + .skip_lock_node_ids + .into_iter() + .map(|(ct, node_ids)| (component_type_to_u8(ct), node_ids)) + .collect(), + } + } + + /// Decrease the reference count on a node's host-side component locks. + fn dec_host_lock_ref( + &self, + py: Python<'_>, + node_id: NodeId, + params: Option<&DecLockRefParamsBinding>, + ) -> PyResult<()> { + let params = params.map(|p| p.to_dec_lock_ref_params()).transpose()?; + py.allow_threads(|| self.core().dec_host_lock_ref(node_id, params.as_ref())); + Ok(()) + } + + /// Set the write-back (vs write-through) policy; decided at HiCache init. + fn set_is_write_back(&self, py: Python<'_>, is_write_back: bool) { + py.allow_threads(|| self.core().is_write_back = is_write_back); + } + + /// The current write-back (vs write-through) policy. + fn is_write_back(&self, py: Python<'_>) -> bool { + py.allow_threads(|| self.core().is_write_back) + } + + /// Set the write-through backup hit threshold; decided at HiCache init. + fn set_write_through_threshold(&self, py: Python<'_>, threshold: i64) { + py.allow_threads(|| self.core().write_through_threshold = threshold); + } + + /// The current write-through backup hit threshold. + fn write_through_threshold(&self, py: Python<'_>) -> i64 { + py.allow_threads(|| self.core().write_through_threshold) + } + + /// Mark the storage tier (L3) wired; storage attaches after tree construction. + fn set_enable_storage(&self, py: Python<'_>, value: bool) { + py.allow_threads(|| self.core().set_enable_storage(value)); + } + + /// Whether the storage tier (L3) is wired. + fn enable_storage(&self, py: Python<'_>) -> bool { + py.allow_threads(|| self.core().enable_storage) + } + + /// Queue the all-cleared placement event. + fn record_all_cleared_event(&self, py: Python<'_>) { + py.allow_threads(|| self.core().record_all_cleared_event()); + } + + /// Drain the queued placement events as tagged tuples. + fn take_events(&self, py: Python<'_>) -> PyResult> + where + Vec: IntoPy>, + { + let events = py.allow_threads(|| self.core().take_events()); + let list = PyList::empty_bound(py); + for event in events { + match event { + KvCacheEvent::BlockStored { + block_hashes, + parent_block_hash, + token_ids, + block_size, + medium, + cache_salt, + } => { + let item: Py = ( + "block_stored", + block_hashes, + parent_block_hash, + token_ids, + block_size, + medium.as_str(), + cache_salt.map(|salt| salt.to_string()), + ) + .into_py(py); + list.append(item)?; + } + KvCacheEvent::BlockRemoved { + block_hashes, + medium, + } => { + let item: Py = + ("block_removed", block_hashes, medium.as_str()).into_py(py); + list.append(item)?; + } + KvCacheEvent::AllBlocksCleared => { + let item: Py = ("all_blocks_cleared",).into_py(py); + list.append(item)?; + } + } + } + Ok(list.unbind()) + } + + /// Drop the subtree rooted at an unbacked D-leaf; not dropped when a lock + /// blocks it. + fn drop_subtree_no_host( + &self, + py: Python<'_>, + node_id: NodeId, + ) -> PyResult { + let (dropped, result) = py.allow_threads(move || self.core().drop_subtree_no_host(node_id)); + Ok(DropSubtreeResultBinding { + dropped, + tracker: tracker_to_py(result.tracker), + new_device_frees: frees_to_py(py, result.device_frees)?, + new_host_frees: frees_to_py(py, result.host_frees)?, + }) + } + + /// Mark a node as having an in-flight write-through backup. + fn mark_write_through_pending(&self, py: Python<'_>, node_id: NodeId) { + py.allow_threads(|| self.core().mark_write_through_pending(node_id)); + } + + /// Clear the write-through-pending mark on the acked nodes. + fn finish_write_through(&self, py: Python<'_>, node_ids: Vec, ack_id: NodeId) { + py.allow_threads(|| self.core().finish_write_through(node_ids, ack_id)); + } + + /// Clear the in-flight H->D marks on the anchor's root path at ack time. + fn finish_load_back(&self, py: Python<'_>, anchor_node_id: NodeId) { + py.allow_threads(|| self.core().finish_load_back(anchor_node_id)); + } + + /// Order-sensitive digest of reclaimed coexisting host values. + fn write_back_coexist_reclaim_digest(&self, py: Python<'_>) -> i64 { + py.allow_threads(|| self.core().write_back_coexist_reclaim_digest) + } + + /// Whether the component's data is device-evicted but host-backed. + fn component_has_host_value_only( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + ) -> PyResult { + let ct = parse_component_type(component_type)?; + Ok(py.allow_threads(|| self.core().component_has_host_value_only(node_id, ct))) + } +} + +#[cfg(feature = "inspection")] +impl TreeCoreBinding { + fn inspect_contains_node(&self, py: Python<'_>, node_id: NodeId) -> bool { + py.allow_threads(|| self.core().inspect_contains_node(node_id)) + } + + fn inspect_get_parent_node_id(&self, py: Python<'_>, node_id: NodeId) -> Option { + py.allow_threads(|| self.core().inspect_get_parent_node_id(node_id)) + } + + fn inspect_get_child_node_ids(&self, py: Python<'_>, node_id: NodeId) -> Vec { + py.allow_threads(|| self.core().inspect_get_child_node_ids(node_id)) + } + + fn inspect_get_node_key_length(&self, py: Python<'_>, node_id: NodeId) -> usize { + py.allow_threads(|| self.core().inspect_get_node_key_length(node_id)) + } + + fn inspect_get_node_token_ids(&self, py: Python<'_>, node_id: NodeId) -> Vec { + py.allow_threads(|| self.core().inspect_get_node_token_ids(node_id)) + } + + fn inspect_is_node_key_bigram(&self, py: Python<'_>, node_id: NodeId) -> bool { + py.allow_threads(|| self.core().inspect_is_node_key_bigram(node_id)) + } + + fn inspect_get_component_host_value( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + ) -> PyResult> { + let component_type = parse_component_type(component_type)?; + Ok(py + .allow_threads(|| { + self.core() + .inspect_get_component_host_value(node_id, component_type) + }) + .map(PyTensor)) + } + + fn inspect_get_component_device_lock_ref( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + ) -> PyResult { + let component_type = parse_component_type(component_type)?; + Ok(py.allow_threads(|| { + self.core() + .inspect_get_component_device_lock_ref(node_id, component_type) + })) + } + + fn inspect_get_node_hit_count(&self, py: Python<'_>, node_id: NodeId) -> i64 { + py.allow_threads(|| self.core().inspect_get_node_hit_count(node_id)) + } + + fn inspect_get_write_through_pending_id( + &self, + py: Python<'_>, + node_id: NodeId, + ) -> Option { + py.allow_threads(|| self.core().inspect_get_write_through_pending_id(node_id)) + } + + fn inspect_is_node_in_device_lru( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + ) -> PyResult { + let component_type = parse_component_type(component_type)?; + Ok(py.allow_threads(|| { + self.core() + .inspect_is_node_in_device_lru(node_id, component_type) + })) + } + + fn inspect_is_node_in_host_lru( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + ) -> PyResult { + let component_type = parse_component_type(component_type)?; + Ok(py.allow_threads(|| { + self.core() + .inspect_is_node_in_host_lru(node_id, component_type) + })) + } + + fn inspect_get_component_device_lru_node_ids( + &self, + py: Python<'_>, + component_type: u8, + ) -> PyResult> { + let component_type = parse_component_type(component_type)?; + Ok(py.allow_threads(|| { + self.core() + .inspect_get_component_device_lru_node_ids(component_type) + })) + } + + fn inspect_is_device_evictable_leaf(&self, py: Python<'_>, node_id: NodeId) -> bool { + py.allow_threads(|| self.core().inspect_is_device_evictable_leaf(node_id)) + } + + fn inspect_is_host_evictable_leaf(&self, py: Python<'_>, node_id: NodeId) -> bool { + py.allow_threads(|| self.core().inspect_is_host_evictable_leaf(node_id)) + } + + fn inspect_is_device_leaf(&self, py: Python<'_>, node_id: NodeId) -> bool { + py.allow_threads(|| self.core().inspect_is_device_leaf(node_id)) + } + + fn inspect_get_all_node_ids(&self, py: Python<'_>) -> Vec { + py.allow_threads(|| self.core().inspect_get_all_node_ids()) + } + + fn inspect_component_protected_size( + &self, + py: Python<'_>, + component_type: u8, + ) -> PyResult { + let component_type = parse_component_type(component_type)?; + Ok(py.allow_threads(|| self.core().inspect_component_protected_size(component_type))) + } + + fn inspect_set_node_hash_values( + &self, + py: Python<'_>, + node_id: NodeId, + hash_values: Option>, + ) { + py.allow_threads(move || { + self.core() + .inspect_set_node_hash_values(node_id, hash_values) + }); + } + + fn inspect_set_component_device_value_raw( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + value: Option, + ) -> PyResult<()> { + let component_type = parse_component_type(component_type)?; + let value = value.map(|value| value.0); + py.allow_threads(move || { + self.core() + .inspect_set_component_device_value_raw(node_id, component_type, value) + }); + Ok(()) + } + + fn inspect_set_component_host_value_raw( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + value: Option, + ) -> PyResult<()> { + let component_type = parse_component_type(component_type)?; + let value = value.map(|value| value.0); + py.allow_threads(move || { + self.core() + .inspect_set_component_host_value_raw(node_id, component_type, value) + }); + Ok(()) + } + + fn inspect_set_component_device_lock_ref( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + lock_ref: u32, + ) -> PyResult<()> { + let component_type = parse_component_type(component_type)?; + py.allow_threads(|| { + self.core() + .inspect_set_component_device_lock_ref(node_id, component_type, lock_ref) + }); + Ok(()) + } + + fn inspect_remove_node_from_device_lru( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + ) -> PyResult<()> { + let component_type = parse_component_type(component_type)?; + py.allow_threads(|| { + self.core() + .inspect_remove_node_from_device_lru(node_id, component_type) + }); + Ok(()) + } + + fn inspect_insert_node_into_host_lru( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + ) -> PyResult<()> { + let component_type = parse_component_type(component_type)?; + py.allow_threads(|| { + self.core() + .inspect_insert_node_into_host_lru(node_id, component_type) + }); + Ok(()) + } + + fn inspect_set_component_evictable_size( + &self, + py: Python<'_>, + component_type: u8, + value: usize, + ) -> PyResult<()> { + let component_type = parse_component_type(component_type)?; + py.allow_threads(|| { + self.core() + .inspect_set_component_evictable_size(component_type, value) + }); + Ok(()) + } + + fn inspect_set_component_protected_size( + &self, + py: Python<'_>, + component_type: u8, + value: usize, + ) -> PyResult<()> { + let component_type = parse_component_type(component_type)?; + py.allow_threads(|| { + self.core() + .inspect_set_component_protected_size(component_type, value) + }); + Ok(()) + } + + fn inspect_update_duplicate_tracking(&self, py: Python<'_>, node_id: NodeId) { + py.allow_threads(|| self.core().inspect_update_duplicate_tracking(node_id)); + } + + fn inspect_advance_insert_walk_once(&self, py: Python<'_>) -> PyResult<()> { + py.allow_threads(|| self.core().inspect_advance_insert_walk_once()) + .map_err(PyRuntimeError::new_err) + } + + fn inspect_evict_component( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + target: u8, + ) -> PyResult { + let component_type = parse_component_type(component_type)?; + let target = parse_evict_layer(target)?; + let result = py.allow_threads(|| { + self.core() + .inspect_evict_component(node_id, component_type, target) + }); + HostEvictionResultBinding::from_eviction_step(py, result) + } + + fn inspect_validate_cascade_evict( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + target: u8, + ) -> PyResult<()> { + let component_type = parse_component_type(component_type)?; + let target = parse_evict_layer(target)?; + py.allow_threads(|| { + self.core() + .inspect_validate_cascade_evict(node_id, component_type, target) + }) + .map_err(PyAssertionError::new_err) + } + + fn inspect_cleanup_tombstone_ancestors( + &self, + py: Python<'_>, + node_id: NodeId, + ) -> PyResult { + let result = py.allow_threads(|| self.core().inspect_cleanup_tombstone_ancestors(node_id)); + HostEvictionResultBinding::from_eviction_step(py, result) + } + + #[allow(clippy::too_many_arguments)] + fn inspect_finalize_component_match_result( + &self, + py: Python<'_>, + component_type: u8, + result: InspectionMatchResultInput, + key: &Bound<'_, PyAny>, + extra_key: Option, + cache_salt: Option, + value_chunks: Vec, + best_value_len: usize, + ) -> PyResult { + let component_type = parse_component_type(component_type)?; + let key = K::key_from(Cow::Owned(py_array_to_vec_i64(py, key)?)).into_owned(); + let InspectionMatchResultInput { + device_indices, + last_device_node: last_device_node_id, + last_host_node: last_host_node_id, + best_match_node: best_match_node_id, + host_hit_length, + swa_host_hit_length, + mamba_host_hit_length, + mamba_branching_seqlen, + full_kv_hit_length, + } = result; + let result = MatchResult { + device_indices: device_indices.0, + last_device_node_id, + last_host_node_id, + best_match_node_id, + host_hit_length, + swa_host_hit_length, + mamba_host_hit_length, + mamba_branching_seqlen, + full_kv_hit_length, + cache_actions: Vec::new(), + }; + let value_chunks = value_chunks + .into_iter() + .map(|value| value.0) + .collect::>(); + let result = py.allow_threads(move || { + let params = MatchPrefixParams { + key: &key, + namespace: KeyNamespaceRef::new(extra_key.as_deref(), cache_salt.as_deref()), + }; + self.core().inspect_finalize_component_match_result( + component_type, + result, + ¶ms, + &value_chunks, + best_value_len, + ) + }); + MatchResultBinding::from_match_result(py, result) + } + + fn inspect_build_backup_node_ids( + &self, + py: Python<'_>, + node_id: NodeId, + write_back: bool, + ) -> Vec { + py.allow_threads(|| { + self.core() + .inspect_build_backup_node_ids(node_id, write_back) + }) + } +} + +impl TreeCoreBinding { + /// Print the tree structure for debugging. + fn pretty_print(&self, py: Python<'_>) { + py.allow_threads(|| self.core().pretty_print()); + } +} + +// The delegate surface is identical for every key type; the macro stamps the +// pyclass + pymethods pair per concrete key. +macro_rules! tree_core_binding { + ($(#[$doc:meta])* $name:ident, $key:ty) => { + $(#[$doc])* + #[pyclass] + pub struct $name { + inner: TreeCoreBinding<$key>, + } + + #[pymethods] + impl $name { + /// Build a tree core for the given component types from the cache's + /// init params. + #[new] + fn new(init_params: &TreeCoreInitParamsBinding, component_types: Vec) -> PyResult { + Ok($name { + inner: TreeCoreBinding::new(init_params, component_types)?, + }) + } + + /// Drop the entire tree and reinitialize empty state. + fn reset(&self, py: Python<'_>) { + self.inner.reset(py) + } + + /// Match a key against the tree. + fn match_prefix( + &self, + py: Python<'_>, + params: &MatchParamsBinding, + ) -> PyResult { + self.inner.match_prefix(py, params) + } + + /// The empty match result anchored at the root. + fn empty_match_result(&self, py: Python<'_>) -> PyResult { + self.inner.empty_match_result(py) + } + + /// Insert device values into the tree per the provided key. + fn insert( + &self, + py: Python<'_>, + params: &InsertParamsBinding, + ) -> PyResult { + self.inner.insert(py, params) + } + + /// Start the resumable insert, running to its first barrier or completion. + fn begin_insert( + &self, + py: Python<'_>, + params: &InsertParamsBinding, + ) -> PyResult { + self.inner.begin_insert(py, params) + } + + /// Continue the suspended insert after its step actions were applied. + fn resume_insert(&self, py: Python<'_>) -> PyResult { + self.inner.resume_insert(py) + } + + /// Whether an insert walk is suspended at a barrier. + fn has_ongoing_insert(&self, py: Python<'_>) -> bool { + self.inner.has_ongoing_insert(py) + } + + /// Finish the insert (idempotent); returns still-pending actions to drain. + fn end_insert(&self, py: Python<'_>) -> PyResult> { + self.inner.end_insert(py) + } + + /// Bump the reference count on a node's component locks. + #[pyo3(signature = (node_id, skip_lock_components = None))] + fn inc_lock_ref( + &self, + py: Python<'_>, + node_id: NodeId, + skip_lock_components: Option>, + ) -> PyResult { + self.inner.inc_lock_ref(py, node_id, skip_lock_components) + } + + /// Decrease the reference count on a node's component locks. + #[pyo3(signature = (node_id, params = None, skip_swa = false))] + fn dec_lock_ref( + &self, + py: Python<'_>, + node_id: NodeId, + params: Option<&DecLockRefParamsBinding>, + skip_swa: bool, + ) -> PyResult<()> { + self.inner.dec_lock_ref(py, node_id, params, skip_swa) + } + + /// Early-release the SWA portion of a request's tree lock; returns this + /// release's per-component (device_frees, host_frees). + #[pyo3(signature = (node_id, swa_uuid_for_lock = None, skip_lock_node_ids = None))] + fn dec_swa_lock_only( + &self, + py: Python<'_>, + node_id: NodeId, + swa_uuid_for_lock: Option, + skip_lock_node_ids: Option>>, + ) -> PyResult<(Py, Py)> { + self.inner.dec_swa_lock_only( + py, + node_id, + swa_uuid_for_lock, + skip_lock_node_ids, + ) + } + + /// Store a component's device value on a node (the SWA rebuild write-back). + fn set_component_device_value( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + value: PyTensor, + ) -> PyResult<()> { + self.inner + .set_component_device_value(py, node_id, component_type, value) + } + + /// A component's device value on a node, if set. + fn get_component_device_value( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + ) -> PyResult> { + self.inner + .get_component_device_value(py, node_id, component_type) + } + + /// Begin a component's device-eviction walk for up to request_cnt tokens. + fn evict_device_start( + &self, + py: Python<'_>, + component_type: u8, + request_cnt: usize, + ) -> PyResult<()> { + self.inner + .evict_device_start(py, component_type, request_cnt) + } + + /// The next device leaf to evict, or None when the walk is done; the + /// passed running tracker gates the budget, and the result carries + /// this step's deltas. + fn evict_device_next_node( + &self, + py: Python<'_>, + component_type: u8, + tracker: HashMap, + ) -> PyResult { + self.inner + .evict_device_next_node(py, component_type, tracker) + } + + /// Evict one device leaf; an unbacked write-back leaf returns its backup + /// action for the caller to execute before demoting. + fn evict_device_leaf( + &self, + py: Python<'_>, + node_id: NodeId, + ) -> PyResult { + self.inner.evict_device_leaf(py, node_id) + } + + /// Finish a component's device-eviction walk. + fn evict_device_end(&self, py: Python<'_>, component_type: u8) -> PyResult<()> { + self.inner.evict_device_end(py, component_type) + } + + /// Verify tree-structure, leaf-set, LRU, size, and ongoing-op invariants; + /// ongoing_* args are (id, node_id) pairs. + fn sanity_check( + &self, + py: Python<'_>, + ongoing_write_through: Vec<(i64, NodeId)>, + ongoing_load_back: Vec<(i64, NodeId)>, + ) -> PyResult<()> { + self.inner + .sanity_check(py, ongoing_write_through, ongoing_load_back) + } + + /// Concatenated FULL device values from from_node up to (exclusive) until_node. + fn collect_full_device_indices( + &self, + py: Python<'_>, + from_node_id: NodeId, + until_node_id: NodeId, + ) -> PyTensor { + self.inner + .collect_full_device_indices(py, from_node_id, until_node_id) + } + + /// Every FULL device value in the tree, concatenated. + fn all_values_flatten(&self, py: Python<'_>) -> PyTensor { + self.inner.all_values_flatten(py) + } + + /// Every Mamba device value in the tree, concatenated. + fn all_mamba_values_flatten(&self, py: Python<'_>) -> PyTensor { + self.inner.all_mamba_values_flatten(py) + } + + /// Flatten every FULL device slot into (slot, position, prev-slot) rows for the KV-canary sweep. + fn walk_for_kv_canary( + &self, + py: Python<'_>, + unlocked_only: bool, + swa_resident_only: bool, + ) -> PyResult { + self.inner + .walk_for_kv_canary(py, unlocked_only, swa_resident_only) + } + + /// Evictable token count of the FULL (base) component. + fn evictable_size(&self, py: Python<'_>) -> usize { + self.inner.evictable_size(py) + } + + /// Protected (locked) token count of the FULL (base) component. + fn protected_size(&self, py: Python<'_>) -> usize { + self.inner.protected_size(py) + } + + /// FULL component evictable token count. + fn full_evictable_size(&self, py: Python<'_>) -> usize { + self.inner.full_evictable_size(py) + } + + /// FULL component protected token count. + fn full_protected_size(&self, py: Python<'_>) -> usize { + self.inner.full_protected_size(py) + } + + /// Evictable token count for one component (0 if the component is absent). + fn component_evictable_size(&self, py: Python<'_>, component_type: u8) -> PyResult { + self.inner.component_evictable_size(py, component_type) + } + + /// Protected token count for one component (0 if the component is absent). + fn component_protected_size(&self, py: Python<'_>, component_type: u8) -> PyResult { + self.inner.component_protected_size(py, component_type) + } + + /// (full_tokens, aux_tokens) summed across the whole tree. + fn total_size(&self, py: Python<'_>) -> (usize, usize) { + self.inner.total_size(py) + } + + /// Whether the node's FULL device value has been evicted. + fn is_full_device_evicted(&self, py: Python<'_>, node_id: NodeId) -> bool { + self.inner.is_full_device_evicted(py, node_id) + } + + /// Mark the host tier (HiCache) as wired. + fn set_hicache_enabled(&self, py: Python<'_>) { + self.inner.set_hicache_enabled(py) + } + + /// Whether the host tier (HiCache) is wired. + fn enable_hicache(&self, py: Python<'_>) -> bool { + self.inner.enable_hicache(py) + } + + /// Mark the SWA host pool as wired (HiCache). + fn set_has_swa_host_pool(&self, py: Python<'_>) { + self.inner.set_has_swa_host_pool(py) + } + + /// Whether the SWA host pool is wired. + fn has_swa_host_pool(&self, py: Python<'_>) -> bool { + self.inner.has_swa_host_pool(py) + } + + /// Insert a host-side (backuped) tree path descending from the given node. + #[pyo3(signature = (node_id, extra_key, key, host_value, hash_value, cache_salt = None))] + fn insert_host( + &self, + py: Python<'_>, + node_id: NodeId, + extra_key: Option, + key: &Bound<'_, PyAny>, + host_value: PyTensor, + hash_value: Vec, + cache_salt: Option, + ) -> PyResult { + self.inner.insert_host( + py, + node_id, + extra_key, + key, + host_value, + hash_value, + cache_salt, + ) + } + + /// Gather a node's device value plus per-component BACKUP_HOST transfers. + fn build_backup_spec( + &self, + py: Python<'_>, + node_id: NodeId, + ) -> PyResult<(PyTensor, Py)> { + self.inner.build_backup_spec(py, node_id) + } + + /// Gather a node's device->storage backup spec; None if the node is not backuped. + fn build_storage_backup_spec( + &self, + py: Python<'_>, + node_id: NodeId, + pass_prefix_keys: bool, + ) -> PyResult> { + self.inner + .build_storage_backup_spec(py, node_id, pass_prefix_keys) + } + + /// Route a build_hicache_transfers call to the component for the given type. + #[allow(clippy::too_many_arguments)] + #[pyo3(signature = (component_type, node_id, phase, host_indices = None, token_ids = None, prefetch_tokens = 0, last_hash = None))] + fn build_hicache_transfers( + &self, + py: Python<'_>, + component_type: u8, + node_id: NodeId, + phase: &str, + host_indices: Option, + token_ids: Option>, + prefetch_tokens: usize, + last_hash: Option, + ) -> PyResult>>> { + self.inner.build_hicache_transfers( + py, + component_type, + node_id, + phase, + host_indices, + token_ids, + prefetch_tokens, + last_hash, + ) + } + + /// The anchor node's caller-defined key and cache salt. + fn prefetch_anchor_info( + &self, + py: Python<'_>, + node_id: NodeId, + ) -> PyResult<(Option, Option)> { + self.inner.prefetch_anchor_info(py, node_id) + } + + /// Whether the node's Full KV is present on host. + fn node_backuped(&self, py: Python<'_>, node_id: NodeId) -> PyResult { + self.inner.node_backuped(py, node_id) + } + + /// Whether the node is a (default or named) root. + fn is_root(&self, py: Python<'_>, node_id: NodeId) -> PyResult { + self.inner.is_root(py, node_id) + } + + /// The node's last page hash, or None when it was never hashed. + fn get_last_hash_value( + &self, + py: Python<'_>, + node_id: NodeId, + ) -> PyResult> { + self.inner.get_last_hash_value(py, node_id) + } + + /// The hash chain of the node's ancestors, in root-to-parent order. + fn get_prefix_hash_values( + &self, + py: Python<'_>, + node_id: NodeId, + ) -> PyResult> { + self.inner.get_prefix_hash_values(py, node_id) + } + + fn get_hash_values(&self, py: Python<'_>, node_id: NodeId) -> PyResult> { + self.inner.get_hash_values(py, node_id) + } + + fn snapshot_buffer_backup( + &self, + py: Python<'_>, + node_id: NodeId, + pass_prefix_keys: bool, + ) -> Option { + self.inner + .snapshot_buffer_backup(py, node_id, pass_prefix_keys) + } + + fn validate_buffer_backup( + &self, + py: Python<'_>, + node_id: NodeId, + expected_key_length: usize, + ) -> Option { + self.inner + .validate_buffer_backup(py, node_id, expected_key_length) + } + + /// Hash every node built while storage was disabled. + fn backfill_missing_hash_values(&self, py: Python<'_>) -> usize { + self.inner.backfill_missing_hash_values(py) + } + + #[pyo3(signature = (extra_key = None))] + fn root_node_handle(&self, py: Python<'_>, extra_key: Option) -> NodeId { + self.inner.root_node_handle(py, extra_key) + } + + fn dfs_weight_order( + &self, + py: Python<'_>, + node_ids: Vec, + ) -> PyResult> { + self.inner.dfs_weight_order(py, node_ids) + } + + /// Commit each component's HiCache transfers; returns the new cache actions. + #[pyo3(signature = (node_id, phase, comp_xfers, insert_result = None, pool_storage_result = None))] + fn commit_hicache_transfers( + &self, + py: Python<'_>, + node_id: NodeId, + phase: &str, + comp_xfers: HashMap>, + insert_result: Option<(usize, Option, bool)>, + pool_storage_result: Option<(usize, HashMap)>, + ) -> PyResult<(Py, Option)> { + self.inner.commit_hicache_transfers( + py, + node_id, + phase, + comp_xfers, + insert_result, + pool_storage_result, + ) + } + + /// Commit a successful backup to the node. + fn commit_backup( + &self, + py: Python<'_>, + node_id: NodeId, + host_indices: PyTensor, + comp_xfers: HashMap>, + ) -> PyResult<()> { + self.inner + .commit_backup(py, node_id, host_indices, comp_xfers) + } + + /// Build the H->D load-back KV transfer plus per-component aux transfers. + #[pyo3(signature = (node_id, mamba_pool_idx = None))] + fn build_load_back_spec( + &self, + py: Python<'_>, + node_id: NodeId, + mamba_pool_idx: Option, + ) -> PyResult<(Py, Py)> { + self.inner.build_load_back_spec(py, node_id, mamba_pool_idx) + } + + /// Commit a successful H->D load-back onto the node; returns its actions. + fn commit_load_back( + &self, + py: Python<'_>, + node_id: NodeId, + device_indices: PyTensor, + kv_xfer: TransferArgs, + comp_xfers: HashMap>, + ) -> PyResult> { + self.inner + .commit_load_back(py, node_id, device_indices, kv_xfer, comp_xfers) + } + + /// Release a node's device KV once its host copy exists. + fn demote(&self, py: Python<'_>, node_id: NodeId) -> PyResult { + self.inner.demote(py, node_id) + } + + /// Evict up to num_tokens of one component's host resources. + fn drive_host_eviction( + &self, + py: Python<'_>, + component_type: u8, + num_tokens: usize, + ) -> PyResult { + self.inner.drive_host_eviction(py, component_type, num_tokens) + } + + /// Evict shallow Mamba device checkpoints beyond the per-path cap + /// on the tail's root path. + fn evict_excess_path_states( + &self, + py: Python<'_>, + tail_node_id: NodeId, + ) -> PyResult { + self.inner.evict_excess_path_states(py, tail_node_id) + } + + /// Bump the reference count on a node's host-side component locks. + fn inc_host_lock_ref(&self, py: Python<'_>, node_id: NodeId) -> IncLockRefResultBinding { + self.inner.inc_host_lock_ref(py, node_id) + } + + /// Decrease the reference count on a node's host-side component locks. + #[pyo3(signature = (node_id, params = None))] + fn dec_host_lock_ref( + &self, + py: Python<'_>, + node_id: NodeId, + params: Option<&DecLockRefParamsBinding>, + ) -> PyResult<()> { + self.inner.dec_host_lock_ref(py, node_id, params) + } + + /// Set the write-back (vs write-through) policy; decided at HiCache init. + fn set_is_write_back(&self, py: Python<'_>, is_write_back: bool) { + self.inner.set_is_write_back(py, is_write_back) + } + + /// The current write-back (vs write-through) policy. + fn is_write_back(&self, py: Python<'_>) -> bool { + self.inner.is_write_back(py) + } + + /// Set the write-through backup hit threshold; decided at HiCache init. + fn set_write_through_threshold(&self, py: Python<'_>, threshold: i64) { + self.inner.set_write_through_threshold(py, threshold) + } + + /// The current write-through backup hit threshold. + fn write_through_threshold(&self, py: Python<'_>) -> i64 { + self.inner.write_through_threshold(py) + } + + /// Mark the storage tier (L3) wired; storage attaches after tree construction. + fn set_enable_storage(&self, py: Python<'_>, value: bool) { + self.inner.set_enable_storage(py, value) + } + + /// Whether the storage tier (L3) is wired. + fn enable_storage(&self, py: Python<'_>) -> bool { + self.inner.enable_storage(py) + } + + /// Queue the all-cleared placement event. + fn record_all_cleared_event(&self, py: Python<'_>) { + self.inner.record_all_cleared_event(py) + } + + /// Drain the queued placement events as tagged tuples. + fn take_events(&self, py: Python<'_>) -> PyResult> { + self.inner.take_events(py) + } + + /// Drop the subtree rooted at an unbacked D-leaf; not dropped when a lock + /// blocks it. + fn drop_subtree_no_host( + &self, + py: Python<'_>, + node_id: NodeId, + ) -> PyResult { + self.inner.drop_subtree_no_host(py, node_id) + } + + /// Mark a node as having an in-flight write-through backup. + fn mark_write_through_pending(&self, py: Python<'_>, node_id: NodeId) { + self.inner.mark_write_through_pending(py, node_id) + } + + /// Clear the write-through-pending mark on the acked nodes. + fn finish_write_through(&self, py: Python<'_>, node_ids: Vec, ack_id: NodeId) { + self.inner.finish_write_through(py, node_ids, ack_id) + } + + /// Clear the in-flight H->D marks on the anchor's root path at ack time. + fn finish_load_back(&self, py: Python<'_>, anchor_node_id: NodeId) { + self.inner.finish_load_back(py, anchor_node_id) + } + + /// Order-sensitive digest of reclaimed coexisting host values. + #[pyo3(name = "write_back_duplicate_reclaim_digest")] + fn write_back_coexist_reclaim_digest(&self, py: Python<'_>) -> i64 { + self.inner.write_back_coexist_reclaim_digest(py) + } + + /// Whether the component's data is device-evicted but host-backed. + fn component_has_host_value_only( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + ) -> PyResult { + self.inner + .component_has_host_value_only(py, node_id, component_type) + } + + // ==== Test-only inspection surface ==== + + #[cfg(feature = "inspection")] + fn inspect_contains_node(&self, py: Python<'_>, node_id: NodeId) -> bool { + self.inner.inspect_contains_node(py, node_id) + } + + #[cfg(feature = "inspection")] + fn inspect_get_parent_node_id( + &self, + py: Python<'_>, + node_id: NodeId, + ) -> Option { + self.inner.inspect_get_parent_node_id(py, node_id) + } + + #[cfg(feature = "inspection")] + fn inspect_get_child_node_ids( + &self, + py: Python<'_>, + node_id: NodeId, + ) -> Vec { + self.inner.inspect_get_child_node_ids(py, node_id) + } + + #[cfg(feature = "inspection")] + fn inspect_get_node_key_length(&self, py: Python<'_>, node_id: NodeId) -> usize { + self.inner.inspect_get_node_key_length(py, node_id) + } + + #[cfg(feature = "inspection")] + fn inspect_get_node_token_ids(&self, py: Python<'_>, node_id: NodeId) -> Vec { + self.inner.inspect_get_node_token_ids(py, node_id) + } + + #[cfg(feature = "inspection")] + fn inspect_is_node_key_bigram(&self, py: Python<'_>, node_id: NodeId) -> bool { + self.inner.inspect_is_node_key_bigram(py, node_id) + } + + #[cfg(feature = "inspection")] + fn inspect_get_component_host_value( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + ) -> PyResult> { + self.inner + .inspect_get_component_host_value(py, node_id, component_type) + } + + #[cfg(feature = "inspection")] + fn inspect_get_component_device_lock_ref( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + ) -> PyResult { + self.inner + .inspect_get_component_device_lock_ref(py, node_id, component_type) + } + + #[cfg(feature = "inspection")] + fn inspect_get_node_hit_count(&self, py: Python<'_>, node_id: NodeId) -> i64 { + self.inner.inspect_get_node_hit_count(py, node_id) + } + + #[cfg(feature = "inspection")] + fn inspect_get_write_through_pending_id( + &self, + py: Python<'_>, + node_id: NodeId, + ) -> Option { + self.inner + .inspect_get_write_through_pending_id(py, node_id) + } + + #[cfg(feature = "inspection")] + fn inspect_is_node_in_device_lru( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + ) -> PyResult { + self.inner + .inspect_is_node_in_device_lru(py, node_id, component_type) + } + + #[cfg(feature = "inspection")] + fn inspect_is_node_in_host_lru( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + ) -> PyResult { + self.inner + .inspect_is_node_in_host_lru(py, node_id, component_type) + } + + #[cfg(feature = "inspection")] + fn inspect_get_component_device_lru_node_ids( + &self, + py: Python<'_>, + component_type: u8, + ) -> PyResult> { + self.inner + .inspect_get_component_device_lru_node_ids(py, component_type) + } + + #[cfg(feature = "inspection")] + fn inspect_is_device_evictable_leaf( + &self, + py: Python<'_>, + node_id: NodeId, + ) -> bool { + self.inner.inspect_is_device_evictable_leaf(py, node_id) + } + + #[cfg(feature = "inspection")] + fn inspect_is_host_evictable_leaf( + &self, + py: Python<'_>, + node_id: NodeId, + ) -> bool { + self.inner.inspect_is_host_evictable_leaf(py, node_id) + } + + #[cfg(feature = "inspection")] + fn inspect_is_device_leaf(&self, py: Python<'_>, node_id: NodeId) -> bool { + self.inner.inspect_is_device_leaf(py, node_id) + } + + #[cfg(feature = "inspection")] + fn inspect_get_all_node_ids(&self, py: Python<'_>) -> Vec { + self.inner.inspect_get_all_node_ids(py) + } + + #[cfg(feature = "inspection")] + fn inspect_component_protected_size( + &self, + py: Python<'_>, + component_type: u8, + ) -> PyResult { + self.inner + .inspect_component_protected_size(py, component_type) + } + + #[cfg(feature = "inspection")] + #[pyo3(signature = (node_id, hash_values = None))] + fn inspect_set_node_hash_values( + &self, + py: Python<'_>, + node_id: NodeId, + hash_values: Option>, + ) { + self.inner + .inspect_set_node_hash_values(py, node_id, hash_values) + } + + #[cfg(feature = "inspection")] + #[pyo3(signature = (node_id, component_type, value = None))] + fn inspect_set_component_device_value_raw( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + value: Option, + ) -> PyResult<()> { + self.inner.inspect_set_component_device_value_raw( + py, + node_id, + component_type, + value, + ) + } + + #[cfg(feature = "inspection")] + #[pyo3(signature = (node_id, component_type, value = None))] + fn inspect_set_component_host_value_raw( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + value: Option, + ) -> PyResult<()> { + self.inner.inspect_set_component_host_value_raw( + py, + node_id, + component_type, + value, + ) + } + + #[cfg(feature = "inspection")] + fn inspect_set_component_device_lock_ref( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + lock_ref: u32, + ) -> PyResult<()> { + self.inner.inspect_set_component_device_lock_ref( + py, + node_id, + component_type, + lock_ref, + ) + } + + #[cfg(feature = "inspection")] + fn inspect_remove_node_from_device_lru( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + ) -> PyResult<()> { + self.inner + .inspect_remove_node_from_device_lru(py, node_id, component_type) + } + + #[cfg(feature = "inspection")] + fn inspect_insert_node_into_host_lru( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + ) -> PyResult<()> { + self.inner + .inspect_insert_node_into_host_lru(py, node_id, component_type) + } + + #[cfg(feature = "inspection")] + fn inspect_set_component_evictable_size( + &self, + py: Python<'_>, + component_type: u8, + value: usize, + ) -> PyResult<()> { + self.inner + .inspect_set_component_evictable_size(py, component_type, value) + } + + #[cfg(feature = "inspection")] + fn inspect_set_component_protected_size( + &self, + py: Python<'_>, + component_type: u8, + value: usize, + ) -> PyResult<()> { + self.inner + .inspect_set_component_protected_size(py, component_type, value) + } + + #[cfg(feature = "inspection")] + fn inspect_update_duplicate_tracking(&self, py: Python<'_>, node_id: NodeId) { + self.inner.inspect_update_duplicate_tracking(py, node_id) + } + + #[cfg(feature = "inspection")] + fn inspect_advance_insert_walk_once(&self, py: Python<'_>) -> PyResult<()> { + self.inner.inspect_advance_insert_walk_once(py) + } + + #[cfg(feature = "inspection")] + fn inspect_evict_component( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + target: u8, + ) -> PyResult { + self.inner + .inspect_evict_component(py, node_id, component_type, target) + } + + #[cfg(feature = "inspection")] + fn inspect_validate_cascade_evict( + &self, + py: Python<'_>, + node_id: NodeId, + component_type: u8, + target: u8, + ) -> PyResult<()> { + self.inner + .inspect_validate_cascade_evict(py, node_id, component_type, target) + } + + #[cfg(feature = "inspection")] + fn inspect_cleanup_tombstone_ancestors( + &self, + py: Python<'_>, + node_id: NodeId, + ) -> PyResult { + self.inner + .inspect_cleanup_tombstone_ancestors(py, node_id) + } + + #[cfg(feature = "inspection")] + #[allow(clippy::too_many_arguments)] + #[pyo3(signature = (component_type, result, key, extra_key, cache_salt, value_chunks, best_value_len))] + fn inspect_finalize_component_match_result( + &self, + py: Python<'_>, + component_type: u8, + result: InspectionMatchResultInput, + key: &Bound<'_, PyAny>, + extra_key: Option, + cache_salt: Option, + value_chunks: Vec, + best_value_len: usize, + ) -> PyResult { + self.inner.inspect_finalize_component_match_result( + py, + component_type, + result, + key, + extra_key, + cache_salt, + value_chunks, + best_value_len, + ) + } + + #[cfg(feature = "inspection")] + #[pyo3(signature = (node_id, write_back = false))] + fn inspect_build_backup_node_ids( + &self, + py: Python<'_>, + node_id: NodeId, + write_back: bool, + ) -> Vec { + self.inner + .inspect_build_backup_node_ids(py, node_id, write_back) + } + + /// Print the tree structure for debugging. + fn pretty_print(&self, py: Python<'_>) { + self.inner.pretty_print(py) + } + } + }; +} + +tree_core_binding!( + /// The UnifiedTreeCore Python adapter over single-token (unigram) child keys. + RustUnifiedTreeCoreBinding, + Vec +); + +tree_core_binding!( + /// The UnifiedTreeCore Python adapter over bigram (EAGLE) child keys; keys + /// cross the boundary as raw token ids and pair up rust-side. + RustBigramUnifiedTreeCoreBinding, + Vec<(i64, i64)> +); + +/// Per-page chained hashes over raw token ids. +#[pyfunction] +#[pyo3(signature = (token_ids, prior_hash, page_size, is_bigram = false))] +fn get_hash_str( + py: Python<'_>, + token_ids: &Bound<'_, PyAny>, + prior_hash: Option, + page_size: usize, + is_bigram: bool, +) -> PyResult> { + let raw = py_array_to_vec_i64(py, token_ids)?; + if page_size == 0 { + return Err(PyValueError::new_err("page_size must be positive")); + } + if let Some(prior_hash) = prior_hash.as_deref().filter(|hash| !hash.is_empty()) + && (prior_hash.len() != 64 || !prior_hash.bytes().all(|byte| byte.is_ascii_hexdigit())) + { + return Err(PyValueError::new_err( + "prior_hash must be a 64-character hexadecimal digest", + )); + } + if let Some(token_id) = raw + .iter() + .find(|token_id| u32::try_from(**token_id).is_err()) + { + return Err(PyValueError::new_err(format!( + "token id {token_id} does not fit in uint32" + ))); + } + Ok(py.allow_threads(move || { + if is_bigram { + let key = as ChildKeyType>::key_from(Cow::Owned(raw)).into_owned(); + crate::node::get_hash_str::>(&key, prior_hash.as_deref(), page_size) + } else { + crate::node::get_hash_str::>(&raw, prior_hash.as_deref(), page_size) + } + })) +} + +fn register_mem_cache_module(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_function(wrap_pyfunction!(get_hash_str, m)?)?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + Ok(()) +} + +/// The production TreeCore extension module. +#[pymodule] +fn mem_cache(m: &Bound<'_, PyModule>) -> PyResult<()> { + register_mem_cache_module(m) +} + +/// White-box variant used only by the shared test inspector. +#[cfg(feature = "inspection")] +#[pymodule] +fn mem_cache_inspection(m: &Bound<'_, PyModule>) -> PyResult<()> { + register_mem_cache_module(m) +} diff --git a/rust/mem-cache/src/tests/components/base.rs b/rust/mem-cache/src/tests/components/base.rs new file mode 100644 index 000000000..b0700bf6c --- /dev/null +++ b/rust/mem-cache/src/tests/components/base.rs @@ -0,0 +1,160 @@ +use super::*; +use crate::components::FULL; +use crate::unified_tree_core::CacheInitParams; + +// Test-only component exercising the trait defaults; abstract hooks stay unimplemented. +struct DefaultComponentForTest; + +impl TreeComponent> for DefaultComponentForTest { + fn component_type(&self) -> ComponentType { + FULL + } + + fn create_match_validator( + &self, + _tree_core: &UnifiedTreeCore>, + match_device_only: bool, + ) -> Box>, NodeIdx_) -> bool> { + unimplemented!() + } + + fn redistribute_on_node_split( + &self, + tree_core: &mut UnifiedTreeCore>, + new_parent_id: NodeIdx_, + child_id: NodeIdx_, + ) { + unimplemented!() + } + + fn evict_component( + &self, + tree_core: &mut UnifiedTreeCore>, + node_id: NodeIdx_, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + target: EvictLayer, + ) -> (usize, usize) { + unimplemented!() + } + + fn evict_device_start(&self, tree_core: &mut UnifiedTreeCore>, request_cnt: usize) { + unimplemented!() + } + + fn evict_device_next_node( + &self, + tree_core: &mut UnifiedTreeCore>, + tracker: &mut HashMap, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + ) -> Option { + unimplemented!() + } + + fn evict_device_end(&self, tree_core: &mut UnifiedTreeCore>) { + unimplemented!() + } + + fn acquire_component_lock( + &self, + tree_core: &mut UnifiedTreeCore>, + node_id: NodeIdx_, + result: IncLockRefResult, + lock_host: bool, + ) -> IncLockRefResult { + unimplemented!() + } + + fn release_component_lock( + &self, + tree_core: &mut UnifiedTreeCore>, + node_id: NodeIdx_, + params: Option<&DecLockRefParams>, + lock_host: bool, + ) { + unimplemented!() + } +} + +#[test] +fn insert_overlap_default_consumes_nothing() { + let mut tc: UnifiedTreeCore> = + UnifiedTreeCore::new(CacheInitParams::default(), vec![FULL]); + let root = tc.arena.root(); + let consumed_from = DefaultComponentForTest.update_component_on_insert_overlap( + &mut tc, + root, + /* prefix_len = */ 3, + /* total_prefix_len = */ 0, + Tensor::from_slice(&[0i64, 1, 2]), + &InsertParams { + key: &vec![0, 1, 2], + namespace: Default::default(), + value: Tensor::from_slice(&[0i64, 1, 2]), + mamba_value: None, + prev_prefix_len: 0, + swa_evicted_seqlen: 0, + chunked: false, + priority: 0, + track_adopted_ranges: false, + }, + &mut InsertResult::default(), + &mut Vec::new(), + ); + // Nothing consumed: the whole overlap stays freeable as duplicates. + assert_eq!(consumed_from, 3); +} + +#[test] +fn finalize_match_result_default_returns_result_unchanged() { + let tc: UnifiedTreeCore> = + UnifiedTreeCore::new(CacheInitParams::default(), vec![FULL]); + let result = MatchResult { + last_device_node_id: 3, + best_match_node_id: 7, + host_hit_length: 11, + ..tc.empty_match_result() + }; + let out = DefaultComponentForTest.finalize_match_result_in_tree_core( + &tc, + result, + &MatchPrefixParams { + key: &Vec::new(), + namespace: Default::default(), + }, + &[], + 0, + ); + assert_eq!(out.last_device_node_id, 3); + assert_eq!(out.best_match_node_id, 7); + assert_eq!(out.host_hit_length, 11); +} + +#[test] +fn drive_host_eviction_default_is_a_noop() { + let mut tc: UnifiedTreeCore> = + UnifiedTreeCore::new(CacheInitParams::default(), vec![FULL]); + let mut tracker = HashMap::from([(FULL, 5usize)]); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + DefaultComponentForTest.drive_host_eviction( + &mut tc, + /* num_tokens = */ 100, + &mut tracker, + &mut device_frees, + &mut host_frees, + ); + assert_eq!(tracker[&FULL], 5); + assert!(device_frees.is_empty()); + assert!(host_frees.is_empty()); +} + +// Component types. + +#[test] +fn idx_matches_discriminants() { + assert_eq!(ComponentType::Full.idx(), 0); + assert_eq!(ComponentType::Swa.idx(), 1); + assert_eq!(ComponentType::Mamba.idx(), 2); +} diff --git a/rust/mem-cache/src/tests/components/full.rs b/rust/mem-cache/src/tests/components/full.rs new file mode 100644 index 000000000..0794cc636 --- /dev/null +++ b/rust/mem-cache/src/tests/components/full.rs @@ -0,0 +1,2828 @@ +use super::*; +use crate::components::FULL; +use crate::test_utils::accumulate_step; +use crate::unified_tree_core::CacheInitParams; + +fn core() -> UnifiedTreeCore> { + UnifiedTreeCore::new(CacheInitParams::default(), vec![FULL]) +} + +// Raw seeding for states set_value rejects: mid-split (key trimmed before the value +// splits) and present-but-empty semantics pins. +fn set_value_no_check( + tc: &mut UnifiedTreeCore, + node: NodeIdx_, + slot: ValueSlotIdx, + value: Tensor, +) { + tc.arena.node_mut(node).state_mut_(slot).value = Some(value); +} + +#[test] +fn eviction_priority_is_lower_for_leaf_than_internal() { + let full = FullComponent; + assert_eq!( + >>::eviction_priority(&full, true), + 0 + ); + assert_eq!( + >>::eviction_priority(&full, false), + 2 + ); +} + +// Three value-bearing root children in the D-leaf set; ticks anti-correlated +// with allocation order so priority order differs from NodeIdx_ order. +fn evict_walk_setup(tc: &mut UnifiedTreeCore>) -> (NodeIdx_, NodeIdx_, NodeIdx_) { + let root = tc.arena.root(); + let ticks = [30i64, 10, 20]; + let mut nodes = [NodeIdx_(0); 3]; + for (i, node) in nodes.iter_mut().enumerate() { + let id = tc + .arena + .alloc_child( + root, + /* key = */ vec![i as i64], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(id, FULL, Tensor::from_slice(&[i as i64])); + tc.arena.node_mut(id).last_access_counter = ticks[i]; + tc.evictable_device_leaves.add(id); + *node = id; + } + (nodes[0], nodes[1], nodes[2]) +} + +fn match_params(key: &Vec) -> MatchPrefixParams<'_, Vec> { + MatchPrefixParams { + key, + namespace: Default::default(), + } +} + +fn insert(tc: &mut UnifiedTreeCore>, key: &Vec, value: &[i64]) { + tc.insert(&crate::unified_tree_core::InsertParams { + key, + namespace: Default::default(), + value: Tensor::from_slice(value), + mamba_value: None, + prev_prefix_len: 0, + swa_evicted_seqlen: 0, + chunked: false, + priority: 0, + track_adopted_ranges: false, + }); +} + +fn tracker() -> HashMap { + HashMap::from([(FULL, 0)]) +} + +fn frees() -> HashMap> { + HashMap::new() +} + +#[test] +fn evict_walk_pops_leaves_lowest_priority_first() { + let mut tc = core(); + let (a, b, c) = evict_walk_setup(&mut tc); + let (mut tr, mut df, mut hf) = (tracker(), frees(), frees()); + tc.evict_device_start(FULL, /* request_cnt = */ 100); + let (node, step) = tc.evict_device_next_node(FULL, &tr); + accumulate_step(step, &mut tr, &mut df, &mut hf); + assert_eq!(node, Some(tc.arena.node(b).id)); + // The driver evicts each returned leaf before asking for the next. + tc.evictable_device_leaves.discard(b); + let (node, step) = tc.evict_device_next_node(FULL, &tr); + accumulate_step(step, &mut tr, &mut df, &mut hf); + assert_eq!(node, Some(tc.arena.node(c).id)); + tc.evictable_device_leaves.discard(c); + let (node, step) = tc.evict_device_next_node(FULL, &tr); + accumulate_step(step, &mut tr, &mut df, &mut hf); + assert_eq!(node, Some(tc.arena.node(a).id)); + tc.evictable_device_leaves.discard(a); + let (node, step) = tc.evict_device_next_node(FULL, &tr); + accumulate_step(step, &mut tr, &mut df, &mut hf); + assert_eq!(node, None); + tc.evict_device_end(FULL); +} + +#[test] +fn evict_walk_stops_at_the_token_budget() { + let mut tc = core(); + let (_a, b, _c) = evict_walk_setup(&mut tc); + let (mut tr, mut df, mut hf) = (tracker(), frees(), frees()); + tc.evict_device_start(FULL, /* request_cnt = */ 5); + let (node, step) = tc.evict_device_next_node(FULL, &tr); + accumulate_step(step, &mut tr, &mut df, &mut hf); + assert_eq!(node, Some(tc.arena.node(b).id)); + // The driver's tracker reaching the budget ends the walk. + *tr.get_mut(&FULL).unwrap() = 5; + let (node, step) = tc.evict_device_next_node(FULL, &tr); + accumulate_step(step, &mut tr, &mut df, &mut hf); + assert_eq!(node, None); + // The budget-None return still resets the cursor. + assert_eq!(tc.component_state(FULL).evict_device_cursor, None); + tc.evict_device_end(FULL); +} + +#[test] +fn evict_walk_stops_when_the_tracker_overshoots_the_budget() { + let mut tc = core(); + let (_a, b, _c) = evict_walk_setup(&mut tc); + let (mut tr, mut df, mut hf) = (tracker(), frees(), frees()); + tc.evict_device_start(FULL, /* request_cnt = */ 5); + let (node, step) = tc.evict_device_next_node(FULL, &tr); + accumulate_step(step, &mut tr, &mut df, &mut hf); + assert_eq!(node, Some(tc.arena.node(b).id)); + // Multi-token evictions jump past the budget; the gate must still fire. + *tr.get_mut(&FULL).unwrap() = 7; + let (node, step) = tc.evict_device_next_node(FULL, &tr); + accumulate_step(step, &mut tr, &mut df, &mut hf); + assert_eq!(node, None); + assert_eq!(tc.component_state(FULL).evict_device_cursor, None); + tc.evict_device_end(FULL); +} + +#[test] +fn evict_walk_reports_done_when_the_baseline_already_meets_the_budget() { + let mut tc = core(); + let (_a, _b, _c) = evict_walk_setup(&mut tc); + tc.evict_device_start(FULL, /* request_cnt = */ 5); + // Prior steps' evictions reach the gate through the baseline. + let baseline = HashMap::from([(FULL, 5)]); + let (node, _step) = tc.evict_device_next_node(FULL, &baseline); + assert_eq!(node, None); + tc.evict_device_end(FULL); +} + +#[test] +fn evict_walk_step_tracker_stays_empty_when_nothing_was_evicted() { + let mut tc = core(); + let (_a, b, _c) = evict_walk_setup(&mut tc); + tc.evict_device_start(FULL, /* request_cnt = */ 100); + let (node, step) = tc.evict_device_next_node(FULL, &tracker()); + assert_eq!(node, Some(tc.arena.node(b).id)); + // The FULL walk frees nothing itself: no zero-delta entries leak out. + assert!(step.tracker.is_empty()); + tc.evict_device_end(FULL); +} + +#[test] +fn evict_walk_skips_nodes_that_left_the_leaf_set() { + let mut tc = core(); + let (_a, b, c) = evict_walk_setup(&mut tc); + let (mut tr, mut df, mut hf) = (tracker(), frees(), frees()); + tc.evict_device_start(FULL, /* request_cnt = */ 100); + // The lowest-priority b stops being a D-leaf after the heap was built + // (e.g. locked): the walk skips it for the next candidate. + tc.evictable_device_leaves.discard(b); + let (node, step) = tc.evict_device_next_node(FULL, &tr); + accumulate_step(step, &mut tr, &mut df, &mut hf); + assert_eq!(node, Some(tc.arena.node(c).id)); + tc.evict_device_end(FULL); +} + +#[test] +fn evict_walk_readmits_the_freed_leafs_parent() { + let mut tc = core(); + let root = tc.arena.root(); + let p = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let c = tc + .arena + .alloc_child( + p, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + // Two live sibling leaves whose ticks straddle the parent's tick, so + // the readmitted parent must compete on its own priority. + let s1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![3], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let s2 = tc + .arena + .alloc_child( + root, + /* key = */ vec![4], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + for (id, tick) in [(c, 5i64), (s1, 10), (p, 20), (s2, 30)] { + tc.arena + .set_device_value(id, FULL, Tensor::from_slice(&[tick])); + tc.arena.node_mut(id).last_access_counter = tick; + } + tc.evictable_device_leaves.add(c); + tc.evictable_device_leaves.add(s1); + tc.evictable_device_leaves.add(s2); + let (mut tr, mut df, mut hf) = (tracker(), frees(), frees()); + tc.evict_device_start(FULL, /* request_cnt = */ 100); + let (node, step) = tc.evict_device_next_node(FULL, &tr); + accumulate_step(step, &mut tr, &mut df, &mut hf); + assert_eq!(node, Some(tc.arena.node(c).id)); + // The driver deletes the write-through leaf outright and its parent + // becomes the new D-leaf; the walk must still find it via the parent + // captured before the free, ordered between the surviving siblings. + tc.evictable_device_leaves.discard(c); + let _ = tc.arena.take_device_value(c, FULL); + tc.arena.free_leaf(c).unwrap(); + tc.evictable_device_leaves.add(p); + let (node, step) = tc.evict_device_next_node(FULL, &tr); + accumulate_step(step, &mut tr, &mut df, &mut hf); + assert_eq!(node, Some(tc.arena.node(s1).id)); + tc.evictable_device_leaves.discard(s1); + let (node, step) = tc.evict_device_next_node(FULL, &tr); + accumulate_step(step, &mut tr, &mut df, &mut hf); + assert_eq!(node, Some(tc.arena.node(p).id)); + tc.evictable_device_leaves.discard(p); + let (node, step) = tc.evict_device_next_node(FULL, &tr); + accumulate_step(step, &mut tr, &mut df, &mut hf); + assert_eq!(node, Some(tc.arena.node(s2).id)); + tc.evict_device_end(FULL); +} + +#[test] +fn evict_end_clears_the_walk_and_allows_a_restart() { + let mut tc = core(); + let (_a, b, _c) = evict_walk_setup(&mut tc); + let (mut tr, mut df, mut hf) = (tracker(), frees(), frees()); + tc.evict_device_start(FULL, /* request_cnt = */ 100); + tc.evict_device_end(FULL); + // A fresh walk rebuilds the heap from the leaf set. + tc.evict_device_start(FULL, /* request_cnt = */ 100); + let (node, step) = tc.evict_device_next_node(FULL, &tr); + accumulate_step(step, &mut tr, &mut df, &mut hf); + assert_eq!(node, Some(tc.arena.node(b).id)); + tc.evict_device_end(FULL); +} + +#[test] +#[should_panic(expected = "Full device eviction already in progress")] +fn evict_start_panics_when_already_ongoing() { + let mut tc = core(); + tc.evict_device_start(FULL, /* request_cnt = */ 1); + tc.evict_device_start(FULL, /* request_cnt = */ 1); +} + +#[test] +#[should_panic(expected = "Full device eviction not started")] +fn evict_next_panics_before_start() { + let mut tc = core(); + let tr = tracker(); + tc.evict_device_next_node(FULL, &tr); +} + +// Three host-backed root children in the H-leaf set; ticks anti-correlated +// with allocation order so priority order differs from NodeIdx_ order. +fn host_walk_setup(tc: &mut UnifiedTreeCore>) -> (NodeIdx_, NodeIdx_, NodeIdx_) { + let root = tc.arena.root(); + let ticks = [30i64, 10, 20]; + let mut nodes = [NodeIdx_(0); 3]; + for (i, node) in nodes.iter_mut().enumerate() { + let id = tc + .arena + .alloc_child( + root, + /* key = */ vec![i as i64], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(id, FULL, Tensor::from_slice(&[i as i64])); + tc.arena.node_mut(id).last_access_counter = ticks[i]; + tc.evictable_host_leaves.add(id); + *node = id; + } + (nodes[0], nodes[1], nodes[2]) +} + +#[test] +fn host_drive_evicts_leaves_lowest_priority_first_until_the_budget() { + let mut tc = core(); + let (a, b, c) = host_walk_setup(&mut tc); + let (mut tr, mut df, mut hf) = (tracker(), frees(), frees()); + accumulate_step( + tc.drive_host_eviction(FULL, /* num_tokens = */ 2), + &mut tr, + &mut df, + &mut hf, + ); + // b (tick 10) and c (tick 20) go; a (tick 30) survives. + assert_eq!(tr[&FULL], 2); + assert!(!tc.evictable_host_leaves.contains(b)); + assert!(!tc.evictable_host_leaves.contains(c)); + assert!(tc.evictable_host_leaves.contains(a)); + assert_eq!(hf[&FULL].len(), 2); + assert!(df.is_empty()); + assert_eq!(tc.arena.len(), 2); + tc.sanity_check(&[], &[]); +} + +#[test] +fn host_drive_stops_when_a_leaf_overshoots_the_budget() { + let mut tc = core(); + let root = tc.arena.root(); + let big = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2, 3], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let small = tc + .arena + .alloc_child( + root, + /* key = */ vec![9], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(big, FULL, Tensor::from_slice(&[10i64, 11, 12])); + tc.arena + .set_host_value(small, FULL, Tensor::from_slice(&[20i64])); + tc.arena.node_mut(big).last_access_counter = 1; + tc.arena.node_mut(small).last_access_counter = 2; + tc.evictable_host_leaves.add(big); + tc.evictable_host_leaves.add(small); + let (mut tr, mut df, mut hf) = (tracker(), frees(), frees()); + accumulate_step( + tc.drive_host_eviction(FULL, /* num_tokens = */ 2), + &mut tr, + &mut df, + &mut hf, + ); + // The 3-token leaf jumps past the 2-token budget; the walk still stops. + assert_eq!(tr[&FULL], 3); + assert!(tc.evictable_host_leaves.contains(small)); + assert_eq!(hf[&FULL].len(), 1); + tc.sanity_check(&[], &[]); +} + +#[test] +fn host_drive_reclaims_coexisting_host_values_while_sparing_the_device_leaf() { + let mut tc = write_back_core(); + insert(&mut tc, &vec![1, 2], &[10, 11]); + insert(&mut tc, &vec![1, 2, 3], &[10, 11, 12]); + let leaf_handle = tc + .match_prefix(&match_params(&vec![1, 2, 3])) + .best_match_node_id; + let leaf = tc.arena.resolve(leaf_handle); + let parent = tc.arena.node(leaf).parent(); + tc.commit_backup( + tc.arena.node(parent).id, + Tensor::from_slice(&[20i64, 21]), + HashMap::new(), + ); + tc.commit_backup(leaf_handle, Tensor::from_slice(&[22i64]), HashMap::new()); + assert!(tc.evictable_host_leaves.is_empty()); + + let (mut tr, mut df, mut hf) = (tracker(), frees(), frees()); + accumulate_step( + tc.drive_host_eviction(FULL, /* num_tokens = */ 2), + &mut tr, + &mut df, + &mut hf, + ); + assert_eq!(tr[&FULL], 2); + assert!(!tc.arena.node(parent).has_host_value(FULL)); + assert!(tc.arena.node(parent).has_device_value(FULL)); + assert!(tc.arena.node(leaf).has_host_value(FULL)); + assert_ne!(tc.write_back_coexist_reclaim_digest, 0); + tc.sanity_check(&[], &[]); +} + +#[test] +fn host_drive_spares_coexisting_host_values_under_an_in_flight_transfer() { + let mut tc = write_back_core(); + insert(&mut tc, &vec![1, 2], &[10, 11]); + let handle = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + tc.commit_backup(handle, Tensor::from_slice(&[20i64, 21]), HashMap::new()); + tc.mark_write_through_pending(handle); + + let (mut tr, mut df, mut hf) = (tracker(), frees(), frees()); + accumulate_step( + tc.drive_host_eviction(FULL, /* num_tokens = */ 2), + &mut tr, + &mut df, + &mut hf, + ); + assert_eq!(tr[&FULL], 0); + assert!(tc.arena.node(tc.arena.resolve(handle)).has_host_value(FULL)); + + tc.finish_write_through(vec![handle], handle); + accumulate_step( + tc.drive_host_eviction(FULL, /* num_tokens = */ 2), + &mut tr, + &mut df, + &mut hf, + ); + assert_eq!(tr[&FULL], 2); + assert!(!tc.arena.node(tc.arena.resolve(handle)).has_host_value(FULL)); + tc.sanity_check(&[], &[]); +} + +#[test] +fn host_drive_is_a_noop_without_host_leaves() { + let mut tc = core(); + tc.insert(&crate::unified_tree_core::InsertParams { + key: &vec![1, 2], + namespace: Default::default(), + value: Tensor::from_slice(&[10i64, 11]), + mamba_value: None, + prev_prefix_len: 0, + swa_evicted_seqlen: 0, + chunked: false, + priority: 0, + track_adopted_ranges: false, + }); + let (mut tr, mut df, mut hf) = (tracker(), frees(), frees()); + accumulate_step( + tc.drive_host_eviction(FULL, /* num_tokens = */ 5), + &mut tr, + &mut df, + &mut hf, + ); + assert_eq!(tr[&FULL], 0); + assert!(hf.is_empty()); + assert_eq!(tc.arena.len(), 2); + tc.sanity_check(&[], &[]); +} + +#[test] +fn host_drive_readmits_the_freed_leafs_parent() { + let mut tc = core(); + let root = tc.arena.root(); + let p = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let c1 = tc + .arena + .alloc_child( + p, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let c2 = tc + .arena + .alloc_child( + p, + /* key = */ vec![3], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + for (id, tick) in [(p, 30i64), (c1, 10), (c2, 20)] { + tc.arena + .set_host_value(id, FULL, Tensor::from_slice(&[tick])); + tc.arena.node_mut(id).last_access_counter = tick; + } + tc.evictable_host_leaves.add(c1); + tc.evictable_host_leaves.add(c2); + let (mut tr, mut df, mut hf) = (tracker(), frees(), frees()); + // p only becomes an H-leaf once both children are gone; the readmission + // after c2 lets one drive drain the whole chain. + accumulate_step( + tc.drive_host_eviction(FULL, /* num_tokens = */ 100), + &mut tr, + &mut df, + &mut hf, + ); + assert_eq!(tr[&FULL], 3); + assert_eq!(hf[&FULL].len(), 3); + assert_eq!(tc.arena.len(), 1); + assert!(tc.evictable_host_leaves.is_empty()); + tc.sanity_check(&[], &[]); +} + +#[test] +fn host_drive_skips_a_stale_heap_entry_for_an_already_freed_leaf() { + let mut tc = core(); + let root = tc.arena.root(); + let p = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let c = tc + .arena + .alloc_child( + p, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(p, FULL, Tensor::from_slice(&[10i64])); + tc.arena + .set_host_value(c, FULL, Tensor::from_slice(&[20i64])); + tc.arena.node_mut(p).last_access_counter = 2; + tc.arena.node_mut(c).last_access_counter = 1; + // p staged in the set despite its child: the initial heap entry goes + // stale once c's eviction readmits (and then frees) p. + tc.evictable_host_leaves.add(p); + tc.evictable_host_leaves.add(c); + let (mut tr, mut df, mut hf) = (tracker(), frees(), frees()); + accumulate_step( + tc.drive_host_eviction(FULL, /* num_tokens = */ 100), + &mut tr, + &mut df, + &mut hf, + ); + // The duplicate p entry is skipped instead of double-freeing. + assert_eq!(tr[&FULL], 2); + assert_eq!(hf[&FULL].len(), 2); + assert_eq!(tc.arena.len(), 1); + tc.sanity_check(&[], &[]); +} + +// Chain root -> n1 (len 2) -> n2 (len 3) with FULL device values; evictable seeded to 5. +fn lock_chain(tc: &mut UnifiedTreeCore>) -> (NodeIdx_, NodeIdx_) { + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 11], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let n2 = tc + .arena + .alloc_child( + n1, + /* key = */ vec![2, 22, 222], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1])); + tc.arena + .set_device_value(n2, FULL, Tensor::from_slice(&[0i64, 1, 2])); + tc.component_state_mut(FULL).evictable_size = 5; + tc.evictable_device_leaves.add(n2); + (n1, n2) +} + +#[test] +fn inc_lock_ref_locks_the_device_path() { + let mut tc = core(); + let (n1, n2) = lock_chain(&mut tc); + let result = tc.inc_lock_ref(tc.arena.node(n2).id); + assert_eq!(result.delta, Some(5)); + assert!(result.skip_lock_node_ids.is_empty()); + assert_eq!(tc.arena.device_lock_ref(n1, FULL), 1); + assert_eq!(tc.arena.device_lock_ref(n2, FULL), 1); + let state = tc.component_state(FULL); + assert_eq!(state.evictable_size, 0); + assert_eq!(state.protected_size, 5); + assert!(!tc.evictable_device_leaves.contains(n2)); +} + +#[test] +fn inc_lock_ref_again_only_bumps_the_refs() { + let mut tc = core(); + let (n1, n2) = lock_chain(&mut tc); + tc.inc_lock_ref(tc.arena.node(n2).id); + let result = tc.inc_lock_ref(tc.arena.node(n2).id); + assert_eq!(result.delta, Some(0)); + assert_eq!(tc.arena.device_lock_ref(n1, FULL), 2); + assert_eq!(tc.arena.device_lock_ref(n2, FULL), 2); + let state = tc.component_state(FULL); + assert_eq!(state.evictable_size, 0); + assert_eq!(state.protected_size, 5); +} + +#[test] +fn inc_lock_ref_counts_only_newly_locked_nodes() { + // n1 is already locked via its own path; locking n2 moves only n2's tokens. + let mut tc = core(); + let (n1, n2) = lock_chain(&mut tc); + tc.inc_lock_ref(tc.arena.node(n1).id); + let result = tc.inc_lock_ref(tc.arena.node(n2).id); + assert_eq!(result.delta, Some(3)); + assert_eq!(tc.arena.device_lock_ref(n1, FULL), 2); + assert_eq!(tc.arena.device_lock_ref(n2, FULL), 1); + let state = tc.component_state(FULL); + assert_eq!(state.evictable_size, 0); + assert_eq!(state.protected_size, 5); +} + +#[test] +fn inc_lock_ref_collects_the_evicted_bottom_segment() { + // n2 and n3 are evicted (no device value): the walk records both and locks only n1. + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 11], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let n2 = tc + .arena + .alloc_child( + n1, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let n3 = tc + .arena + .alloc_child( + n2, + /* key = */ vec![3], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1])); + tc.component_state_mut(FULL).evictable_size = 2; + tc.evictable_device_leaves.add(n1); + let result = tc.inc_lock_ref(tc.arena.node(n3).id); + assert_eq!(result.delta, Some(2)); + assert_eq!( + result.skip_lock_node_ids[&FULL], + HashSet::from([tc.arena.node(n2).id, tc.arena.node(n3).id]) + ); + assert_eq!(tc.arena.device_lock_ref(n1, FULL), 1); + assert_eq!(tc.arena.device_lock_ref(n2, FULL), 0); + assert_eq!(tc.arena.device_lock_ref(n3, FULL), 0); + // The locked ancestor leaves the D-leaf set. + assert!(!tc.evictable_device_leaves.contains(n1)); +} + +#[test] +fn lock_round_trips_on_a_root_anchor_are_noops() { + let mut tc = core(); + let root = tc.arena.root(); + let result = tc.inc_lock_ref(tc.arena.node(root).id); + assert_eq!(result.delta, Some(0)); + assert!(result.skip_lock_node_ids.is_empty()); + // The protected root keeps its construction-time lock through the pair. + assert_eq!(tc.arena.device_lock_ref(root, FULL), 1); + tc.dec_lock_ref( + tc.arena.node(root).id, + /* params = */ None, + /* skip_swa = */ false, + ); + assert_eq!(tc.arena.device_lock_ref(root, FULL), 1); +} + +#[test] +fn lock_walks_stop_at_the_root_of_a_salted_chain() { + let mut tc = core(); + let lora = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + lora, + /* key = */ vec![1, 11], + /* priority = */ 0, + Some("lora-1"), + ) + .unwrap(); + tc.arena + .set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1])); + tc.component_state_mut(FULL).evictable_size = 2; + let result = tc.inc_lock_ref(tc.arena.node(n1).id); + assert_eq!(result.delta, Some(2)); + assert_eq!(tc.arena.device_lock_ref(n1, FULL), 1); + // The root keeps its construction-time lock untouched. + assert_eq!(tc.arena.device_lock_ref(lora, FULL), 1); + // The release walk stops at the same boundary. + tc.dec_lock_ref( + tc.arena.node(n1).id, + /* params = */ None, + /* skip_swa = */ false, + ); + assert_eq!(tc.arena.device_lock_ref(n1, FULL), 0); + assert_eq!(tc.arena.device_lock_ref(lora, FULL), 1); +} + +#[test] +fn lock_walks_treat_a_present_but_empty_value_as_device_on() { + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let empty: [i64; 0] = []; + set_value_no_check( + &mut tc, + n1, + ValueSlotIdx::device(FULL), + Tensor::from_slice(&empty), + ); + let result = tc.inc_lock_ref(tc.arena.node(n1).id); + // A present-but-empty value is device-on (Python `value is not None`): + // locked, zero tokens moved. + assert_eq!(result.delta, Some(0)); + assert_eq!(tc.arena.device_lock_ref(n1, FULL), 1); + assert!(result.skip_lock_node_ids.is_empty()); + // The release side moves the same zero tokens back. + tc.dec_lock_ref( + tc.arena.node(n1).id, + /* params = */ None, + /* skip_swa = */ false, + ); + assert_eq!(tc.arena.device_lock_ref(n1, FULL), 0); + let state = tc.component_state(FULL); + assert_eq!(state.evictable_size, 0); + assert_eq!(state.protected_size, 0); +} + +#[test] +fn dec_lock_ref_unlocks_and_restores_sizes() { + let mut tc = core(); + let (n1, n2) = lock_chain(&mut tc); + tc.inc_lock_ref(tc.arena.node(n2).id); + tc.dec_lock_ref( + tc.arena.node(n2).id, + /* params = */ None, + /* skip_swa = */ false, + ); + assert_eq!(tc.arena.device_lock_ref(n1, FULL), 0); + assert_eq!(tc.arena.device_lock_ref(n2, FULL), 0); + let state = tc.component_state(FULL); + assert_eq!(state.evictable_size, 5); + assert_eq!(state.protected_size, 0); + // The unlocked leaf re-enters the D-leaf set; its valued-child parent does not. + assert!(tc.evictable_device_leaves.contains(n2)); + assert!(!tc.evictable_device_leaves.contains(n1)); +} + +#[test] +fn dec_lock_ref_replays_the_skip_set() { + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 11], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let n2 = tc + .arena + .alloc_child( + n1, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let n3 = tc + .arena + .alloc_child( + n2, + /* key = */ vec![3], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1])); + tc.component_state_mut(FULL).evictable_size = 2; + let result = tc.inc_lock_ref(tc.arena.node(n3).id); + let params = DecLockRefParams { + skip_lock_node_ids: result.skip_lock_node_ids, + ..Default::default() + }; + // The still-evicted n2 and n3 are skipped instead of tripping the lock asserts. + tc.dec_lock_ref( + tc.arena.node(n3).id, + Some(¶ms), + /* skip_swa = */ false, + ); + assert_eq!(tc.arena.device_lock_ref(n1, FULL), 0); + assert_eq!(tc.arena.device_lock_ref(n2, FULL), 0); + assert_eq!(tc.arena.device_lock_ref(n3, FULL), 0); + assert_eq!(tc.evictable_size_(FULL), 2); + // The unlocked ancestor (whose child is valueless) re-enters the D-leaf set. + assert!(tc.evictable_device_leaves.contains(n1)); +} + +#[test] +fn temp_lock_skips_the_evicted_anchor_and_mirrors_on_release() { + // Chain root -> a -> y -> anchor with FULL device values; the anchor is evicted. + let mut tc = core(); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let y = tc + .arena + .alloc_child( + a, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let anchor = tc + .arena + .alloc_child( + y, + /* key = */ vec![3], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(a, FULL, Tensor::from_slice(&[0i64])); + tc.arena + .set_device_value(y, FULL, Tensor::from_slice(&[0i64])); + tc.component_state_mut(FULL).evictable_size = 3; + // The temp lock records the evicted anchor and locks only its ancestors. + let temp_lock = tc.inc_lock_ref(tc.arena.node(anchor).id); + assert_eq!(tc.arena.device_lock_ref(anchor, FULL), 0); + assert_eq!(tc.arena.device_lock_ref(y, FULL), 1); + assert_eq!(tc.arena.device_lock_ref(a, FULL), 1); + assert_eq!( + temp_lock.skip_lock_node_ids[&FULL], + HashSet::from([tc.arena.node(anchor).id]) + ); + // A load-back restores the anchor; the second acquire covers it. + tc.arena + .set_device_value(anchor, FULL, Tensor::from_slice(&[0i64])); + let second_lock = tc.inc_lock_ref(tc.arena.node(anchor).id); + assert_eq!(tc.arena.device_lock_ref(anchor, FULL), 1); + assert_eq!(tc.arena.device_lock_ref(y, FULL), 2); + assert_eq!(tc.arena.device_lock_ref(a, FULL), 2); + // Releasing the temp lock mirrors its skip set: the anchor keeps its lock. + let temp_params = DecLockRefParams { + skip_lock_node_ids: temp_lock.skip_lock_node_ids, + ..Default::default() + }; + tc.dec_lock_ref( + tc.arena.node(anchor).id, + Some(&temp_params), + /* skip_swa = */ false, + ); + assert_eq!(tc.arena.device_lock_ref(anchor, FULL), 1); + assert_eq!(tc.arena.device_lock_ref(y, FULL), 1); + assert_eq!(tc.arena.device_lock_ref(a, FULL), 1); + let second_params = DecLockRefParams { + skip_lock_node_ids: second_lock.skip_lock_node_ids, + ..Default::default() + }; + tc.dec_lock_ref( + tc.arena.node(anchor).id, + Some(&second_params), + /* skip_swa = */ false, + ); + assert_eq!(tc.arena.device_lock_ref(anchor, FULL), 0); + assert_eq!(tc.arena.device_lock_ref(y, FULL), 0); + assert_eq!(tc.arena.device_lock_ref(a, FULL), 0); +} + +#[test] +#[should_panic(expected = "has no FULL device value")] +fn dec_lock_ref_panics_without_replaying_the_skip_set() { + // Dropping the acquire's skip set makes the release walk hit the tombstone. + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 11], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let n2 = tc + .arena + .alloc_child( + n1, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1])); + tc.component_state_mut(FULL).evictable_size = 2; + tc.inc_lock_ref(tc.arena.node(n2).id); + tc.dec_lock_ref( + tc.arena.node(n2).id, + /* params = */ None, + /* skip_swa = */ false, + ); +} + +#[test] +fn dec_lock_ref_with_skip_swa_still_releases_full() { + let mut tc = core(); + let (_n1, n2) = lock_chain(&mut tc); + tc.inc_lock_ref(tc.arena.node(n2).id); + tc.dec_lock_ref( + tc.arena.node(n2).id, + /* params = */ None, + /* skip_swa = */ true, + ); + assert_eq!(tc.arena.device_lock_ref(n2, FULL), 0); +} + +#[test] +fn nested_locks_release_pairwise() { + // Two acquires then two releases: sizes move only on the outermost pair. + let mut tc = core(); + let (_n1, n2) = lock_chain(&mut tc); + tc.inc_lock_ref(tc.arena.node(n2).id); + tc.inc_lock_ref(tc.arena.node(n2).id); + tc.dec_lock_ref( + tc.arena.node(n2).id, + /* params = */ None, + /* skip_swa = */ false, + ); + let state = tc.component_state(FULL); + assert_eq!(state.evictable_size, 0); + assert_eq!(state.protected_size, 5); + assert!(!tc.evictable_device_leaves.contains(n2)); + tc.dec_lock_ref( + tc.arena.node(n2).id, + /* params = */ None, + /* skip_swa = */ false, + ); + let state = tc.component_state(FULL); + assert_eq!(state.evictable_size, 5); + assert_eq!(state.protected_size, 0); + assert!(tc.evictable_device_leaves.contains(n2)); +} + +#[test] +#[should_panic(expected = "is not locked")] +fn dec_lock_ref_panics_on_an_unlocked_node() { + let mut tc = core(); + let (_n1, n2) = lock_chain(&mut tc); + tc.dec_lock_ref( + tc.arena.node(n2).id, + /* params = */ None, + /* skip_swa = */ false, + ); +} + +#[test] +#[should_panic(expected = "FULL invariant broken: evicted ancestor")] +fn inc_lock_ref_panics_on_an_evicted_ancestor() { + // A value-less n1 below a valued n2 breaks the FULL bottom-up invariant. + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let n2 = tc + .arena + .alloc_child( + n1, + /* key = */ vec![2, 22, 222], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(n2, FULL, Tensor::from_slice(&[0i64, 1, 2])); + tc.component_state_mut(FULL).evictable_size = 3; + tc.inc_lock_ref(tc.arena.node(n2).id); +} + +#[test] +#[should_panic(expected = "evictable size underflow")] +fn inc_lock_ref_panics_when_evictable_size_is_unaccounted() { + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(n1, FULL, Tensor::from_slice(&[0i64])); + tc.inc_lock_ref(tc.arena.node(n1).id); +} + +#[test] +#[should_panic(expected = "protected size underflow")] +fn dec_lock_ref_panics_on_protected_underflow() { + // A lock ref not accounted through acquire trips the checked release. + let mut tc = core(); + let (_n1, n2) = lock_chain(&mut tc); + tc.arena + .node_mut(n2) + .set_lock_ref_(ValueSlotIdx::device(FULL), 1); + tc.dec_lock_ref( + tc.arena.node(n2).id, + /* params = */ None, + /* skip_swa = */ false, + ); +} + +fn write_back_core() -> UnifiedTreeCore> { + UnifiedTreeCore::new( + CacheInitParams { + is_write_back: true, + ..Default::default() + }, + vec![FULL], + ) +} + +// A backuped anchor: host value only, seeded into the H-leaf set. +fn host_lock_anchor(tc: &mut UnifiedTreeCore>) -> NodeIdx_ { + let root = tc.arena.root(); + let node = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(node, FULL, Tensor::from_slice(&[10i64, 11])); + tc.evictable_host_leaves.add(node); + node +} + +#[test] +fn inc_host_lock_ref_pins_the_backuped_anchor() { + let mut tc = core(); + let node = host_lock_anchor(&mut tc); + tc.component_state_mut(FULL).evictable_size = 7; + let result = tc.inc_host_lock_ref(tc.arena.node(node).id); + assert_eq!(result.delta, None); + assert!(result.skip_lock_node_ids.is_empty()); + assert_eq!(tc.arena.host_lock_ref(node, FULL), 1); + // The pinned anchor leaves the H-leaf set; the device tier is untouched. + assert!(!tc.evictable_host_leaves.contains(node)); + assert_eq!(tc.arena.device_lock_ref(node, FULL), 0); + let state = tc.component_state(FULL); + assert_eq!(state.evictable_size, 7); + assert_eq!(state.protected_size, 0); +} + +#[test] +fn inc_host_lock_ref_again_only_bumps_the_counter() { + let mut tc = core(); + let node = host_lock_anchor(&mut tc); + tc.inc_host_lock_ref(tc.arena.node(node).id); + tc.inc_host_lock_ref(tc.arena.node(node).id); + assert_eq!(tc.arena.host_lock_ref(node, FULL), 2); + assert!(!tc.evictable_host_leaves.contains(node)); +} + +#[test] +fn inc_host_lock_ref_pins_only_the_anchor_not_its_ancestors() { + // Both chain nodes are backuped; only the anchor's host counter moves. + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let n2 = tc + .arena + .alloc_child( + n1, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(n1, FULL, Tensor::from_slice(&[0i64])); + tc.arena + .set_host_value(n2, FULL, Tensor::from_slice(&[0i64])); + tc.inc_host_lock_ref(tc.arena.node(n2).id); + assert_eq!(tc.arena.host_lock_ref(n2, FULL), 1); + assert_eq!(tc.arena.host_lock_ref(n1, FULL), 0); +} + +#[test] +fn inc_host_lock_ref_skips_an_anchor_without_a_host_value() { + let mut tc = core(); + let (_n1, n2) = lock_chain(&mut tc); + tc.inc_host_lock_ref(tc.arena.node(n2).id); + assert_eq!(tc.arena.host_lock_ref(n2, FULL), 0); +} + +#[test] +fn host_lock_round_trips_on_a_root_anchor_are_noops() { + let mut tc = core(); + let root = tc.arena.root(); + let result = tc.inc_host_lock_ref(tc.arena.node(root).id); + assert_eq!(result.delta, None); + assert_eq!(tc.arena.host_lock_ref(root, FULL), 0); + tc.dec_host_lock_ref(tc.arena.node(root).id, /* params = */ None); + assert_eq!(tc.arena.host_lock_ref(root, FULL), 0); +} + +#[test] +fn inc_host_lock_ref_under_write_back_pins_a_device_only_anchor() { + let mut tc = write_back_core(); + let (_n1, n2) = lock_chain(&mut tc); + tc.inc_host_lock_ref(tc.arena.node(n2).id); + assert_eq!(tc.arena.host_lock_ref(n2, FULL), 1); + // The write-back host lock is a pure counter: no size shifts. + let state = tc.component_state(FULL); + assert_eq!(state.evictable_size, 5); + assert_eq!(state.protected_size, 0); +} + +#[test] +fn dec_host_lock_ref_unpins_and_restores_the_h_leaf_set() { + let mut tc = core(); + let node = host_lock_anchor(&mut tc); + tc.component_state_mut(FULL).evictable_size = 7; + tc.inc_host_lock_ref(tc.arena.node(node).id); + tc.dec_host_lock_ref(tc.arena.node(node).id, /* params = */ None); + assert_eq!(tc.arena.host_lock_ref(node, FULL), 0); + assert!(tc.evictable_host_leaves.contains(node)); + let state = tc.component_state(FULL); + assert_eq!(state.evictable_size, 7); + assert_eq!(state.protected_size, 0); +} + +#[test] +fn dec_host_lock_ref_on_an_unlocked_anchor_is_a_noop() { + let mut tc = core(); + let node = host_lock_anchor(&mut tc); + tc.dec_host_lock_ref(tc.arena.node(node).id, /* params = */ None); + assert_eq!(tc.arena.host_lock_ref(node, FULL), 0); +} + +#[test] +fn dec_host_lock_ref_keeps_the_counter_when_the_host_value_is_gone() { + // A host-evicted anchor keeps its pin count under write-through. + let mut tc = core(); + let node = host_lock_anchor(&mut tc); + tc.inc_host_lock_ref(tc.arena.node(node).id); + let _ = tc.arena.take_host_value(node, FULL); + tc.dec_host_lock_ref(tc.arena.node(node).id, /* params = */ None); + assert_eq!(tc.arena.host_lock_ref(node, FULL), 1); +} + +#[test] +fn host_lock_round_trip_under_write_back_is_a_pure_counter() { + let mut tc = write_back_core(); + let (_n1, n2) = lock_chain(&mut tc); + tc.inc_host_lock_ref(tc.arena.node(n2).id); + tc.dec_host_lock_ref(tc.arena.node(n2).id, /* params = */ None); + assert_eq!(tc.arena.host_lock_ref(n2, FULL), 0); + let state = tc.component_state(FULL); + assert_eq!(state.evictable_size, 5); + assert_eq!(state.protected_size, 0); +} + +#[test] +fn acquire_host_arm_updates_the_h_leaf_set_without_the_dispatcher() { + let mut tc = core(); + let node = host_lock_anchor(&mut tc); + FullComponent.acquire_component_lock( + &mut tc, + node, + IncLockRefResult::default(), + /* lock_host = */ true, + ); + assert!(!tc.evictable_host_leaves.contains(node)); +} + +#[test] +fn release_host_arm_updates_the_h_leaf_set_without_the_dispatcher() { + let mut tc = core(); + let node = host_lock_anchor(&mut tc); + tc.inc_host_lock_ref(tc.arena.node(node).id); + FullComponent.release_component_lock( + &mut tc, node, /* params = */ None, /* lock_host = */ true, + ); + assert!(tc.evictable_host_leaves.contains(node)); +} + +#[test] +fn nested_host_locks_release_pairwise() { + let mut tc = core(); + let node = host_lock_anchor(&mut tc); + tc.inc_host_lock_ref(tc.arena.node(node).id); + tc.inc_host_lock_ref(tc.arena.node(node).id); + tc.dec_host_lock_ref(tc.arena.node(node).id, /* params = */ None); + assert_eq!(tc.arena.host_lock_ref(node, FULL), 1); + assert!(!tc.evictable_host_leaves.contains(node)); + tc.dec_host_lock_ref(tc.arena.node(node).id, /* params = */ None); + assert_eq!(tc.arena.host_lock_ref(node, FULL), 0); + assert!(tc.evictable_host_leaves.contains(node)); +} + +#[test] +fn node_has_component_data_tracks_each_slot() { + let mut tc = core(); + let root = tc.arena.root(); + let node = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + assert!(!crate::components::node_has_component_data( + &tc.arena, + node, + FULL, + EvictLayer::Device + )); + assert!(!crate::components::node_has_component_data( + &tc.arena, + node, + FULL, + EvictLayer::Host + )); + + tc.arena + .set_device_value(node, FULL, Tensor::from_slice(&[10i64])); + assert!(crate::components::node_has_component_data( + &tc.arena, + node, + FULL, + EvictLayer::Device + )); + assert!(!crate::components::node_has_component_data( + &tc.arena, + node, + FULL, + EvictLayer::Host + )); + + tc.arena + .set_host_value(node, FULL, Tensor::from_slice(&[20i64])); + assert!(crate::components::node_has_component_data( + &tc.arena, + node, + FULL, + EvictLayer::Host + )); +} + +#[test] +fn match_validator_device_only_accepts_device_backed_only() { + let mut tc = core(); + let root = tc.arena.root(); + let dev = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let host = tc + .arena + .alloc_child( + root, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let empty = tc + .arena + .alloc_child( + root, + /* key = */ vec![3], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(dev, FULL, Tensor::from_slice(&[0i64])); + tc.arena + .set_host_value(host, FULL, Tensor::from_slice(&[0i64])); + let mut validator = >>::create_match_validator( + &FullComponent, + &tc, + true, + ); + assert!(validator(&tc, dev)); + assert!(!validator(&tc, host)); + assert!(!validator(&tc, empty)); +} + +#[test] +fn match_validator_hicache_accepts_device_or_host_backed() { + let mut tc = core(); + let root = tc.arena.root(); + let dev = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let host = tc + .arena + .alloc_child( + root, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let empty = tc + .arena + .alloc_child( + root, + /* key = */ vec![3], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(dev, FULL, Tensor::from_slice(&[0i64])); + tc.arena + .set_host_value(host, FULL, Tensor::from_slice(&[0i64])); + let mut validator = >>::create_match_validator( + &FullComponent, + &tc, + false, + ); + assert!(validator(&tc, dev)); + assert!(validator(&tc, host)); + assert!(!validator(&tc, empty)); +} + +#[test] +fn match_validator_accepts_node_with_both_device_and_host() { + let mut tc = core(); + let root = tc.arena.root(); + let both = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(both, FULL, Tensor::from_slice(&[0i64])); + tc.arena + .set_host_value(both, FULL, Tensor::from_slice(&[0i64])); + let mut device_only = >>::create_match_validator( + &FullComponent, + &tc, + true, + ); + let mut hicache = >>::create_match_validator( + &FullComponent, + &tc, + false, + ); + assert!(device_only(&tc, both)); + assert!(hicache(&tc, both)); +} + +#[test] +fn match_validator_verdict_is_per_node_not_stateful() { + let mut tc = core(); + let root = tc.arena.root(); + let dev = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let empty = tc + .arena + .alloc_child( + root, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(dev, FULL, Tensor::from_slice(&[0i64])); + // One validator reused across nodes returns each node's own verdict (FULL is stateless). + let mut validator = >>::create_match_validator( + &FullComponent, + &tc, + true, + ); + assert!(validator(&tc, dev)); + assert!(!validator(&tc, empty)); + assert!(validator(&tc, dev)); + assert!(!validator(&tc, empty)); +} + +#[test] +fn match_validator_accepts_present_but_empty_device_value() { + let mut tc = core(); + let root = tc.arena.root(); + let empty_dev = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let no_value = tc + .arena + .alloc_child( + root, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let empty: [i64; 0] = []; + set_value_no_check( + &mut tc, + empty_dev, + ValueSlotIdx::device(FULL), + Tensor::from_slice(&empty), + ); + // A present-but-empty device value is a boundary (Python `value is not None`); a + // truly value-less node is not. + let mut device_only = >>::create_match_validator( + &FullComponent, + &tc, + true, + ); + assert!(device_only(&tc, empty_dev)); + assert!(!device_only(&tc, no_value)); +} + +// Chain root -> n1 -> n2 -> n3 with FULL host values on n2 (len 2) and n3 (len 3). +fn host_hit_chain() -> (UnifiedTreeCore>, NodeIdx_, NodeIdx_) { + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let n2 = tc + .arena + .alloc_child( + n1, + /* key = */ vec![2, 22], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let n3 = tc + .arena + .alloc_child( + n2, + /* key = */ vec![3, 33, 333], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(n2, FULL, Tensor::from_slice(&[0i64, 1])); + tc.arena + .set_host_value(n3, FULL, Tensor::from_slice(&[0i64, 1, 2])); + (tc, n1, n3) +} + +fn finalize(tc: &UnifiedTreeCore>, result: MatchResult) -> MatchResult { + FullComponent.finalize_match_result_in_tree_core( + tc, + result, + &MatchPrefixParams { + key: &Vec::new(), + namespace: Default::default(), + }, + &[], + 0, + ) +} + +#[test] +fn finalize_sums_full_host_hits_between_best_and_last_device() { + let (tc, n1, n3) = host_hit_chain(); + let out = finalize( + &tc, + MatchResult { + last_device_node_id: tc.arena.node(n1).id, + best_match_node_id: tc.arena.node(n3).id, + host_hit_length: 0, + ..tc.empty_match_result() + }, + ); + assert_eq!(out.host_hit_length, 5); + // Other MatchResult fields pass through unchanged (Python returns via _replace). + assert_eq!(out.last_device_node_id, tc.arena.node(n1).id); + assert_eq!(out.best_match_node_id, tc.arena.node(n3).id); +} + +#[test] +fn finalize_keeps_larger_existing_host_hit_length() { + let (tc, n1, n3) = host_hit_chain(); + let out = finalize( + &tc, + MatchResult { + last_device_node_id: tc.arena.node(n1).id, + best_match_node_id: tc.arena.node(n3).id, + host_hit_length: 100, + ..tc.empty_match_result() + }, + ); + assert_eq!(out.host_hit_length, 100); +} + +#[test] +fn finalize_overrides_smaller_existing_host_hit_length() { + let (tc, n1, n3) = host_hit_chain(); + let out = finalize( + &tc, + MatchResult { + last_device_node_id: tc.arena.node(n1).id, + best_match_node_id: tc.arena.node(n3).id, + host_hit_length: 2, + ..tc.empty_match_result() + }, + ); + assert_eq!(out.host_hit_length, 5); +} + +#[test] +fn finalize_is_noop_without_full_host_values() { + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let n2 = tc + .arena + .alloc_child( + n1, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let out = finalize( + &tc, + MatchResult { + last_device_node_id: tc.arena.node(root).id, + best_match_node_id: tc.arena.node(n2).id, + host_hit_length: 4, + ..tc.empty_match_result() + }, + ); + assert_eq!(out.host_hit_length, 4); +} + +#[test] +fn finalize_walks_a_salted_chain_up_to_the_root() { + let mut tc = core(); + let lora = tc.arena.root(); + let a = tc + .arena + .alloc_child( + lora, + /* key = */ vec![1, 11], + /* priority = */ 0, + Some("lora-1"), + ) + .unwrap(); + let b = tc + .arena + .alloc_child( + a, + /* key = */ vec![2, 22, 222], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(a, FULL, Tensor::from_slice(&[0i64, 1])); + tc.arena + .set_host_value(b, FULL, Tensor::from_slice(&[0i64, 1, 2])); + // The root is last_device_node_id; its own host value is never counted. + set_value_no_check( + &mut tc, + lora, + ValueSlotIdx::host(FULL), + Tensor::from_slice(&[0i64, 1, 2, 3]), + ); + let out = finalize( + &tc, + MatchResult { + last_device_node_id: tc.arena.node(lora).id, + best_match_node_id: tc.arena.node(b).id, + host_hit_length: 0, + ..tc.empty_match_result() + }, + ); + assert_eq!(out.host_hit_length, 5); +} + +#[test] +#[should_panic(expected = "hit root")] +fn finalize_panics_when_walk_hits_root_before_last_device() { + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + // last_device_node_id is not an ancestor of best_match_node_id -> corrupt result. + finalize( + &tc, + MatchResult { + last_device_node_id: tc.arena.node(n1).id, + best_match_node_id: tc.arena.node(root).id, + host_hit_length: 0, + ..tc.empty_match_result() + }, + ); +} + +#[test] +fn match_validator_hicache_accepts_present_but_empty_host_value() { + let mut tc = core(); + let root = tc.arena.root(); + let host = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let empty: [i64; 0] = []; + set_value_no_check( + &mut tc, + host, + ValueSlotIdx::host(FULL), + Tensor::from_slice(&empty), + ); + // A present-but-empty host value is backuped (Python `host_value is not None`). + let mut hicache = >>::create_match_validator( + &FullComponent, + &tc, + false, + ); + let mut device_only = >>::create_match_validator( + &FullComponent, + &tc, + true, + ); + assert!(hicache(&tc, host)); + assert!(!device_only(&tc, host)); +} + +#[test] +#[should_panic(expected = "out of bounds")] +fn match_validator_panics_on_missing_node() { + let tc = core(); + let mut validator = >>::create_match_validator( + &FullComponent, + &tc, + true, + ); + validator(&tc, NodeIdx_(999)); +} + +#[test] +#[should_panic(expected = "is not allocated")] +fn finalize_panics_on_missing_best_match_node() { + let tc = core(); + let root = tc.arena.root(); + finalize( + &tc, + MatchResult { + last_device_node_id: tc.arena.node(root).id, + best_match_node_id: 999, + host_hit_length: 0, + ..tc.empty_match_result() + }, + ); +} + +#[test] +fn finalize_empty_walk_when_best_equals_last_device() { + let (tc, _n1, n3) = host_hit_chain(); + // best_match_node_id == last_device_node_id -> no hops -> unchanged. + let out = finalize( + &tc, + MatchResult { + last_device_node_id: tc.arena.node(n3).id, + best_match_node_id: tc.arena.node(n3).id, + host_hit_length: 2, + ..tc.empty_match_result() + }, + ); + assert_eq!(out.host_hit_length, 2); +} + +#[test] +fn finalize_excludes_last_device_nodes_own_host_value() { + // root -> n1 -> n2 -> n3, host on all three; last_device=n2 counts only n3. + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 11, 111, 1111], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let n2 = tc + .arena + .alloc_child( + n1, + /* key = */ vec![2, 22], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let n3 = tc + .arena + .alloc_child( + n2, + /* key = */ vec![3, 33, 333], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(n1, FULL, Tensor::from_slice(&[0i64, 1, 2, 3])); + tc.arena + .set_host_value(n2, FULL, Tensor::from_slice(&[0i64, 1])); + tc.arena + .set_host_value(n3, FULL, Tensor::from_slice(&[0i64, 1, 2])); + let out = finalize( + &tc, + MatchResult { + last_device_node_id: tc.arena.node(n2).id, + best_match_node_id: tc.arena.node(n3).id, + host_hit_length: 0, + ..tc.empty_match_result() + }, + ); + assert_eq!(out.host_hit_length, 3); +} + +#[test] +fn finalize_skips_nodes_without_host_value() { + // root -> n1 -> n2 -> n3, host only on n1 and n3; a value-less n2 contributes 0. + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let n2 = tc + .arena + .alloc_child( + n1, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let n3 = tc + .arena + .alloc_child( + n2, + /* key = */ vec![3, 33, 333], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(n1, FULL, Tensor::from_slice(&[0i64])); + tc.arena + .set_host_value(n3, FULL, Tensor::from_slice(&[0i64, 1, 2])); + let out = finalize( + &tc, + MatchResult { + last_device_node_id: tc.arena.node(root).id, + best_match_node_id: tc.arena.node(n3).id, + host_hit_length: 0, + ..tc.empty_match_result() + }, + ); + assert_eq!(out.host_hit_length, 4); +} + +// Test-only split setup: root -> parent (key len 2) -> child, as _split_node leaves them. +fn nodes(tc: &mut UnifiedTreeCore>) -> (NodeIdx_, NodeIdx_) { + let root = tc.arena.root(); + let parent = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let child = tc + .arena + .alloc_child( + parent, + /* key = */ vec![3], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + (parent, child) +} + +#[test] +fn redistribute_copies_device_lock_ref_to_new_parent() { + let mut tc = core(); + let (parent, child) = nodes(&mut tc); + tc.arena + .node_mut(child) + .set_lock_ref_(ValueSlotIdx::device(FULL), 3); + FullComponent.redistribute_on_node_split(&mut tc, parent, child); + assert_eq!(tc.arena.device_lock_ref(parent, FULL), 3); + assert_eq!(tc.arena.device_lock_ref(child, FULL), 3); +} + +#[test] +fn redistribute_splits_device_value() { + let mut tc = core(); + let (parent, child) = nodes(&mut tc); + let original = Tensor::from_slice(&[10i64, 11, 12]); + // Mid-split state: the child's key is already trimmed, its value not yet split. + set_value_no_check(&mut tc, child, ValueSlotIdx::device(FULL), original.copy()); + FullComponent.redistribute_on_node_split(&mut tc, parent, child); + let parent_node = tc.arena.node(parent); + assert!( + parent_node + .device_value(FULL) + .equal(&original.narrow(0, 0, 2)) + ); + let child_node = tc.arena.node(child); + assert!( + child_node + .device_value(FULL) + .equal(&original.narrow(0, 2, 1)) + ); +} + +#[test] +fn redistribute_splits_host_value() { + let mut tc = core(); + let (parent, child) = nodes(&mut tc); + let original = Tensor::from_slice(&[20i64, 21, 22]); + set_value_no_check(&mut tc, child, ValueSlotIdx::host(FULL), original.copy()); + FullComponent.redistribute_on_node_split(&mut tc, parent, child); + let parent_node = tc.arena.node(parent); + assert!( + parent_node + .host_value(FULL) + .equal(&original.narrow(0, 0, 2)) + ); + let child_node = tc.arena.node(child); + assert!(child_node.host_value(FULL).equal(&original.narrow(0, 2, 1))); +} + +#[test] +fn redistribute_splits_device_value_with_bigram_key() { + let mut tc: UnifiedTreeCore> = + UnifiedTreeCore::new(CacheInitParams::default(), vec![FULL]); + let root = tc.arena.root(); + let parent = tc + .arena + .alloc_child( + root, + /* key = */ vec![(1, 2), (3, 4)], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let child = tc + .arena + .alloc_child( + parent, + /* key = */ vec![(5, 6)], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + // One value row per atom: a 2-atom bigram parent key takes 2 rows. + let original = Tensor::from_slice(&[10i64, 11, 12]); + // Mid-split state: the child's key is already trimmed, its value not yet split. + set_value_no_check(&mut tc, child, ValueSlotIdx::device(FULL), original.copy()); + FullComponent.redistribute_on_node_split(&mut tc, parent, child); + let parent_node = tc.arena.node(parent); + assert!( + parent_node + .device_value(FULL) + .equal(&original.narrow(0, 0, 2)) + ); + let child_node = tc.arena.node(child); + assert!( + child_node + .device_value(FULL) + .equal(&original.narrow(0, 2, 1)) + ); +} + +#[test] +fn redistribute_leaves_tombstoned_child_values_none() { + let mut tc = core(); + let (parent, child) = nodes(&mut tc); + FullComponent.redistribute_on_node_split(&mut tc, parent, child); + let parent_node = tc.arena.node(parent); + assert!(!parent_node.has_device_value(FULL)); + assert!(!parent_node.has_host_value(FULL)); + assert_eq!(parent_node.device_lock_ref(FULL), 0); + let child_node = tc.arena.node(child); + assert!(!child_node.has_device_value(FULL)); + assert!(!child_node.has_host_value(FULL)); +} + +#[test] +fn redistribute_device_only_child_leaves_host_none() { + let mut tc = core(); + let (parent, child) = nodes(&mut tc); + set_value_no_check( + &mut tc, + child, + ValueSlotIdx::device(FULL), + Tensor::from_slice(&[10i64, 11, 12]), + ); + FullComponent.redistribute_on_node_split(&mut tc, parent, child); + let parent_node = tc.arena.node(parent); + assert!(parent_node.has_device_value(FULL)); + assert!(!parent_node.has_host_value(FULL)); + let child_node = tc.arena.node(child); + assert!(!child_node.has_host_value(FULL)); +} + +#[test] +fn redistribute_halves_do_not_alias_source() { + let mut tc = core(); + let (parent, child) = nodes(&mut tc); + let mut source = Tensor::from_slice(&[10i64, 11, 12]); + let expected = source.copy(); + set_value_no_check( + &mut tc, + child, + ValueSlotIdx::device(FULL), + source.shallow_clone(), + ); + FullComponent.redistribute_on_node_split(&mut tc, parent, child); + // Writing through the source storage must not leak into either half. + let _ = source.fill_(99); + let parent_node = tc.arena.node(parent); + assert!( + parent_node + .device_value(FULL) + .equal(&expected.narrow(0, 0, 2)) + ); + let child_node = tc.arena.node(child); + assert!( + child_node + .device_value(FULL) + .equal(&expected.narrow(0, 2, 1)) + ); +} + +#[test] +fn redistribute_preserves_preexisting_parent_slot_value() { + let mut tc = core(); + let (parent, child) = nodes(&mut tc); + let sentinel = Tensor::from_slice(&[7i64, 8]); + tc.arena.set_device_value(parent, FULL, sentinel.copy()); + // Child has no device value: the parent's existing slot must be left untouched. + FullComponent.redistribute_on_node_split(&mut tc, parent, child); + let parent_node = tc.arena.node(parent); + assert!(parent_node.device_value(FULL).equal(&sentinel)); +} + +#[test] +fn redistribute_does_not_copy_host_lock_ref() { + let mut tc = core(); + let (parent, child) = nodes(&mut tc); + tc.arena + .node_mut(child) + .set_lock_ref_(ValueSlotIdx::host(FULL), 5); + FullComponent.redistribute_on_node_split(&mut tc, parent, child); + assert_eq!(tc.arena.host_lock_ref(parent, FULL), 0); +} + +#[test] +#[should_panic(expected = "out of bounds")] +fn redistribute_panics_on_missing_new_parent() { + let mut tc = core(); + let (_parent, child) = nodes(&mut tc); + FullComponent.redistribute_on_node_split(&mut tc, NodeIdx_(999), child); +} + +#[test] +fn evict_device_pushes_value_and_decrements_size() { + let mut tc = core(); + let root = tc.arena.root(); + let node = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2, 3], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let original = Tensor::from_slice(&[10i64, 11, 12]); + tc.arena.set_device_value(node, FULL, original.copy()); + tc.component_state_mut(FULL).evictable_size = 5; + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + let (freed, host_freed) = FullComponent.evict_component( + &mut tc, + node, + &mut device_frees, + &mut host_frees, + EvictLayer::Device, + ); + assert_eq!(freed, 3); + assert_eq!(host_freed, 0); + assert_eq!(tc.evictable_size_(FULL), 2); + let pushed = &device_frees[&FULL]; + assert_eq!(pushed.len(), 1); + assert!(pushed[0].equal(&original)); + // The device value is NOT tombstoned here (deferred to the cascade). + assert!(tc.arena.has_device_value(node, FULL)); +} + +#[test] +fn evict_device_on_valueless_node_is_noop() { + let mut tc = core(); + let root = tc.arena.root(); + let node = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.component_state_mut(FULL).evictable_size = 5; + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + let (freed, host_freed) = FullComponent.evict_component( + &mut tc, + node, + &mut device_frees, + &mut host_frees, + EvictLayer::Device, + ); + assert_eq!(freed, 0); + assert_eq!(host_freed, 0); + assert_eq!(tc.evictable_size_(FULL), 5); + assert!(!device_frees.contains_key(&FULL)); +} + +#[test] +fn evict_device_pushes_a_present_but_empty_value() { + let mut tc = core(); + let root = tc.arena.root(); + let node = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let empty: [i64; 0] = []; + set_value_no_check( + &mut tc, + node, + ValueSlotIdx::device(FULL), + Tensor::from_slice(&empty), + ); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + let (freed, _) = FullComponent.evict_component( + &mut tc, + node, + &mut device_frees, + &mut host_frees, + EvictLayer::Device, + ); + assert_eq!(freed, 0); + assert_eq!(device_frees[&FULL].len(), 1); +} + +#[test] +fn evict_host_pushes_host_value_and_tombstones() { + let mut tc = core(); + let root = tc.arena.root(); + let node = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let host_original = Tensor::from_slice(&[20i64, 21]); + tc.arena + .set_device_value(node, FULL, Tensor::from_slice(&[10i64, 11])); + tc.arena.set_host_value(node, FULL, host_original.copy()); + tc.component_state_mut(FULL).evictable_size = 5; + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + let (freed, host_freed) = FullComponent.evict_component( + &mut tc, + node, + &mut device_frees, + &mut host_frees, + EvictLayer::Host, + ); + assert_eq!(freed, 0); + assert_eq!(host_freed, 2); + let pushed = &host_frees[&FULL]; + assert_eq!(pushed.len(), 1); + assert!(pushed[0].equal(&host_original)); + // Host is tombstoned; device data and the evictable counter are untouched. + let node_ref = tc.arena.node(node); + assert!(!node_ref.has_host_value(FULL)); + assert!(node_ref.has_device_value(FULL)); + assert!(device_frees.is_empty()); + assert_eq!(tc.evictable_size_(FULL), 5); +} + +#[test] +fn evict_host_on_node_without_host_value_is_noop() { + let mut tc = core(); + let root = tc.arena.root(); + let node = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + let (freed, host_freed) = FullComponent.evict_component( + &mut tc, + node, + &mut device_frees, + &mut host_frees, + EvictLayer::Host, + ); + assert_eq!(freed, 0); + assert_eq!(host_freed, 0); + assert!(host_frees.is_empty()); +} + +#[test] +fn evict_host_pushes_a_present_but_empty_host_value() { + let mut tc = core(); + let root = tc.arena.root(); + let node = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let empty: [i64; 0] = []; + set_value_no_check( + &mut tc, + node, + ValueSlotIdx::host(FULL), + Tensor::from_slice(&empty), + ); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + let (_, host_freed) = FullComponent.evict_component( + &mut tc, + node, + &mut device_frees, + &mut host_frees, + EvictLayer::Host, + ); + assert_eq!(host_freed, 0); + assert_eq!(host_frees[&FULL].len(), 1); + assert!(!tc.arena.has_host_value(node, FULL)); +} + +#[test] +fn evict_host_accumulates_into_shared_host_frees() { + let mut tc = core(); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let b = tc + .arena + .alloc_child( + root, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let host_a = Tensor::from_slice(&[20i64, 21]); + let host_b = Tensor::from_slice(&[30i64]); + tc.arena.set_host_value(a, FULL, host_a.copy()); + tc.arena.set_host_value(b, FULL, host_b.copy()); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + let freed_a = FullComponent.evict_component( + &mut tc, + a, + &mut device_frees, + &mut host_frees, + EvictLayer::Host, + ); + let freed_b = FullComponent.evict_component( + &mut tc, + b, + &mut device_frees, + &mut host_frees, + EvictLayer::Host, + ); + assert_eq!(freed_a, (0, 2)); + assert_eq!(freed_b, (0, 1)); + let pushed = &host_frees[&FULL]; + assert_eq!(pushed.len(), 2); + assert!(pushed[0].equal(&host_a)); + assert!(pushed[1].equal(&host_b)); +} + +#[test] +fn evict_device_leaves_host_value_untouched() { + let mut tc = core(); + let root = tc.arena.root(); + let node = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(node, FULL, Tensor::from_slice(&[10i64])); + tc.arena + .set_host_value(node, FULL, Tensor::from_slice(&[20i64])); + tc.component_state_mut(FULL).evictable_size = 1; + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + let (freed, host_freed) = FullComponent.evict_component( + &mut tc, + node, + &mut device_frees, + &mut host_frees, + EvictLayer::Device, + ); + assert_eq!(freed, 1); + assert_eq!(host_freed, 0); + assert!(host_frees.is_empty()); + assert!(tc.arena.has_host_value(node, FULL)); +} + +#[test] +fn evict_all_frees_both_layers() { + let mut tc = core(); + let root = tc.arena.root(); + let node = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2, 3], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let device_original = Tensor::from_slice(&[10i64, 11, 12]); + let host_original = Tensor::from_slice(&[20i64, 21, 22]); + tc.arena + .set_device_value(node, FULL, device_original.copy()); + tc.arena.set_host_value(node, FULL, host_original.copy()); + tc.component_state_mut(FULL).evictable_size = 3; + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + let (freed, host_freed) = FullComponent.evict_component( + &mut tc, + node, + &mut device_frees, + &mut host_frees, + EvictLayer::All, + ); + assert_eq!(freed, 3); + assert_eq!(host_freed, 3); + assert!(device_frees[&FULL][0].equal(&device_original)); + assert!(host_frees[&FULL][0].equal(&host_original)); + // Device stays for the cascade to tombstone; host is tombstoned inline. + let node_ref = tc.arena.node(node); + assert!(node_ref.has_device_value(FULL)); + assert!(!node_ref.has_host_value(FULL)); + assert_eq!(tc.evictable_size_(FULL), 0); +} + +#[test] +fn evict_device_accumulates_into_shared_device_frees() { + let mut tc = core(); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let b = tc + .arena + .alloc_child( + root, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let value_a = Tensor::from_slice(&[10i64, 11]); + let value_b = Tensor::from_slice(&[20i64]); + tc.arena.set_device_value(a, FULL, value_a.copy()); + tc.arena.set_device_value(b, FULL, value_b.copy()); + tc.component_state_mut(FULL).evictable_size = 3; + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + FullComponent.evict_component( + &mut tc, + a, + &mut device_frees, + &mut host_frees, + EvictLayer::Device, + ); + FullComponent.evict_component( + &mut tc, + b, + &mut device_frees, + &mut host_frees, + EvictLayer::Device, + ); + let pushed = &device_frees[&FULL]; + assert_eq!(pushed.len(), 2); + assert!(pushed[0].equal(&value_a)); + assert!(pushed[1].equal(&value_b)); + assert_eq!(tc.evictable_size_(FULL), 0); +} + +#[test] +fn evict_device_pushed_entry_aliases_the_node_value() { + let mut tc = core(); + let root = tc.arena.root(); + let node = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2, 3], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(node, FULL, Tensor::from_slice(&[10i64, 11, 12])); + tc.component_state_mut(FULL).evictable_size = 3; + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + FullComponent.evict_component( + &mut tc, + node, + &mut device_frees, + &mut host_frees, + EvictLayer::Device, + ); + // The pushed entry shares storage with the node's value (Python appends cd.value). + let _ = device_frees.get_mut(&FULL).unwrap()[0].fill_(99); + assert!( + tc.arena + .device_value(node, FULL) + .equal(&Tensor::from_slice(&[99i64, 99, 99])) + ); +} + +#[test] +#[should_panic(expected = "out of bounds")] +fn evict_panics_on_missing_node() { + let mut tc = core(); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + FullComponent.evict_component( + &mut tc, + NodeIdx_(999), + &mut device_frees, + &mut host_frees, + EvictLayer::Device, + ); +} + +#[test] +#[should_panic(expected = "evictable size underflow")] +fn evict_panics_when_size_would_underflow() { + let mut tc = core(); + let root = tc.arena.root(); + let node = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2, 3], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(node, FULL, Tensor::from_slice(&[10i64, 11, 12])); + // Counter still 0: evicting 3 tokens must fail loudly, not wrap. + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + FullComponent.evict_component( + &mut tc, + node, + &mut device_frees, + &mut host_frees, + EvictLayer::Device, + ); +} + +#[test] +#[should_panic(expected = "EvictLayer::All")] +fn node_has_component_data_panics_on_all_layer() { + let mut tc = core(); + let root = tc.arena.root(); + let node = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + crate::components::node_has_component_data(&tc.arena, node, FULL, EvictLayer::All); +} + +#[test] +#[should_panic(expected = "out of bounds")] +fn node_has_component_data_panics_on_missing_node() { + let tc = core(); + crate::components::node_has_component_data(&tc.arena, NodeIdx_(999), FULL, EvictLayer::Device); +} + +// Chain root -> a (device-on, backuped) -> b (host-only) -> c (host-only). +fn load_back_chain(tc: &mut UnifiedTreeCore>) -> (NodeIdx_, NodeIdx_, NodeIdx_) { + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(a, FULL, Tensor::from_slice(&[10i64])); + tc.arena + .set_host_value(a, FULL, Tensor::from_slice(&[20i64])); + let b = tc + .arena + .alloc_child( + a, + /* key = */ vec![2, 3], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(b, FULL, Tensor::from_slice(&[21i64, 22])); + let c = tc + .arena + .alloc_child( + b, + /* key = */ vec![4], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(c, FULL, Tensor::from_slice(&[23i64])); + (a, b, c) +} + +#[test] +fn build_hicache_transfers_returns_none_for_non_load_back_phases() { + let mut tc = core(); + let (a, _b, _c) = load_back_chain(&mut tc); + for phase in [ + CacheTransferPhase::BackupHost, + CacheTransferPhase::BackupStorage, + CacheTransferPhase::Prefetch, + ] { + let transfers = FullComponent + .build_hicache_transfers( + &tc, a, phase, /* mamba_pool_idx = */ None, /* host_indices = */ None, + /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* last_hash = */ None, + ) + .unwrap(); + assert!(transfers.is_none()); + } +} + +#[test] +fn load_back_build_collects_the_evicted_suffix_ancestors_first() { + let mut tc = core(); + let (_a, b, c) = load_back_chain(&mut tc); + let transfers = FullComponent + .build_hicache_transfers( + &tc, + c, + CacheTransferPhase::LoadBack, + /* mamba_pool_idx = */ None, + /* host_indices = */ None, + /* token_ids = */ None, + /* prefetch_tokens = */ 0, + /* last_hash = */ None, + ) + .unwrap() + .unwrap(); + assert_eq!(transfers.len(), 1); + let xfer = &transfers[0]; + assert_eq!(xfer.name, PoolName::Kv); + // The walk stops at a's device value; b's host pages precede c's. + assert!( + xfer.host_indices + .as_ref() + .unwrap() + .equal(&Tensor::from_slice(&[21i64, 22, 23])) + ); + assert!(xfer.device_indices.is_none()); + assert_eq!( + xfer.nodes_to_load, + Some(vec![tc.arena.node(b).id, tc.arena.node(c).id]) + ); +} + +#[test] +fn load_back_build_returns_an_empty_cpu_transfer_for_a_device_backed_node() { + let mut tc = core(); + let (a, _b, _c) = load_back_chain(&mut tc); + let transfers = FullComponent + .build_hicache_transfers( + &tc, + a, + CacheTransferPhase::LoadBack, + /* mamba_pool_idx = */ None, + /* host_indices = */ None, + /* token_ids = */ None, + /* prefetch_tokens = */ 0, + /* last_hash = */ None, + ) + .unwrap() + .unwrap(); + let host_indices = transfers[0].host_indices.as_ref().unwrap(); + assert_eq!(host_indices.numel(), 0); + assert_eq!(host_indices.kind(), Kind::Int64); + assert_eq!(host_indices.device(), tch::Device::Cpu); + assert_eq!(transfers[0].nodes_to_load, Some(vec![])); +} + +#[test] +#[should_panic(expected = "value: Full/host slot has no value")] +fn load_back_build_panics_on_an_evicted_unbacked_node() { + let mut tc = core(); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let _ = FullComponent.build_hicache_transfers( + &tc, + a, + CacheTransferPhase::LoadBack, + /* mamba_pool_idx = */ None, + /* host_indices = */ None, + /* token_ids = */ None, + /* prefetch_tokens = */ 0, + /* last_hash = */ None, + ); +} + +#[test] +fn load_back_commit_attaches_device_slices_in_chain_order() { + let mut tc = core(); + let (_a, b, c) = load_back_chain(&mut tc); + let mut cache_actions = Vec::new(); + let transfer = PoolTransfer { + name: PoolName::Kv, + host_indices: Some(Tensor::from_slice(&[21i64, 22, 23])), + device_indices: Some(Tensor::from_slice(&[50i64, 51, 52])), + nodes_to_load: Some(vec![tc.arena.node(b).id, tc.arena.node(c).id]), + ..Default::default() + }; + FullComponent.commit_hicache_transfer( + &mut tc, + c, + CacheTransferPhase::LoadBack, + vec![transfer], + &mut cache_actions, + /* insert_result = */ None, + /* pool_storage_result = */ None, + ); + assert!( + tc.arena + .device_value(b, FULL) + .equal(&Tensor::from_slice(&[50i64, 51])) + ); + assert!( + tc.arena + .device_value(c, FULL) + .equal(&Tensor::from_slice(&[52i64])) + ); + assert_eq!(tc.full_evictable_size(), 3); + assert!(cache_actions.is_empty()); + assert!(tc.evictable_device_leaves.contains(c)); + // b entered the D-leaf set while c was still evicted and the per-node update + // never revisits it; the orchestrator's post-commit path re-lock corrects it. + assert!(tc.evictable_device_leaves.contains(b)); +} + +#[test] +fn load_back_commit_without_transfers_leaves_the_node_evicted() { + let mut tc = core(); + let (_a, _b, c) = load_back_chain(&mut tc); + let mut cache_actions = Vec::new(); + FullComponent.commit_hicache_transfer( + &mut tc, + c, + CacheTransferPhase::LoadBack, + vec![], + &mut cache_actions, + /* insert_result = */ None, + /* pool_storage_result = */ None, + ); + assert!(tc.arena.node(c).evicted()); + assert!(cache_actions.is_empty()); +} + +#[test] +fn load_back_commit_without_device_indices_leaves_the_node_evicted() { + let mut tc = core(); + let (_a, b, c) = load_back_chain(&mut tc); + let mut cache_actions = Vec::new(); + let transfer = PoolTransfer { + name: PoolName::Kv, + host_indices: Some(Tensor::from_slice(&[21i64, 22, 23])), + nodes_to_load: Some(vec![tc.arena.node(b).id, tc.arena.node(c).id]), + ..Default::default() + }; + FullComponent.commit_hicache_transfer( + &mut tc, + c, + CacheTransferPhase::LoadBack, + vec![transfer], + &mut cache_actions, + /* insert_result = */ None, + /* pool_storage_result = */ None, + ); + assert!(tc.arena.node(b).evicted()); + assert!(tc.arena.node(c).evicted()); + assert_eq!(tc.full_evictable_size(), 0); +} diff --git a/rust/mem-cache/src/tests/components/mamba.rs b/rust/mem-cache/src/tests/components/mamba.rs new file mode 100644 index 000000000..c749aa0d1 --- /dev/null +++ b/rust/mem-cache/src/tests/components/mamba.rs @@ -0,0 +1,1784 @@ +use super::*; +use crate::components::{FULL, MAMBA, SWA}; +use crate::test_utils::{accumulate_step, action_kinds}; +use crate::unified_lru_list::UnifiedLRUList; + +fn mamba_core(page_size: usize) -> UnifiedTreeCore> { + mamba_core_with_chunk(page_size, /* chunk = */ 256) +} + +fn mamba_core_with_cap(cap: usize) -> UnifiedTreeCore> { + UnifiedTreeCore::new( + CacheInitParams { + page_size: 1, + mamba_cache_chunk_size: Some(256), + mamba_max_states_per_path: Some(cap), + ..CacheInitParams::default() + }, + vec![FULL, MAMBA], + ) +} + +fn mamba_core_with_chunk(page_size: usize, chunk: usize) -> UnifiedTreeCore> { + UnifiedTreeCore::new( + CacheInitParams { + page_size, + mamba_cache_chunk_size: Some(chunk), + ..CacheInitParams::default() + }, + vec![FULL, MAMBA], + ) +} + +fn hybrid_lock_core() -> (UnifiedTreeCore>, NodeIdx_, NodeIdx_) { + let mut tc = UnifiedTreeCore::new( + CacheInitParams { + page_size: 1, + swa_sliding_window_size: Some(2), + mamba_cache_chunk_size: Some(256), + ..CacheInitParams::default() + }, + vec![FULL, SWA, MAMBA], + ); + let [parent, leaf] = chain::<2>(&mut tc); + for (node, full_slot, swa_slot, mamba_slot) in [(parent, 10, 20, 30), (leaf, 11, 21, 31)] { + tc.arena + .set_device_value(node, FULL, Tensor::from_slice(&[full_slot])); + tc.set_component_device_value(tc.arena.node(node).id, SWA, Tensor::from_slice(&[swa_slot])); + set_mamba_device(&mut tc, node, mamba_slot); + tc.update_evictable_leaf_sets_(node); + } + tc.component_state_mut(FULL).evictable_size = 2; + (tc, parent, leaf) +} + +fn insert_params_mamba<'k>( + key: &'k Vec, + value: &[i64], + mamba_slot: Option, +) -> InsertParams<'k, Vec> { + InsertParams { + key, + namespace: Default::default(), + value: Tensor::from_slice(value), + mamba_value: mamba_slot.map(|slot| Tensor::from_slice(&[slot])), + prev_prefix_len: 0, + swa_evicted_seqlen: 0, + chunked: false, + priority: 0, + track_adopted_ranges: false, + } +} + +fn match_params(key: &Vec) -> MatchPrefixParams<'_, Vec> { + MatchPrefixParams { + key, + namespace: Default::default(), + } +} + +// A two-node arena-built path: A[1,2] and B[3,4], both with FULL device +// values; mamba data is seeded by the caller. +fn two_node_path(tc: &mut UnifiedTreeCore>) -> (NodeIdx_, NodeIdx_) { + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(a, FULL, Tensor::from_slice(&[10i64, 11])); + let b = tc + .arena + .alloc_child( + a, + /* key = */ vec![3, 4], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(b, FULL, Tensor::from_slice(&[12i64, 13])); + (a, b) +} + +fn mamba_component() -> MambaComponent { + MambaComponent::new(&CacheInitParams { + mamba_cache_chunk_size: Some(256), + ..CacheInitParams::default() + }) +} + +// A chain of single-atom children under the default root; returns the node ids. +fn chain(tc: &mut UnifiedTreeCore>) -> [NodeIdx_; N] { + let mut parent = tc.arena.root(); + let mut nodes = [NodeIdx_(0); N]; + for (i, node) in nodes.iter_mut().enumerate() { + let id = tc + .arena + .alloc_child( + parent, + /* key = */ vec![i as i64 + 1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + *node = id; + parent = id; + } + nodes +} + +// Seed a one-slot mamba device value with its LRU and size bookkeeping. +fn set_mamba_device(tc: &mut UnifiedTreeCore>, node: NodeIdx_, slot: i64) { + tc.arena + .set_device_value(node, MAMBA, Tensor::from_slice(&[slot])); + tc.device_lru_list_mut(MAMBA).insert_mru(node); + tc.inc_evictable_size(MAMBA, 1); +} + +fn set_mamba_host(tc: &mut UnifiedTreeCore>, node: NodeIdx_, slot: i64) { + tc.arena + .set_host_value(node, MAMBA, Tensor::from_slice(&[slot])); +} + +fn set_full_host(tc: &mut UnifiedTreeCore>, node: NodeIdx_, slot: i64) { + tc.arena + .set_host_value(node, FULL, Tensor::from_slice(&[slot])); +} + +fn lru_order(lru: &UnifiedLRUList) -> Vec { + lru.iter().collect() +} + +#[test] +fn node_has_component_data_reads_each_layer() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + let mamba = mamba_component(); + assert!(!crate::components::node_has_component_data( + &tc.arena, + a, + MAMBA, + EvictLayer::Device + )); + assert!(!crate::components::node_has_component_data( + &tc.arena, + a, + MAMBA, + EvictLayer::Host + )); + set_mamba_device(&mut tc, a, 7); + set_mamba_host(&mut tc, a, 8); + assert!(crate::components::node_has_component_data( + &tc.arena, + a, + MAMBA, + EvictLayer::Device + )); + assert!(crate::components::node_has_component_data( + &tc.arena, + a, + MAMBA, + EvictLayer::Host + )); +} + +#[test] +fn refresh_lru_walkdown_is_a_no_op() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a, b] = chain::<2>(&mut tc); + set_mamba_device(&mut tc, a, 7); + set_mamba_device(&mut tc, b, 8); + let mamba = mamba_component(); + // The walk never reorders mamba states; commit and match own the stamps. + mamba.refresh_lru(&mut tc, LRURefreshPhase::Walkdown, a); + assert_eq!(lru_order(tc.device_lru_list(MAMBA)), vec![b, a]); +} + +#[test] +fn refresh_lru_match_end_touches_only_the_matched_node() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a, b, c] = chain::<3>(&mut tc); + set_mamba_device(&mut tc, c, 9); + set_mamba_device(&mut tc, b, 8); + set_mamba_device(&mut tc, a, 7); + let mamba = mamba_component(); + // Only the consumed state re-ranks; its valued ancestors stay put. + mamba.refresh_lru(&mut tc, LRURefreshPhase::MatchEnd, c); + assert_eq!(lru_order(tc.device_lru_list(MAMBA)), vec![c, a, b]); +} + +#[test] +fn refresh_lru_insert_end_is_a_noop() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a, b] = chain::<2>(&mut tc); + set_mamba_device(&mut tc, a, 7); + set_mamba_device(&mut tc, b, 8); + mamba_component().refresh_lru(&mut tc, LRURefreshPhase::InsertEnd, a); + assert_eq!(lru_order(tc.device_lru_list(MAMBA)), vec![b, a]); +} + +#[test] +fn device_value_round_trips_through_the_component() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + let mamba = mamba_component(); + assert!(tc.arena.try_device_value(a, MAMBA).is_none()); + tc.set_component_device_value(tc.arena.node(a).id, MAMBA, Tensor::from_slice(&[42i64])); + assert!( + tc.arena + .try_device_value(a, MAMBA) + .unwrap() + .equal(&Tensor::from_slice(&[42i64])) + ); +} + +#[test] +fn match_validator_accepts_device_and_optionally_host() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a, b, c] = chain::<3>(&mut tc); + set_mamba_device(&mut tc, a, 7); + set_mamba_host(&mut tc, b, 8); + let mamba = mamba_component(); + let mut hicache = mamba.create_match_validator(&tc, /* match_device_only = */ false); + assert!(hicache(&tc, a)); + assert!(hicache(&tc, b)); + assert!(!hicache(&tc, c)); + let mut device_only = mamba.create_match_validator(&tc, /* match_device_only = */ true); + assert!(device_only(&tc, a)); + assert!(!device_only(&tc, b)); +} + +#[test] +fn split_keeps_mamba_data_on_the_leaf() { + let mut tc = mamba_core(/* page_size = */ 1); + let root = tc.arena.root(); + let leaf = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(leaf, FULL, Tensor::from_slice(&[10i64, 11])); + tc.update_evictable_leaf_sets_(leaf); + set_mamba_device(&mut tc, leaf, 7); + let (new_parent, action) = tc.split_node_(leaf, /* split_len = */ 1); + assert!(action.is_none()); + assert!(tc.arena.node(leaf).has_device_value(MAMBA)); + assert!(!tc.arena.node(new_parent).has_device_value(MAMBA)); + assert!(!tc.arena.node(new_parent).has_host_value(MAMBA)); +} + +#[test] +fn device_lock_moves_the_slot_between_evictable_and_protected_once() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + set_mamba_device(&mut tc, a, 7); + let mamba = mamba_component(); + let result = mamba.acquire_component_lock( + &mut tc, + a, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + assert!(result.skip_lock_node_ids.is_empty()); + assert_eq!(tc.evictable_size_(MAMBA), 0); + assert_eq!(tc.protected_size_(MAMBA), 1); + mamba.acquire_component_lock( + &mut tc, + a, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + assert_eq!(tc.protected_size_(MAMBA), 1); + assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 2); + mamba.release_component_lock(&mut tc, a, None, /* lock_host = */ false); + assert_eq!(tc.protected_size_(MAMBA), 1); + mamba.release_component_lock(&mut tc, a, None, /* lock_host = */ false); + assert_eq!(tc.evictable_size_(MAMBA), 1); + assert_eq!(tc.protected_size_(MAMBA), 0); + assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 0); +} + +#[test] +fn skip_aware_lock_records_only_the_mamba_target() { + let (mut tc, parent, leaf) = hybrid_lock_core(); + let leaf_handle = tc.arena.node(leaf).id; + + let result = tc.inc_lock_ref_with_skip(leaf_handle, &[MAMBA]); + + assert_eq!(result.skip_lock_node_ids[&MAMBA].len(), 1); + assert!(result.skip_lock_node_ids[&MAMBA].contains(&leaf_handle)); + assert_eq!(tc.arena.node(parent).device_lock_ref(MAMBA), 0); + assert_eq!(tc.arena.node(leaf).device_lock_ref(MAMBA), 0); + assert_eq!(tc.evictable_size_(MAMBA), 2); + assert_eq!(tc.protected_size_(MAMBA), 0); + assert_eq!(tc.arena.node(parent).device_lock_ref(FULL), 1); + assert_eq!(tc.arena.node(leaf).device_lock_ref(FULL), 1); + + tc.dec_lock_ref( + leaf_handle, + Some(&DecLockRefParams { + swa_uuid_for_lock: result.swa_uuid_for_lock, + skip_lock_node_ids: result.skip_lock_node_ids, + ..Default::default() + }), + /* skip_swa = */ false, + ); + assert_eq!(tc.arena.node(parent).device_lock_ref(FULL), 0); + assert_eq!(tc.arena.node(leaf).device_lock_ref(FULL), 0); +} + +#[test] +fn swa_only_release_honors_a_skipped_mamba_target() { + let (mut tc, _parent, leaf) = hybrid_lock_core(); + let leaf_handle = tc.arena.node(leaf).id; + let owner = tc.inc_lock_ref(leaf_handle); + let skipped = tc.inc_lock_ref_with_skip(leaf_handle, &[MAMBA]); + assert_eq!(tc.arena.node(leaf).device_lock_ref(MAMBA), 1); + + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + tc.dec_swa_lock_only_with_skip( + leaf_handle, + skipped.swa_uuid_for_lock, + Some(&skipped.skip_lock_node_ids), + &mut device_frees, + &mut host_frees, + ); + + assert!(device_frees.is_empty()); + assert!(host_frees.is_empty()); + assert_eq!(tc.arena.node(leaf).device_lock_ref(MAMBA), 1); + assert_eq!(tc.protected_size_(MAMBA), 1); + + let skipped_params = DecLockRefParams { + swa_uuid_for_lock: skipped.swa_uuid_for_lock, + skip_lock_node_ids: skipped.skip_lock_node_ids, + ..Default::default() + }; + tc.dec_lock_ref( + leaf_handle, + Some(&skipped_params), + /* skip_swa = */ true, + ); + let owner_params = DecLockRefParams { + swa_uuid_for_lock: owner.swa_uuid_for_lock, + skip_lock_node_ids: owner.skip_lock_node_ids, + ..Default::default() + }; + tc.dec_lock_ref( + leaf_handle, + Some(&owner_params), + /* skip_swa = */ false, + ); + assert_eq!(tc.protected_size_(MAMBA), 0); +} + +#[test] +fn tombstone_lock_is_recorded_and_replayed_at_release() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + let mamba = mamba_component(); + let result = mamba.acquire_component_lock( + &mut tc, + a, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + assert!(result.skip_lock_node_ids[&MAMBA].contains(&tc.arena.node(a).id)); + assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 0); + // The replayed skip set keeps the release from touching the node. + let params = DecLockRefParams { + skip_lock_node_ids: result.skip_lock_node_ids.clone(), + ..DecLockRefParams::default() + }; + mamba.release_component_lock(&mut tc, a, Some(¶ms), /* lock_host = */ false); + assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 0); + assert_eq!(tc.evictable_size_(MAMBA), 0); +} + +#[test] +fn root_locks_are_noops() { + let mut tc = mamba_core(/* page_size = */ 1); + let root = tc.arena.root(); + let mamba = mamba_component(); + let result = mamba.acquire_component_lock( + &mut tc, + root, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + assert!(result.skip_lock_node_ids.is_empty()); + mamba.release_component_lock(&mut tc, root, None, /* lock_host = */ false); + assert_eq!(tc.evictable_size_(MAMBA), 0); +} + +#[test] +fn host_lock_detaches_and_reattaches_the_host_lru() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + set_mamba_host(&mut tc, a, 8); + tc.host_lru_list_mut(MAMBA).insert_mru(a); + let mamba = mamba_component(); + mamba.acquire_component_lock( + &mut tc, + a, + IncLockRefResult::default(), + /* lock_host = */ true, + ); + assert!(!tc.host_lru_list(MAMBA).in_list(Some(a))); + assert_eq!(tc.arena.node(a).host_lock_ref(MAMBA), 1); + mamba.release_component_lock(&mut tc, a, None, /* lock_host = */ true); + assert!(tc.host_lru_list(MAMBA).in_list(Some(a))); + assert_eq!(tc.arena.node(a).host_lock_ref(MAMBA), 0); +} + +#[test] +fn host_unlock_skips_the_lru_for_device_backed_nodes() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + set_mamba_host(&mut tc, a, 8); + set_mamba_device(&mut tc, a, 7); + let mamba = mamba_component(); + mamba.acquire_component_lock( + &mut tc, + a, + IncLockRefResult::default(), + /* lock_host = */ true, + ); + mamba.release_component_lock(&mut tc, a, None, /* lock_host = */ true); + assert!(!tc.host_lru_list(MAMBA).in_list(Some(a))); +} + +#[test] +fn eviction_priority_is_the_lowest_tier_everywhere() { + let mamba = mamba_component(); + assert_eq!( + TreeComponent::>::eviction_priority(&mamba, /* is_leaf = */ true), + 0 + ); + assert_eq!( + TreeComponent::>::eviction_priority(&mamba, /* is_leaf = */ false), + 0 + ); +} + +#[test] +fn insert_attaches_the_donated_slot_to_the_new_leaf() { + let mut tc = mamba_core(/* page_size = */ 1); + let result = tc.insert(&insert_params_mamba(&vec![1, 2], &[10, 11], Some(7))); + assert!(!result.mamba_exist); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + assert!( + tc.arena + .node(tc.arena.resolve(leaf)) + .try_device_value(MAMBA) + .unwrap() + .equal(&Tensor::from_slice(&[7i64])) + ); + assert!( + tc.device_lru_list(MAMBA) + .in_list(Some(tc.arena.resolve(leaf))) + ); + assert_eq!(tc.evictable_size_(MAMBA), 1); +} + +#[test] +fn reinsert_keeps_the_existing_slot_and_flags_the_caller() { + let mut tc = mamba_core(/* page_size = */ 1); + tc.insert(&insert_params_mamba(&vec![1, 2], &[10, 11], Some(7))); + let result = tc.insert(&insert_params_mamba(&vec![1, 2], &[10, 11], Some(8))); + assert!(result.mamba_exist); + assert_eq!(tc.evictable_size_(MAMBA), 1); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + // The original slot stays; the caller frees the unused donated one. + assert!( + tc.arena + .node(tc.arena.resolve(leaf)) + .try_device_value(MAMBA) + .unwrap() + .equal(&Tensor::from_slice(&[7i64])) + ); +} + +#[test] +fn reinsert_full_backed_target_schedules_mamba_only_backup() { + let mut tc = mamba_core(/* page_size = */ 1); + tc.set_hicache_enabled(); + let key = vec![1, 2]; + tc.insert(&insert_params_mamba(&key, &[10, 11], Some(7))); + let leaf = tc.match_prefix(&match_params(&key)).best_match_node_id; + tc.commit_backup(leaf, Tensor::from_slice(&[100i64, 101]), HashMap::new()); + + let result = tc.insert(&insert_params_mamba(&key, &[20, 21], Some(8))); + let backups = result + .cache_actions + .iter() + .filter_map(|action| match action { + CacheAction::BackupKV(backup) => Some(backup), + _ => None, + }) + .collect::>(); + assert_eq!(backups.len(), 1); + assert_eq!(backups[0].node_ids, vec![leaf]); + + let (full_device_indices, comp_xfers) = tc.build_backup_spec(leaf); + assert_eq!(full_device_indices.numel(), 0); + let mamba_xfers = &comp_xfers[&MAMBA]; + assert_eq!(mamba_xfers.len(), 1); + assert!( + mamba_xfers[0] + .device_indices + .as_ref() + .unwrap() + .equal(&Tensor::from_slice(&[7i64])) + ); + + tc.mark_write_through_pending(leaf); + let pending = tc.insert(&insert_params_mamba(&key, &[30, 31], Some(9))); + assert!( + !pending + .cache_actions + .iter() + .any(|action| matches!(action, CacheAction::BackupKV(_))) + ); +} + +#[test] +fn tombstone_refill_moves_the_node_from_host_to_device_lru() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + set_mamba_host(&mut tc, a, 8); + tc.host_lru_list_mut(MAMBA).insert_mru(a); + let mut result = InsertResult::default(); + mamba_component().commit_insert_component_data( + &mut tc, + a, + /* is_new_leaf = */ false, + &insert_params_mamba(&vec![1], &[10], Some(9)), + &mut result, + &mut Vec::new(), + ); + assert!(!result.mamba_exist); + assert!( + tc.arena + .node(a) + .try_device_value(MAMBA) + .unwrap() + .equal(&Tensor::from_slice(&[9i64])) + ); + assert!(!tc.host_lru_list(MAMBA).in_list(Some(a))); + assert!(tc.device_lru_list(MAMBA).in_list(Some(a))); + assert_eq!(tc.evictable_size_(MAMBA), 1); +} + +#[test] +#[should_panic(expected = "requires a donated mamba_value")] +fn insert_without_a_mamba_value_panics() { + let mut tc = mamba_core(/* page_size = */ 1); + tc.insert(&insert_params_mamba(&vec![1], &[10], None)); +} + +#[test] +fn match_reports_the_chunk_aligned_branching_seqlen() { + let mut tc = mamba_core_with_chunk(/* page_size = */ 1, /* chunk = */ 3); + let (a, _b) = two_node_path(&mut tc); + set_mamba_device(&mut tc, a, 7); + let result = tc.match_prefix(&match_params(&vec![1, 2, 3, 4])); + // The walk covers 4 tokens past the mamba anchor; 4 aligns down to 3. + assert_eq!(result.best_match_node_id, tc.arena.node(a).id); + assert_eq!(result.mamba_branching_seqlen, Some(3)); + assert_eq!(result.mamba_host_hit_length, 0); +} + +#[test] +fn branching_seqlen_uses_the_joint_chunk_and_tree_page_grid() { + let mut tc = mamba_core_with_chunk(/* page_size = */ 3, /* chunk = */ 2); + let root = tc.arena.root(); + let node = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2, 3, 4, 5, 6, 7, 8, 9], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena.set_device_value( + node, + FULL, + Tensor::from_slice(&[10i64, 11, 12, 13, 14, 15, 16, 17, 18]), + ); + + let result = tc.match_prefix(&match_params(&vec![1, 2, 3, 4, 5, 6, 7, 8, 9])); + + assert_eq!(result.full_kv_hit_length, 9); + // lcm(chunk=2, page=3) is 6; chunk-only alignment would incorrectly yield 8. + assert_eq!(result.mamba_branching_seqlen, Some(6)); +} + +#[test] +fn short_walks_have_no_branching_seqlen() { + let mut tc = mamba_core_with_chunk(/* page_size = */ 1, /* chunk = */ 8); + let (a, _b) = two_node_path(&mut tc); + set_mamba_device(&mut tc, a, 7); + let result = tc.match_prefix(&match_params(&vec![1, 2, 3, 4])); + // 4 walked tokens align down to zero at chunk 8. + assert_eq!(result.mamba_branching_seqlen, None); +} + +#[test] +fn matches_ending_on_the_mamba_anchor_have_no_branching_seqlen() { + let mut tc = mamba_core_with_chunk(/* page_size = */ 1, /* chunk = */ 2); + let (a, b) = two_node_path(&mut tc); + set_mamba_device(&mut tc, a, 7); + set_mamba_device(&mut tc, b, 9); + let result = tc.match_prefix(&match_params(&vec![1, 2, 3, 4])); + assert_eq!(result.best_match_node_id, tc.arena.node(b).id); + assert_eq!(result.mamba_branching_seqlen, None); +} + +#[test] +fn hicache_branching_seqlen_uses_the_full_kv_hit() { + let mut tc = mamba_core_with_chunk(/* page_size = */ 1, /* chunk = */ 3); + tc.set_hicache_enabled(); + let (a, _b) = two_node_path(&mut tc); + set_mamba_device(&mut tc, a, 7); + let result = tc.match_prefix(&match_params(&vec![1, 2, 3, 4])); + // The full walk hit 4 tokens; chunk-3 alignment lands past the 2-token + // mamba boundary, so the branch point fills even under HiCache. + assert_eq!(result.full_kv_hit_length, 4); + assert_eq!(result.mamba_branching_seqlen, Some(3)); +} + +#[test] +fn host_only_anchor_bumps_the_mamba_host_hit() { + let mut tc = mamba_core(/* page_size = */ 1); + tc.set_hicache_enabled(); + let (a, _b) = two_node_path(&mut tc); + set_mamba_host(&mut tc, a, 8); + let result = tc.match_prefix(&match_params(&vec![1, 2])); + assert_eq!(result.best_match_node_id, tc.arena.node(a).id); + assert_eq!(result.mamba_host_hit_length, 1); +} + +#[test] +fn evict_component_device_frees_and_tombstones_the_slot() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + set_mamba_device(&mut tc, a, 7); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + let (freed, host_freed) = mamba_component().evict_component( + &mut tc, + a, + &mut device_frees, + &mut host_frees, + EvictLayer::Device, + ); + assert_eq!((freed, host_freed), (1, 0)); + assert!(!tc.arena.node(a).has_device_value(MAMBA)); + assert_eq!(tc.evictable_size_(MAMBA), 0); + assert!(device_frees[&MAMBA][0].equal(&Tensor::from_slice(&[7i64]))); + assert!(host_frees.is_empty()); + // No host backup: the node does not enter the host LRU. + assert!(!tc.host_lru_list(MAMBA).in_list(Some(a))); +} + +#[test] +fn evict_component_device_moves_a_host_backed_node_into_the_host_lru() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + set_mamba_device(&mut tc, a, 7); + set_mamba_host(&mut tc, a, 8); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + mamba_component().evict_component( + &mut tc, + a, + &mut device_frees, + &mut host_frees, + EvictLayer::Device, + ); + assert!(tc.arena.node(a).has_host_value(MAMBA)); + assert!(tc.host_lru_list(MAMBA).in_list(Some(a))); +} + +#[test] +fn evict_component_host_frees_and_leaves_the_host_lru() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + set_mamba_host(&mut tc, a, 8); + tc.host_lru_list_mut(MAMBA).insert_mru(a); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + let (freed, host_freed) = mamba_component().evict_component( + &mut tc, + a, + &mut device_frees, + &mut host_frees, + EvictLayer::Host, + ); + assert_eq!((freed, host_freed), (0, 1)); + assert!(!tc.arena.node(a).has_host_value(MAMBA)); + assert!(!tc.host_lru_list(MAMBA).in_list(Some(a))); + assert!(host_frees[&MAMBA][0].equal(&Tensor::from_slice(&[8i64]))); +} + +#[test] +fn evict_component_all_frees_both_tiers_without_the_host_lru() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + set_mamba_device(&mut tc, a, 7); + set_mamba_host(&mut tc, a, 8); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + let (freed, host_freed) = mamba_component().evict_component( + &mut tc, + a, + &mut device_frees, + &mut host_frees, + EvictLayer::All, + ); + assert_eq!((freed, host_freed), (1, 1)); + assert!(!tc.arena.node(a).has_device_value(MAMBA)); + assert!(!tc.arena.node(a).has_host_value(MAMBA)); + // ALL means the node dies: it must not re-enter the host LRU. + assert!(!tc.host_lru_list(MAMBA).in_list(Some(a))); +} + +#[test] +fn device_walk_advances_one_allocator_mutation_per_call() { + let mut tc = mamba_core(/* page_size = */ 1); + tc.insert(&insert_params_mamba(&vec![1], &[10], Some(7))); + tc.insert(&insert_params_mamba(&vec![1, 2], &[10, 11], Some(8))); + let a = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; + let b = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + let mut tracker = HashMap::from([(MAMBA, 0)]); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + tc.evict_device_start(MAMBA, /* request_cnt = */ 2); + let (first, step) = tc.evict_device_next_node(MAMBA, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + // The internal node is a complete step so its free can be reused before + // the walk hands out another victim. + assert_eq!(first, None); + assert!(!tc.arena.node(tc.arena.resolve(a)).has_device_value(MAMBA)); + assert!(tc.arena.node(tc.arena.resolve(b)).has_device_value(MAMBA)); + assert!(tc.arena.has_device_value(tc.arena.resolve(a), FULL)); + assert_eq!(tracker[&MAMBA], 1); + assert!(!tc.device_lru_list(MAMBA).in_list(Some(tc.arena.resolve(a)))); + + let (second, step) = tc.evict_device_next_node(MAMBA, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + assert_eq!(second, Some(b)); + assert_eq!(tracker[&MAMBA], 1); + tc.evict_device_end(MAMBA); +} + +#[test] +fn device_walk_skips_locked_nodes() { + let mut tc = mamba_core(/* page_size = */ 1); + tc.insert(&insert_params_mamba(&vec![1], &[10], Some(7))); + tc.insert(&insert_params_mamba(&vec![1, 2], &[10, 11], Some(8))); + let a = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; + let b = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + let a_idx = tc.arena.resolve(a); + mamba_component().acquire_component_lock( + &mut tc, + a_idx, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + let mut tracker = HashMap::from([(MAMBA, 0)]); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + tc.evict_device_start(MAMBA, /* request_cnt = */ 2); + let (next, step) = tc.evict_device_next_node(MAMBA, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + // The locked internal node stays; the cursor starts on the leaf. + assert_eq!(next, Some(b)); + assert!(tc.arena.node(tc.arena.resolve(a)).has_device_value(MAMBA)); + assert_eq!(tracker[&MAMBA], 0); + tc.evict_device_end(MAMBA); +} + +#[test] +#[should_panic(expected = "Mamba device eviction not started")] +fn device_walk_requires_a_start() { + let mut tc = mamba_core(/* page_size = */ 1); + let tracker = HashMap::from([(MAMBA, 0)]); + tc.evict_device_next_node(MAMBA, &tracker); +} + +#[test] +#[should_panic(expected = "Mamba eviction cursor on a valueless node")] +fn device_walk_asserts_a_valued_cursor_node() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + set_mamba_device(&mut tc, a, 7); + // a stays in the Mamba LRU but loses its device value out of band. + let _ = tc.arena.node_mut(a).take_device_value(MAMBA); + let tracker = HashMap::from([(MAMBA, 0)]); + tc.evict_device_start(MAMBA, /* request_cnt = */ 100); + tc.evict_device_next_node(MAMBA, &tracker); +} + +#[test] +fn host_eviction_tombstones_internal_host_values() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a, _b] = chain::<2>(&mut tc); + set_mamba_host(&mut tc, a, 8); + tc.host_lru_list_mut(MAMBA).insert_mru(a); + let mut tracker = HashMap::from([(MAMBA, 0)]); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + mamba_component().drive_host_eviction( + &mut tc, + /* num_tokens = */ 1, + &mut tracker, + &mut device_frees, + &mut host_frees, + ); + assert!(!tc.arena.node(a).has_host_value(MAMBA)); + assert!(!tc.host_lru_list(MAMBA).in_list(Some(a))); + assert_eq!(tracker[&MAMBA], 1); + assert!(host_frees[&MAMBA][0].equal(&Tensor::from_slice(&[8i64]))); +} + +#[test] +fn host_eviction_takes_a_host_leaf_atomically() { + let mut tc = mamba_core(/* page_size = */ 1); + tc.set_hicache_enabled(); + let root = tc.arena.root(); + let leaf = tc + .insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1], + Tensor::from_slice(&[100i64]), + vec!["h0".to_string()], + ) + .inserted_host_node + .unwrap(); + let leaf_idx = tc.arena.resolve(leaf); + set_mamba_host(&mut tc, leaf_idx, 8); + tc.host_lru_list_mut(MAMBA).insert_mru(leaf_idx); + assert!(tc.evictable_host_leaves.contains(tc.arena.resolve(leaf))); + let mut tracker = HashMap::from([(MAMBA, 0)]); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + mamba_component().drive_host_eviction( + &mut tc, + /* num_tokens = */ 1, + &mut tracker, + &mut device_frees, + &mut host_frees, + ); + // The atomic host-leaf eviction frees both components and the node. + assert_eq!(tracker[&MAMBA], 1); + assert!(host_frees[&MAMBA][0].equal(&Tensor::from_slice(&[8i64]))); + assert!(host_frees[&FULL][0].equal(&Tensor::from_slice(&[100i64]))); + assert!(tc.arena.try_resolve(leaf).is_none()); + assert!(!tc.host_lru_list(MAMBA).in_list(Some(leaf_idx))); +} + +#[test] +#[should_panic(expected = "has no host value")] +fn host_drive_panics_on_an_lru_member_without_a_mamba_host_value() { + let mut tc = mamba_core(/* page_size = */ 1); + let [n] = chain::<1>(&mut tc); + tc.arena + .set_host_value(n, FULL, Tensor::from_slice(&[100i64])); + tc.host_lru_list_mut(MAMBA).insert_mru(n); + let mut tracker = HashMap::from([(MAMBA, 0)]); + mamba_component().drive_host_eviction( + &mut tc, + /* num_tokens = */ 100, + &mut tracker, + &mut HashMap::new(), + &mut HashMap::new(), + ); +} + +#[test] +fn swa_triggered_cascade_takes_the_mamba_slot() { + let mut tc = UnifiedTreeCore::>::new( + CacheInitParams { + page_size: 1, + swa_sliding_window_size: Some(4), + mamba_cache_chunk_size: Some(256), + ..CacheInitParams::default() + }, + vec![FULL, SWA, MAMBA], + ); + let [a] = chain::<1>(&mut tc); + set_mamba_device(&mut tc, a, 7); + let mut tracker = HashMap::from([(SWA, 0)]); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + // The SWA internal tier (1) outranks mamba (0): the cascade takes it. + tc.cascade_evict_( + a, + SWA, + &mut tracker, + &mut device_frees, + &mut host_frees, + EvictLayer::Device, + ); + assert!(!tc.arena.node(a).has_device_value(MAMBA)); + assert_eq!(tracker[&MAMBA], 1); +} + +#[test] +fn mamba_triggered_cascade_spares_the_higher_tiers() { + let mut tc = UnifiedTreeCore::>::new( + CacheInitParams { + page_size: 1, + swa_sliding_window_size: Some(4), + mamba_cache_chunk_size: Some(256), + ..CacheInitParams::default() + }, + vec![FULL, SWA, MAMBA], + ); + let [a] = chain::<1>(&mut tc); + tc.arena + .set_device_value(a, FULL, Tensor::from_slice(&[10i64])); + tc.arena + .set_device_value(a, SWA, Tensor::from_slice(&[20i64])); + let mut tracker = HashMap::from([(MAMBA, 0)]); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + tc.cascade_evict_( + a, + MAMBA, + &mut tracker, + &mut device_frees, + &mut host_frees, + EvictLayer::Device, + ); + // FULL (2) and SWA (1) both outrank the mamba trigger (0). + assert!(tc.arena.has_device_value(a, FULL)); + assert!(tc.arena.has_device_value(a, SWA)); + assert!(device_frees.is_empty()); +} + +#[test] +fn backup_host_build_carries_the_device_slot() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + set_mamba_device(&mut tc, a, 7); + let transfers = tc + .build_hicache_transfers( + MAMBA, + tc.arena.node(a).id, + CacheTransferPhase::BackupHost, + None, + None, + 0, + None, + ) + .unwrap(); + assert_eq!(transfers.len(), 1); + assert_eq!(transfers[0].name, PoolName::Mamba); + assert!( + transfers[0] + .device_indices + .as_ref() + .unwrap() + .equal(&Tensor::from_slice(&[7i64])) + ); + // A tombstone has nothing to back up. + let [b] = [tc + .arena + .alloc_child( + a, + /* key = */ vec![9], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap()]; + assert!( + tc.build_hicache_transfers( + MAMBA, + tc.arena.node(b).id, + CacheTransferPhase::BackupHost, + None, + None, + 0, + None, + ) + .is_none() + ); +} + +#[test] +fn load_back_build_restores_the_host_only_node() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + set_mamba_host(&mut tc, a, 8); + let transfers = tc + .build_hicache_transfers( + MAMBA, + tc.arena.node(a).id, + CacheTransferPhase::LoadBack, + None, + None, + 0, + None, + ) + .unwrap(); + assert_eq!(transfers.len(), 1); + assert!( + transfers[0] + .host_indices + .as_ref() + .unwrap() + .equal(&Tensor::from_slice(&[8i64])) + ); + assert_eq!(transfers[0].nodes_to_load, Some(vec![tc.arena.node(a).id])); +} + +#[test] +fn load_back_build_skips_device_backed_and_bare_nodes() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a, b] = chain::<2>(&mut tc); + set_mamba_device(&mut tc, a, 7); + for node in [a, b] { + assert!( + tc.build_hicache_transfers( + MAMBA, + tc.arena.node(node).id, + CacheTransferPhase::LoadBack, + None, + None, + 0, + None, + ) + .is_none() + ); + } +} + +#[test] +fn load_back_build_adds_the_per_request_cow_transfer() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + set_mamba_host(&mut tc, a, 8); + let transfers = mamba_component() + .build_hicache_transfers( + &tc, + a, + CacheTransferPhase::LoadBack, + /* mamba_pool_idx = */ Some(Tensor::from_slice(&[3i64]).squeeze()), + None, + None, + 0, + None, + ) + .unwrap() + .unwrap(); + assert_eq!(transfers.len(), 2); + // The CoW transfer copies the host slot into the request's device slot. + assert!( + transfers[1] + .device_indices + .as_ref() + .unwrap() + .equal(&Tensor::from_slice(&[3i64])) + ); + assert!(transfers[1].nodes_to_load.is_none()); +} + +#[test] +fn backup_host_commit_stores_the_host_slot_once() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + let mut cache_actions = Vec::new(); + tc.commit_hicache_transfers( + tc.arena.node(a).id, + CacheTransferPhase::BackupHost, + HashMap::from([( + MAMBA, + vec![PoolTransfer { + name: PoolName::Mamba, + host_indices: Some(Tensor::from_slice(&[30i64])), + ..Default::default() + }], + )]), + &mut cache_actions, + None, + None, + ); + assert!( + tc.arena + .node(a) + .try_host_value(MAMBA) + .unwrap() + .equal(&Tensor::from_slice(&[30i64])) + ); + // A second backup keeps the existing host slot. + tc.commit_hicache_transfers( + tc.arena.node(a).id, + CacheTransferPhase::BackupHost, + HashMap::from([( + MAMBA, + vec![PoolTransfer { + name: PoolName::Mamba, + host_indices: Some(Tensor::from_slice(&[31i64])), + ..Default::default() + }], + )]), + &mut cache_actions, + None, + None, + ); + assert!( + tc.arena + .node(a) + .try_host_value(MAMBA) + .unwrap() + .equal(&Tensor::from_slice(&[30i64])) + ); + assert!(cache_actions.is_empty()); +} + +#[test] +fn load_back_commit_moves_the_node_onto_the_device_tier() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + set_mamba_host(&mut tc, a, 8); + tc.host_lru_list_mut(MAMBA).insert_mru(a); + let mut cache_actions = Vec::new(); + tc.commit_hicache_transfers( + tc.arena.node(a).id, + CacheTransferPhase::LoadBack, + HashMap::from([( + MAMBA, + vec![PoolTransfer { + name: PoolName::Mamba, + host_indices: Some(Tensor::from_slice(&[8i64])), + device_indices: Some(Tensor::from_slice(&[40i64])), + ..Default::default() + }], + )]), + &mut cache_actions, + None, + None, + ); + assert!( + tc.arena + .node(a) + .try_device_value(MAMBA) + .unwrap() + .equal(&Tensor::from_slice(&[40i64])) + ); + assert!(!tc.host_lru_list(MAMBA).in_list(Some(a))); + assert!(tc.device_lru_list(MAMBA).in_list(Some(a))); + assert_eq!(tc.evictable_size_(MAMBA), 1); + assert!(cache_actions.is_empty()); +} + +#[test] +fn mamba_device_eviction_skips_a_load_back_pinned_node() { + let mut tc = mamba_core(/* page_size = */ 1); + tc.is_write_back = true; + let [n] = chain::<1>(&mut tc); + set_full_host(&mut tc, n, 10); + set_mamba_host(&mut tc, n, 20); + let (kv_xfer, mut comp_xfers) = + tc.build_load_back_spec(tc.arena.node(n).id, /* req = */ None); + comp_xfers.get_mut(&MAMBA).unwrap()[0].device_indices = Some(Tensor::from_slice(&[40i64])); + tc.commit_load_back( + tc.arena.node(n).id, + Tensor::from_slice(&[30i64]), + kv_xfer, + comp_xfers, + ); + + tc.evict_device_start(MAMBA, /* request_cnt = */ 1); + let (next, _) = tc.evict_device_next_node(MAMBA, &HashMap::new()); + assert_eq!(next, None); + tc.evict_device_end(MAMBA); + assert!(tc.arena.has_device_value(n, MAMBA)); + + tc.finish_load_back(tc.arena.node(n).id); + tc.evict_device_start(MAMBA, /* request_cnt = */ 1); + let (next, _) = tc.evict_device_next_node(MAMBA, &HashMap::new()); + assert_eq!(next, Some(tc.arena.node(n).id)); + tc.evict_device_end(MAMBA); +} + +#[test] +fn mamba_host_eviction_skips_a_load_back_pinned_node() { + let mut tc = mamba_core(/* page_size = */ 1); + tc.is_write_back = true; + let [a, b] = chain::<2>(&mut tc); + set_full_host(&mut tc, a, 10); + set_full_host(&mut tc, b, 11); + set_mamba_host(&mut tc, a, 20); + tc.host_lru_list_mut(MAMBA).insert_mru(a); + let (kv_xfer, comp_xfers) = tc.build_load_back_spec(tc.arena.node(b).id, /* req = */ None); + assert!(comp_xfers.is_empty()); + tc.commit_load_back( + tc.arena.node(b).id, + Tensor::from_slice(&[30i64, 31]), + kv_xfer, + comp_xfers, + ); + + let result = tc.drive_host_eviction(MAMBA, /* num_tokens = */ 1); + assert_eq!(result.tracker[&MAMBA], 0); + assert!(result.host_frees.is_empty()); + assert!(tc.arena.has_host_value(a, MAMBA)); + + tc.finish_load_back(tc.arena.node(b).id); + let result = tc.drive_host_eviction(MAMBA, /* num_tokens = */ 1); + assert_eq!(result.tracker[&MAMBA], 1); + assert_eq!(result.host_frees[&MAMBA].len(), 1); + assert!(!tc.arena.has_host_value(a, MAMBA)); + tc.sanity_check(&[], &[]); +} + +#[test] +fn backup_storage_commit_is_a_noop() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + set_mamba_host(&mut tc, a, 8); + let mut cache_actions = Vec::new(); + tc.commit_hicache_transfers( + tc.arena.node(a).id, + CacheTransferPhase::BackupStorage, + HashMap::from([( + MAMBA, + vec![PoolTransfer { + name: PoolName::Mamba, + host_indices: Some(Tensor::from_slice(&[8i64])), + ..Default::default() + }], + )]), + &mut cache_actions, + None, + None, + ); + assert!(tc.arena.node(a).has_host_value(MAMBA)); + assert!(cache_actions.is_empty()); +} + +#[test] +fn backup_storage_build_keys_the_trailing_hash() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + // No host value yet: nothing to publish. + assert!( + tc.build_hicache_transfers( + MAMBA, + tc.arena.node(a).id, + CacheTransferPhase::BackupStorage, + None, + None, + 0, + None, + ) + .is_none() + ); + set_mamba_host(&mut tc, a, 8); + // Host value but no hash chain: still nothing. + assert!( + tc.build_hicache_transfers( + MAMBA, + tc.arena.node(a).id, + CacheTransferPhase::BackupStorage, + None, + None, + 0, + None, + ) + .is_none() + ); + tc.arena.node_mut(a).hash_value = Some(vec!["h0".to_string(), "h1".to_string()]); + let transfers = tc + .build_hicache_transfers( + MAMBA, + tc.arena.node(a).id, + CacheTransferPhase::BackupStorage, + None, + None, + 0, + None, + ) + .unwrap(); + assert_eq!(transfers.len(), 1); + assert_eq!(transfers[0].keys, Some(vec!["h1".to_string()])); + assert_eq!(transfers[0].hit_policy, PoolHitPolicy::TrailingPages); + assert!( + transfers[0] + .host_indices + .as_ref() + .unwrap() + .equal(&Tensor::from_slice(&[8i64])) + ); +} + +#[test] +fn prefetch_build_wraps_the_host_buffer_with_a_placeholder_key() { + let tc = mamba_core(/* page_size = */ 1); + let transfers = tc + .build_hicache_transfers( + MAMBA, + tc.arena.node(tc.arena.root()).id, + CacheTransferPhase::Prefetch, + Some(Tensor::from_slice(&[30i64])), + None, + 0, + None, + ) + .unwrap(); + assert_eq!(transfers.len(), 1); + assert_eq!(transfers[0].keys, Some(vec!["__placeholder__".to_string()])); + assert_eq!(transfers[0].hit_policy, PoolHitPolicy::TrailingPages); +} + +#[test] +fn prefetch_commit_attaches_the_loaded_slot_to_the_inserted_node() { + let mut tc = mamba_core(/* page_size = */ 1); + tc.set_hicache_enabled(); + let root = tc.arena.root(); + let target = tc + .insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1], + Tensor::from_slice(&[100i64]), + vec!["h0".to_string()], + ) + .inserted_host_node + .unwrap(); + let mut insert_result = InsertResult { + total_len: 1, + inserted_host_node: Some(target), + ..InsertResult::default() + }; + let mut cache_actions = Vec::new(); + tc.commit_hicache_transfers( + tc.arena.node(root).id, + CacheTransferPhase::Prefetch, + HashMap::from([( + MAMBA, + vec![PoolTransfer { + name: PoolName::Mamba, + host_indices: Some(Tensor::from_slice(&[50i64])), + ..Default::default() + }], + )]), + &mut cache_actions, + Some(&mut insert_result), + Some(&PoolTransferResult { + kv_hit_pages: 1, + extra_pool_hit_pages: HashMap::from([(PoolName::Mamba, 1)]), + }), + ); + assert!( + tc.arena + .node(tc.arena.resolve(target)) + .try_host_value(MAMBA) + .unwrap() + .equal(&Tensor::from_slice(&[50i64])) + ); + assert!( + tc.host_lru_list(MAMBA) + .in_list(Some(tc.arena.resolve(target))) + ); + assert!(!insert_result.mamba_exist); + assert!(cache_actions.is_empty()); +} + +#[test] +fn prefetch_commit_frees_the_buffer_when_it_cannot_attach() { + let mut tc = mamba_core(/* page_size = */ 1); + tc.set_hicache_enabled(); + let root = tc.arena.root(); + let target = tc + .insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1], + Tensor::from_slice(&[100i64]), + vec!["h0".to_string()], + ) + .inserted_host_node + .unwrap(); + // Not loaded: the buffer frees and the caller keeps its slot flag. + let mut insert_result = InsertResult { + total_len: 1, + inserted_host_node: Some(target), + ..InsertResult::default() + }; + let mut cache_actions = Vec::new(); + tc.commit_hicache_transfers( + tc.arena.node(root).id, + CacheTransferPhase::Prefetch, + HashMap::from([( + MAMBA, + vec![PoolTransfer { + name: PoolName::Mamba, + host_indices: Some(Tensor::from_slice(&[50i64])), + ..Default::default() + }], + )]), + &mut cache_actions, + Some(&mut insert_result), + Some(&PoolTransferResult { + kv_hit_pages: 1, + extra_pool_hit_pages: HashMap::new(), + }), + ); + assert!( + !tc.arena + .node(tc.arena.resolve(target)) + .has_host_value(MAMBA) + ); + assert!(insert_result.mamba_exist); + let CacheAction::FreeComponentHostSlot { + component_type, + host_indices, + } = &cache_actions[0] + else { + panic!("expected a FreeComponentHostSlot action"); + }; + assert_eq!(*component_type, MAMBA); + assert!(host_indices[0].equal(&Tensor::from_slice(&[50i64]))); + + // An already-hosted target frees the buffer too. + let target_idx = tc.arena.resolve(target); + set_mamba_host(&mut tc, target_idx, 8); + let mut insert_result = InsertResult { + total_len: 1, + inserted_host_node: Some(target), + ..InsertResult::default() + }; + let mut cache_actions = Vec::new(); + tc.commit_hicache_transfers( + tc.arena.node(root).id, + CacheTransferPhase::Prefetch, + HashMap::from([( + MAMBA, + vec![PoolTransfer { + name: PoolName::Mamba, + host_indices: Some(Tensor::from_slice(&[51i64])), + ..Default::default() + }], + )]), + &mut cache_actions, + Some(&mut insert_result), + Some(&PoolTransferResult { + kv_hit_pages: 1, + extra_pool_hit_pages: HashMap::from([(PoolName::Mamba, 1)]), + }), + ); + assert!(insert_result.mamba_exist); + assert_eq!(cache_actions.len(), 1); + assert!( + tc.arena + .node(tc.arena.resolve(target)) + .try_host_value(MAMBA) + .unwrap() + .equal(&Tensor::from_slice(&[8i64])) + ); +} + +#[test] +fn new_combines_the_chunk_and_tree_page_grids() { + let component = MambaComponent::new(&CacheInitParams { + page_size: 6, + mamba_cache_chunk_size: Some(4), + ..CacheInitParams::default() + }); + assert_eq!(component.mamba_checkpoint_grid, 12); +} + +#[test] +#[should_panic(expected = "requires mamba_cache_chunk_size")] +fn new_panics_without_a_chunk_size() { + MambaComponent::new(&CacheInitParams::default()); +} + +#[test] +fn evict_excess_path_states_removes_the_shallowest_states_beyond_the_cap() { + let mut tc = mamba_core_with_cap(2); + let [a, b, c] = chain::<3>(&mut tc); + set_mamba_device(&mut tc, a, 7); + set_mamba_device(&mut tc, b, 8); + set_mamba_device(&mut tc, c, 9); + let mut result = tc.evict_excess_path_states(tc.arena.node(c).id); + let freed = result + .device_frees + .remove(&MAMBA) + .expect("the excess state frees its slot"); + assert_eq!(freed.len(), 1); + assert!(freed[0].equal(&Tensor::from_slice(&[7i64]))); + assert!(result.host_frees.is_empty()); + assert!(tc.arena.node(a).try_device_value(MAMBA).is_none()); + assert!(tc.arena.node(b).try_device_value(MAMBA).is_some()); + assert!(tc.arena.node(c).try_device_value(MAMBA).is_some()); +} + +#[test] +fn evict_excess_path_states_preserves_forks_locked_nodes_and_the_tail() { + let mut tc = mamba_core_with_cap(1); + let [a, b, c] = chain::<3>(&mut tc); + set_mamba_device(&mut tc, a, 7); + set_mamba_device(&mut tc, b, 8); + set_mamba_device(&mut tc, c, 9); + // a forks; b is locked: the cap is soft and neither state is removed. + tc.arena + .alloc_child( + a, + /* key = */ vec![9], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .node_mut(b) + .set_lock_ref_(ValueSlotIdx::device(MAMBA), 1); + let result = tc.evict_excess_path_states(tc.arena.node(c).id); + assert!(result.device_frees.is_empty()); + assert!(result.host_frees.is_empty()); + assert!(tc.arena.node(a).try_device_value(MAMBA).is_some()); + assert!(tc.arena.node(b).try_device_value(MAMBA).is_some()); + assert!(tc.arena.node(c).try_device_value(MAMBA).is_some()); +} + +#[test] +fn evict_excess_path_states_without_a_cap_is_a_no_op() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a, b] = chain::<2>(&mut tc); + set_mamba_device(&mut tc, a, 7); + set_mamba_device(&mut tc, b, 8); + let result = tc.evict_excess_path_states(tc.arena.node(b).id); + assert!(result.device_frees.is_empty()); + assert!(result.host_frees.is_empty()); + assert!(tc.arena.node(a).try_device_value(MAMBA).is_some()); +} + +#[test] +fn insert_commit_emits_the_path_cap_action_only_when_capped() { + let mut tc = mamba_core_with_cap(1); + let result = tc.insert(&insert_params_mamba(&vec![1, 2], &[10, 11], Some(7))); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + let [CacheAction::MambaEvictExcessPathStates { tail_node_id }] = + result.cache_actions.as_slice() + else { + panic!( + "expected one MambaEvictExcessPathStates, got {:?}", + action_kinds(&result.cache_actions) + ); + }; + assert_eq!(*tail_node_id, leaf); + + let mut uncapped = mamba_core(/* page_size = */ 1); + let result = uncapped.insert(&insert_params_mamba(&vec![1, 2], &[10, 11], Some(7))); + assert!(result.cache_actions.is_empty()); +} + +#[test] +fn swa_evict_on_a_full_locked_leaf_sweeps_mamba_and_spares_full() { + let mut tc = UnifiedTreeCore::>::new( + CacheInitParams { + page_size: 1, + swa_sliding_window_size: Some(4), + mamba_cache_chunk_size: Some(256), + ..CacheInitParams::default() + }, + vec![FULL, SWA, MAMBA], + ); + let [a] = chain::<1>(&mut tc); + tc.arena + .set_device_value(a, FULL, Tensor::from_slice(&[10i64])); + tc.set_component_device_value(tc.arena.node(a).id, SWA, Tensor::from_slice(&[20i64])); + set_mamba_device(&mut tc, a, 7); + // The held Full lock keeps the leaf out of the D-leaf set. + tc.arena + .node_mut(a) + .set_lock_ref_(ValueSlotIdx::device(FULL), 1); + tc.update_evictable_leaf_sets_(a); + assert!(!tc.evictable_device_leaves.contains(a)); + let mut tracker = HashMap::from([(FULL, 0), (SWA, 0), (MAMBA, 0)]); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + tc.evict_device_start(SWA, /* request_cnt = */ 10); + let (next, step) = tc.evict_device_next_node(SWA, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + assert_eq!(next, None); + tc.evict_device_end(SWA); + // The walk tombstoned SWA inline and cascaded the lower-tier mamba slot. + assert!(!tc.arena.has_device_value(a, SWA)); + assert!(!tc.arena.node(a).has_device_value(MAMBA)); + assert_eq!(tracker[&SWA], 1); + assert_eq!(tracker[&MAMBA], 1); + assert!(device_frees[&SWA][0].equal(&Tensor::from_slice(&[10i64]))); + assert!(device_frees[&MAMBA][0].equal(&Tensor::from_slice(&[7i64]))); + // The higher-tier locked Full is spared and pins the node in the tree. + assert!( + tc.arena + .device_value(a, FULL) + .equal(&Tensor::from_slice(&[10i64])) + ); + assert_eq!(tc.arena.device_lock_ref(a, FULL), 1); + assert_eq!(tc.arena.len(), 2); + assert!(!tc.device_lru_list(SWA).in_list(Some(a))); + assert!(!tc.device_lru_list(MAMBA).in_list(Some(a))); + assert!(!tc.evictable_device_leaves.contains(a)); + assert_eq!(tc.evictable_size_(SWA), 0); + assert_eq!(tc.evictable_size_(MAMBA), 0); +} + +#[test] +fn host_only_mamba_anchor_disables_branching_under_hicache() { + // Without hicache, the mamba-less walk reports the chunk-aligned branch point. + let mut plain = mamba_core_with_chunk(/* page_size = */ 1, /* chunk = */ 3); + let root = plain.arena.root(); + let n = plain + .arena + .alloc_child( + root, + /* key = */ vec![1, 2, 3, 4], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + plain + .arena + .set_device_value(n, FULL, Tensor::from_slice(&[10i64, 11, 12, 13])); + let result = plain.match_prefix(&match_params(&vec![1, 2, 3, 4])); + assert_eq!(result.best_match_node_id, plain.arena.node(root).id); + assert_eq!(result.last_device_node_id, plain.arena.node(root).id); + assert_eq!(result.mamba_branching_seqlen, Some(3)); + + // HiCache: the host-only anchor advances the best match; no branch point. + let mut tc = mamba_core_with_chunk(/* page_size = */ 1, /* chunk = */ 3); + tc.set_hicache_enabled(); + let root = tc.arena.root(); + let n = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2, 3, 4], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(n, FULL, Tensor::from_slice(&[10i64, 11, 12, 13])); + set_mamba_host(&mut tc, n, 8); + tc.host_lru_list_mut(MAMBA).insert_mru(n); + let result = tc.match_prefix(&match_params(&vec![1, 2, 3, 4])); + assert_eq!(result.best_match_node_id, tc.arena.node(n).id); + assert_eq!(result.last_device_node_id, tc.arena.node(root).id); + assert_eq!(result.mamba_branching_seqlen, None); + assert!(result.mamba_host_hit_length >= 1); + assert_eq!(result.host_hit_length, 4); +} + +#[test] +fn branching_from_a_host_full_hit_is_reusable_after_insert() { + let mut tc = mamba_core_with_chunk(/* page_size = */ 1, /* chunk = */ 3); + tc.set_hicache_enabled(); + tc.insert(&insert_params_mamba(&vec![1, 2, 3], &[10, 11, 12], Some(7))); + tc.insert(&insert_params_mamba( + &vec![1, 2, 3, 4, 5, 6, 7], + &[10, 11, 12, 13, 14, 15, 16], + Some(8), + )); + let a = tc + .match_prefix(&match_params(&vec![1, 2, 3])) + .best_match_node_id; + let b = tc + .match_prefix(&match_params(&vec![1, 2, 3, 4, 5, 6, 7])) + .best_match_node_id; + tc.commit_backup( + b, + Tensor::from_slice(&[100i64, 101, 102, 103]), + HashMap::new(), + ); + tc.demote(b); + // The demote's cascade swept b's mamba slot: b is Full-host-only, no mamba. + assert!(!tc.arena.node(tc.arena.resolve(b)).has_device_value(MAMBA)); + assert!(!tc.arena.node(tc.arena.resolve(b)).has_host_value(MAMBA)); + let result = tc.match_prefix(&match_params(&vec![1, 2, 3, 4, 5, 6, 7])); + assert_eq!(result.best_match_node_id, a); + assert_eq!(result.last_device_node_id, a); + assert_eq!(result.device_indices.numel(), 3); + assert_eq!(result.host_hit_length, 0); + assert_eq!(result.full_kv_hit_length, 7); + assert_eq!(result.mamba_branching_seqlen, Some(6)); + // Re-inserting up to the branch point makes the span device-reusable. + let insert_result = tc.insert(&insert_params_mamba( + &vec![1, 2, 3, 4, 5, 6], + &[20, 21, 22, 23, 24, 25], + Some(9), + )); + assert!(!insert_result.mamba_exist); + let second = tc.match_prefix(&match_params(&vec![1, 2, 3, 4, 5, 6, 7])); + assert_eq!(second.device_indices.numel(), 6); + assert_eq!(second.mamba_branching_seqlen, None); +} + +#[test] +fn skip_set_release_after_a_restore_and_relock_keeps_the_new_lock() { + let mut tc = mamba_core(/* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + let mamba = mamba_component(); + let first = mamba.acquire_component_lock( + &mut tc, + a, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + assert!(first.skip_lock_node_ids[&MAMBA].contains(&tc.arena.node(a).id)); + // The tombstone is restored and a second request locks it before the + // first release replays its skip set. + set_mamba_device(&mut tc, a, 7); + let _ = mamba.acquire_component_lock( + &mut tc, + a, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 1); + assert_eq!(tc.evictable_size_(MAMBA), 0); + assert_eq!(tc.protected_size_(MAMBA), 1); + let params = DecLockRefParams { + skip_lock_node_ids: first.skip_lock_node_ids.clone(), + ..DecLockRefParams::default() + }; + mamba.release_component_lock(&mut tc, a, Some(¶ms), /* lock_host = */ false); + // The replayed skip keeps the restored node's fresh lock intact. + assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 1); + assert_eq!(tc.protected_size_(MAMBA), 1); + mamba.release_component_lock(&mut tc, a, None, /* lock_host = */ false); + assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 0); + assert_eq!(tc.evictable_size_(MAMBA), 1); + assert_eq!(tc.protected_size_(MAMBA), 0); +} diff --git a/rust/mem-cache/src/tests/components/swa.rs b/rust/mem-cache/src/tests/components/swa.rs new file mode 100644 index 000000000..b68e7a7aa --- /dev/null +++ b/rust/mem-cache/src/tests/components/swa.rs @@ -0,0 +1,4894 @@ +use super::*; +use crate::components::{FULL, MAMBA, SWA}; +use crate::test_utils::{accumulate_step, action_kinds}; +use crate::unified_tree_core::CacheInitParams; + +#[test] +fn component_type_is_swa() { + let swa = SwaComponent::new(&swa_params()); + assert_eq!( + >>::component_type(&swa), + SWA + ); +} + +fn swa_params() -> CacheInitParams { + CacheInitParams { + swa_sliding_window_size: Some(4096), + ..Default::default() + } +} + +#[test] +#[should_panic(expected = "requires swa_sliding_window_size")] +fn new_panics_without_a_sliding_window_size() { + SwaComponent::new(&CacheInitParams::default()); +} + +#[test] +fn new_stores_the_sliding_window_size() { + let params = CacheInitParams { + swa_sliding_window_size: Some(4096), + ..Default::default() + }; + assert_eq!(SwaComponent::new(¶ms).sliding_window_size, 4096); +} + +#[test] +fn construction_with_full_and_swa_locks_both_roots() { + let tc: UnifiedTreeCore> = UnifiedTreeCore::new(swa_params(), vec![FULL, SWA]); + let root = tc.arena.root(); + let root_node = tc.arena.node(root); + assert_eq!(root_node.values[FULL.idx()].lock_ref, 1); + assert_eq!(root_node.values[SWA.idx()].lock_ref, 1); + assert_eq!(root_node.values[MAMBA.idx()].lock_ref, 0); +} + +#[test] +fn swa_sizes_read_zero_on_a_fresh_tree() { + let tc: UnifiedTreeCore> = UnifiedTreeCore::new(swa_params(), vec![FULL, SWA]); + assert_eq!(tc.swa_evictable_size(), 0); + assert_eq!(tc.swa_protected_size(), 0); +} + +#[test] +fn node_has_component_data_tracks_each_slot() { + let mut tc: UnifiedTreeCore> = UnifiedTreeCore::new(swa_params(), vec![FULL, SWA]); + let root = tc.arena.root(); + let node = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let swa = SwaComponent::new(&swa_params()); + assert!(!crate::components::node_has_component_data( + &tc.arena, + node, + SWA, + EvictLayer::Device + )); + assert!(!crate::components::node_has_component_data( + &tc.arena, + node, + SWA, + EvictLayer::Host + )); + + tc.arena + .set_device_value(node, SWA, Tensor::from_slice(&[10i64])); + assert!(crate::components::node_has_component_data( + &tc.arena, + node, + SWA, + EvictLayer::Device + )); + assert!(!crate::components::node_has_component_data( + &tc.arena, + node, + SWA, + EvictLayer::Host + )); + + tc.arena + .set_host_value(node, SWA, Tensor::from_slice(&[20i64])); + assert!(crate::components::node_has_component_data( + &tc.arena, + node, + SWA, + EvictLayer::Host + )); +} + +// A [Full, Swa] core with the given sliding window and page size. +fn swa_core(window: usize, page_size: usize) -> UnifiedTreeCore> { + UnifiedTreeCore::new( + CacheInitParams { + page_size, + ..swa_params_with_window(window) + }, + vec![FULL, SWA], + ) +} + +fn swa_params_with_window(window: usize) -> CacheInitParams { + CacheInitParams { + swa_sliding_window_size: Some(window), + ..Default::default() + } +} + +fn swa_component(window: usize) -> SwaComponent { + SwaComponent::new(&swa_params_with_window(window)) +} + +// A chain of single-atom children under the default root; returns the node ids. +fn chain(tc: &mut UnifiedTreeCore>) -> [NodeIdx_; N] { + let mut parent = tc.arena.root(); + let mut nodes = [NodeIdx_(0); N]; + for (i, node) in nodes.iter_mut().enumerate() { + let id = tc + .arena + .alloc_child( + parent, + /* key = */ vec![i as i64 + 1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + *node = id; + parent = id; + } + nodes +} + +fn set_swa_device(tc: &mut UnifiedTreeCore>, node: NodeIdx_) { + let len = tc.arena.node(node).key.atom_len(); + tc.arena + .set_device_value(node, SWA, Tensor::from_slice(&vec![0i64; len])); +} + +fn set_swa_host(tc: &mut UnifiedTreeCore>, node: NodeIdx_) { + let len = tc.arena.node(node).key.atom_len(); + tc.arena + .set_host_value(node, SWA, Tensor::from_slice(&vec![0i64; len])); +} + +fn node_swa_uuid(tc: &UnifiedTreeCore>, node: NodeIdx_) -> Option { + tc.arena.node(node).swa_uuid +} + +fn node_swa_host_uuid(tc: &UnifiedTreeCore>, node: NodeIdx_) -> Option { + tc.arena.node(node).swa_host_uuid +} + +// Give `node` an SWA device value via the auxiliary store (sizes + LRU stamped). +fn store_swa_device(tc: &mut UnifiedTreeCore>, node: NodeIdx_) { + let len = tc.arena.node(node).key.atom_len(); + tc.set_component_device_value( + tc.arena.node(node).id, + SWA, + Tensor::from_slice(&vec![0i64; len]), + ); +} + +#[test] +fn match_validator_accepts_valued_nodes_before_any_gap() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a, b] = chain(&mut tc); + set_swa_device(&mut tc, a); + set_swa_device(&mut tc, b); + // The window starts unbounded: a 2-atom valued span validates under a + // window of 4 because no gap has been seen yet. + let mut validator = + swa_component(4).create_match_validator(&tc, /* match_device_only = */ true); + assert!(validator(&tc, a)); + assert!(validator(&tc, b)); +} + +#[test] +fn match_validator_gap_resets_until_the_window_is_reached() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, t, c, d, e] = chain(&mut tc); + set_swa_device(&mut tc, a); + set_swa_device(&mut tc, c); + set_swa_device(&mut tc, d); + set_swa_device(&mut tc, e); + let mut validator = + swa_component(2).create_match_validator(&tc, /* match_device_only = */ true); + assert!(validator(&tc, a)); + // The tombstone resets the run; below the window stays invalid, the + // exact window boundary revalidates, and beyond it stays valid. + assert!(!validator(&tc, t)); + assert!(!validator(&tc, c)); + assert!(validator(&tc, d)); + assert!(validator(&tc, e)); +} + +#[test] +fn match_validator_window_larger_than_the_remaining_span_never_revalidates() { + let mut tc = swa_core(/* window = */ 3, /* page_size = */ 1); + let [_a, t, c, d] = chain(&mut tc); + set_swa_device(&mut tc, c); + set_swa_device(&mut tc, d); + let mut validator = + swa_component(3).create_match_validator(&tc, /* match_device_only = */ true); + assert!(!validator(&tc, t)); + assert!(!validator(&tc, c)); + assert!(!validator(&tc, d)); +} + +#[test] +fn match_validator_device_only_treats_host_only_swa_as_a_gap() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, h, c] = chain(&mut tc); + set_swa_device(&mut tc, a); + set_swa_host(&mut tc, h); + set_swa_device(&mut tc, c); + let mut validator = + swa_component(2).create_match_validator(&tc, /* match_device_only = */ true); + assert!(validator(&tc, a)); + assert!(!validator(&tc, h)); + assert!(!validator(&tc, c)); +} + +#[test] +fn match_validator_host_backed_nodes_extend_the_window() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [t, h, c] = chain(&mut tc); + set_swa_host(&mut tc, h); + set_swa_device(&mut tc, c); + // Without match_device_only, a host-backed node is not a gap: it + // counts toward the window like a device-backed one. + let mut validator = + swa_component(2).create_match_validator(&tc, /* match_device_only = */ false); + assert!(!validator(&tc, t)); + assert!(!validator(&tc, h)); + assert!(validator(&tc, c)); +} + +#[test] +fn match_validator_hicache_accepts_live_or_backuped_swa_tombstones() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + tc.set_hicache_enabled(); + let [live, backuped, dead] = chain(&mut tc); + tc.arena + .set_device_value(live, FULL, Tensor::from_slice(&[0i64])); + tc.arena + .set_host_value(backuped, FULL, Tensor::from_slice(&[0i64])); + let mut validator = + swa_component(2).create_match_validator(&tc, /* match_device_only = */ true); + assert!(validator(&tc, live)); + assert!(validator(&tc, backuped)); + assert!(!validator(&tc, dead)); +} + +#[test] +fn match_validator_without_hicache_rejects_every_swa_tombstone() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [live, backuped, dead] = chain(&mut tc); + tc.arena + .set_device_value(live, FULL, Tensor::from_slice(&[0i64])); + tc.arena + .set_host_value(backuped, FULL, Tensor::from_slice(&[0i64])); + let mut validator = + swa_component(2).create_match_validator(&tc, /* match_device_only = */ true); + assert!(!validator(&tc, live)); + assert!(!validator(&tc, backuped)); + assert!(!validator(&tc, dead)); +} + +#[test] +fn match_validator_hicache_with_a_host_pool_rejects_swa_tombstones() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + tc.set_hicache_enabled(); + let [live] = chain(&mut tc); + tc.arena + .set_device_value(live, FULL, Tensor::from_slice(&[0i64])); + // A wired host SWA pool means tombstones must gate the match again. + tc.set_has_swa_host_pool(); + let swa = SwaComponent { + sliding_window_size: 2, + }; + let mut validator = >>::create_match_validator( + &swa, &tc, /* match_device_only = */ true, + ); + assert!(!validator(&tc, live)); +} + +#[test] +fn match_validator_hicache_tombstone_acceptance_still_resets_the_window() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + tc.set_hicache_enabled(); + let [live, c] = chain(&mut tc); + tc.arena + .set_device_value(live, FULL, Tensor::from_slice(&[0i64])); + set_swa_device(&mut tc, c); + let mut validator = + swa_component(2).create_match_validator(&tc, /* match_device_only = */ true); + // The accepted tombstone still zeroes the run: the next valued node + // sits below the window. + assert!(validator(&tc, live)); + assert!(!validator(&tc, c)); +} + +// The SWA LRU order, MRU to LRU. +fn swa_lru_order(tc: &UnifiedTreeCore>) -> Vec { + tc.device_lru_list(SWA).iter().collect() +} + +#[test] +fn refresh_lru_walkdown_is_a_noop() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let root = tc.arena.root(); + let [a, b] = chain(&mut tc); + set_swa_device(&mut tc, a); + set_swa_device(&mut tc, b); + tc.device_lru_list_mut(SWA).insert_mru(a); + tc.device_lru_list_mut(SWA).insert_mru(b); + swa_component(2).refresh_lru(&mut tc, LRURefreshPhase::Walkdown, a); + assert_eq!(swa_lru_order(&tc), vec![b, a]); +} + +#[test] +fn refresh_lru_match_end_reranks_the_window_run_deepest_first() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let root = tc.arena.root(); + let [a, b, c, d] = chain(&mut tc); + for node in [d, c, b, a] { + set_swa_device(&mut tc, node); + tc.device_lru_list_mut(SWA).insert_mru(node); + } + // The walk window is sliding_window_size + page_size = 3: d, c, b are + // re-ranked deepest first; a stays beyond the window. + swa_component(2).refresh_lru(&mut tc, LRURefreshPhase::MatchEnd, d); + assert_eq!(swa_lru_order(&tc), vec![d, c, b, a]); +} + +#[test] +fn refresh_lru_insert_end_matches_the_match_end_walk() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let root = tc.arena.root(); + let [a, b, c, d] = chain(&mut tc); + for node in [d, c, b, a] { + set_swa_device(&mut tc, node); + tc.device_lru_list_mut(SWA).insert_mru(node); + } + swa_component(2).refresh_lru(&mut tc, LRURefreshPhase::InsertEnd, d); + assert_eq!(swa_lru_order(&tc), vec![d, c, b, a]); +} + +#[test] +fn refresh_lru_window_walk_skips_tombstones_but_counts_their_span() { + let mut tc = swa_core(/* window = */ 1, /* page_size = */ 1); + let root = tc.arena.root(); + let [a, _t, c] = chain(&mut tc); + let s = tc + .arena + .alloc_child( + root, + /* key = */ vec![9], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + for node in [c, a, s] { + set_swa_device(&mut tc, node); + tc.device_lru_list_mut(SWA).insert_mru(node); + } + // The walk window is 2: the unlisted tombstone t is skipped but its + // atom consumes the window, so a is never re-ranked. + swa_component(1).refresh_lru(&mut tc, LRURefreshPhase::MatchEnd, c); + assert_eq!(swa_lru_order(&tc), vec![c, s, a]); +} + +// Finalize `best` against an otherwise-empty match result carrying a prior +// SWA host hit. +fn finalize( + tc: &UnifiedTreeCore>, + swa: &SwaComponent, + best: NodeIdx_, + prior_swa_host_hit: usize, +) -> MatchResult { + swa.finalize_match_result_in_tree_core( + tc, + MatchResult { + best_match_node_id: tc.arena.node(best).id, + swa_host_hit_length: prior_swa_host_hit, + ..tc.empty_match_result() + }, + &MatchPrefixParams { + key: &Vec::new(), + namespace: Default::default(), + }, + &[], + 0, + ) +} + +#[test] +fn finalize_without_host_chunks_leaves_the_result_unchanged() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a, b] = chain(&mut tc); + set_swa_device(&mut tc, a); + set_swa_device(&mut tc, b); + let out = finalize(&tc, &swa_component(4), b, /* prior = */ 0); + assert_eq!(out.swa_host_hit_length, 0); + assert_eq!(out.best_match_node_id, tc.arena.node(b).id); +} + +#[test] +fn finalize_sums_swa_host_chunks_within_the_window() { + let mut tc = swa_core(/* window = */ 5, /* page_size = */ 1); + let [a, h, b, c] = chain(&mut tc); + set_swa_host(&mut tc, a); + set_swa_host(&mut tc, h); + set_swa_device(&mut tc, b); + set_swa_device(&mut tc, c); + // From c up: device 1 + 1, then host 1 + 1 — all inside the window of 5. + let out = finalize(&tc, &swa_component(5), c, /* prior = */ 0); + assert_eq!(out.swa_host_hit_length, 2); +} + +#[test] +fn finalize_stops_at_the_window_before_higher_host_chunks() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [h, b, c] = chain(&mut tc); + set_swa_host(&mut tc, h); + set_swa_device(&mut tc, b); + set_swa_device(&mut tc, c); + // The device span alone covers the window of 2: the host chunk above + // the boundary is never counted. + let out = finalize(&tc, &swa_component(2), c, /* prior = */ 0); + assert_eq!(out.swa_host_hit_length, 0); +} + +#[test] +fn finalize_counts_the_straddling_host_chunk_in_full() { + let mut tc = swa_core(/* window = */ 3, /* page_size = */ 1); + let root = tc.arena.root(); + let h = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2, 3, 4], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let c = tc + .arena + .alloc_child( + h, + /* key = */ vec![5, 6], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + set_swa_host(&mut tc, h); + set_swa_device(&mut tc, c); + // The host chunk straddles the window boundary (2 of its 4 tokens are + // in-window) and is counted in full, uncapped. + let out = finalize(&tc, &swa_component(3), c, /* prior = */ 0); + assert_eq!(out.swa_host_hit_length, 4); +} + +#[test] +fn finalize_breaks_at_an_swa_gap() { + let mut tc = swa_core(/* window = */ 5, /* page_size = */ 1); + let [h, _t, c] = chain(&mut tc); + set_swa_host(&mut tc, h); + set_swa_device(&mut tc, c); + // The tombstone between c and h ends the walk: the host chunk above + // the gap is unreachable. + let out = finalize(&tc, &swa_component(5), c, /* prior = */ 0); + assert_eq!(out.swa_host_hit_length, 0); +} + +#[test] +fn finalize_keeps_a_larger_existing_swa_host_hit() { + let mut tc = swa_core(/* window = */ 5, /* page_size = */ 1); + let [h, c] = chain(&mut tc); + set_swa_host(&mut tc, h); + set_swa_device(&mut tc, c); + let out = finalize(&tc, &swa_component(5), c, /* prior = */ 100); + assert_eq!(out.swa_host_hit_length, 100); +} + +#[test] +fn finalize_overrides_a_smaller_existing_swa_host_hit() { + let mut tc = swa_core(/* window = */ 5, /* page_size = */ 1); + let [h, hh, c] = chain(&mut tc); + set_swa_host(&mut tc, h); + set_swa_host(&mut tc, hh); + set_swa_device(&mut tc, c); + let out = finalize(&tc, &swa_component(5), c, /* prior = */ 1); + assert_eq!(out.swa_host_hit_length, 2); +} + +#[test] +fn finalize_walk_stops_at_the_root() { + let mut tc = swa_core(/* window = */ 5, /* page_size = */ 1); + let lora = tc.arena.root(); + let c = tc + .arena + .alloc_child( + lora, + /* key = */ vec![1], + /* priority = */ 0, + Some("lora-1"), + ) + .unwrap(); + set_swa_device(&mut tc, c); + // A host value on the root itself is never counted: the walk ends there. + tc.arena + .node_mut(lora) + .state_mut_(ValueSlotIdx::host(SWA)) + .value = Some(Tensor::from_slice(&[0i64])); + let out = finalize(&tc, &swa_component(5), c, /* prior = */ 0); + assert_eq!(out.swa_host_hit_length, 0); +} + +fn insert_params_swa<'k>( + key: &'k Vec, + value: &[i64], + prev_prefix_len: usize, + swa_evicted_seqlen: usize, +) -> InsertParams<'k, Vec> { + InsertParams { + key, + namespace: Default::default(), + value: Tensor::from_slice(value), + mamba_value: None, + prev_prefix_len, + swa_evicted_seqlen, + chunked: false, + priority: 0, + track_adopted_ranges: false, + } +} + +// The child of `node` along the edge keyed by `page` (default namespace). +fn child_of(tc: &UnifiedTreeCore>, node: NodeIdx_, page: &[i64]) -> NodeIdx_ { + tc.arena + .child_on_page(node, /* extra_key = */ None, page) + .expect("child on page") +} + +// Tombstone a leaf's FULL device value the way eviction leaves it. +fn evict_full(tc: &mut UnifiedTreeCore>, leaf: NodeIdx_, remaining_size: usize) { + let _ = tc.arena.take_device_value(leaf, FULL); + tc.component_state_mut(FULL).evictable_size = remaining_size; + tc.evictable_device_leaves.discard(leaf); +} + +#[test] +fn insert_new_leaf_in_window_emits_one_leaf_rebuild() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3], + &[10, 11, 12], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 0, + )); + assert_eq!(result.prefix_len, 0); + let root = tc.arena.root(); + let leaf = child_of(&tc, root, &[1]); + let [ + CacheAction::SwaRebuild { + node_id, + source_value, + }, + ] = result.cache_actions.as_slice() + else { + panic!("expected one SwaRebuild action"); + }; + assert_eq!(*node_id, tc.arena.node(leaf).id); + assert!(source_value.equal(&Tensor::from_slice(&[10i64, 11, 12]))); + // The rebuild is deferred to apply time: the leaf is still an SWA tombstone. + assert!(!tc.arena.has_device_value(leaf, SWA)); +} + +#[test] +fn insert_new_leaf_straddling_the_boundary_splits_and_rebuilds_the_tail() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4], + &[10, 11, 12, 13], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 2, + )); + let root = tc.arena.root(); + let parent = child_of(&tc, root, &[1]); + let child = child_of(&tc, parent, &[3]); + assert_eq!(tc.arena.node(parent).key, vec![1, 2]); + assert_eq!(tc.arena.node(child).key, vec![3, 4]); + // The out-of-window parent is an SWA tombstone holding its own Full span. + assert!(!tc.arena.has_device_value(parent, SWA)); + assert!( + tc.arena + .device_value(parent, FULL) + .equal(&Tensor::from_slice(&[10i64, 11])) + ); + let [ + CacheAction::SwaRebuild { + node_id, + source_value, + }, + ] = result.cache_actions.as_slice() + else { + panic!("expected one SwaRebuild action"); + }; + assert_eq!(*node_id, tc.arena.node(child).id); + assert!(source_value.equal(&Tensor::from_slice(&[12i64, 13]))); +} + +#[test] +fn insert_new_leaf_outside_the_window_stays_a_tombstone() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + // The boundary equals the leaf end: no split, no rebuild. + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3], + &[10, 11, 12], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 3, + )); + assert!(result.cache_actions.is_empty()); + let root = tc.arena.root(); + let leaf = child_of(&tc, root, &[1]); + assert_eq!(tc.arena.node(leaf).key, vec![1, 2, 3]); + assert!(!tc.arena.has_device_value(leaf, SWA)); +} + +#[test] +fn insert_long_leaf_caps_the_window_and_rebuilds_the_older_prefix_first() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4, 5, 6], + &[10, 11, 12, 13, 14, 15], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 0, + )); + let root = tc.arena.root(); + let capped = child_of(&tc, root, &[1]); + let tail = child_of(&tc, capped, &[5]); + assert_eq!(tc.arena.node(capped).key, vec![1, 2, 3, 4]); + assert_eq!(tc.arena.node(tail).key, vec![5, 6]); + let [ + CacheAction::SwaRebuild { + node_id: first_id, + source_value: first_value, + }, + CacheAction::SwaRebuild { + node_id: second_id, + source_value: second_value, + }, + ] = result.cache_actions.as_slice() + else { + panic!("expected two SwaRebuild actions"); + }; + assert_eq!(*first_id, tc.arena.node(capped).id); + assert!(first_value.equal(&Tensor::from_slice(&[10i64, 11, 12, 13]))); + assert_eq!(*second_id, tc.arena.node(tail).id); + assert!(second_value.equal(&Tensor::from_slice(&[14i64, 15]))); +} + +#[test] +fn cap_split_uses_the_page_rounded_window() { + let mut tc = swa_core(/* window = */ 3, /* page_size = */ 2); + tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4, 5, 6], + &[10, 11, 12, 13, 14, 15], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 0, + )); + let root = tc.arena.root(); + let capped = child_of(&tc, root, &[1, 2]); + // tail_size rounds the window of 3 up to 2 pages: the tail keeps 4 + // atoms, the parent 2. + assert_eq!(tc.arena.node(capped).key, vec![1, 2]); + let tail = child_of(&tc, capped, &[3, 4]); + assert_eq!(tc.arena.node(tail).key, vec![3, 4, 5, 6]); +} + +#[test] +fn insert_overlap_with_live_swa_frees_the_whole_duplicate() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0)); + let root = tc.arena.root(); + let leaf = child_of(&tc, root, &[1]); + tc.set_component_device_value( + tc.arena.node(leaf).id, + SWA, + Tensor::from_slice(&[50i64, 51, 52]), + ); + let result = tc.insert(&insert_params_swa(&vec![1, 2, 3], &[20, 21, 22], 0, 0)); + assert_eq!(result.prefix_len, 3); + let [CacheAction::FreeDeviceKV(freed)] = result.cache_actions.as_slice() else { + panic!("expected one FreeDeviceKV action"); + }; + assert_eq!(freed.len(), 1); + assert!(freed[0].equal(&Tensor::from_slice(&[20i64, 21, 22]))); + // The node keeps its original Full and SWA values. + assert!( + tc.arena + .device_value(leaf, FULL) + .equal(&Tensor::from_slice(&[10i64, 11, 12])) + ); + assert!( + tc.arena + .device_value(leaf, SWA) + .equal(&Tensor::from_slice(&[50i64, 51, 52])) + ); +} + +#[test] +fn insert_overlap_recovers_a_tombstone_inside_the_window() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0)); + let root = tc.arena.root(); + let leaf = child_of(&tc, root, &[1]); + let result = tc.insert(&InsertParams { + track_adopted_ranges: true, + ..insert_params_swa(&vec![1, 2, 3], &[20, 21, 22], 0, 0) + }); + // SWA consumed the whole slice: the node adopts the fresh Full KV and + // the old Full is freed instead of the duplicates. + assert!( + tc.arena + .device_value(leaf, FULL) + .equal(&Tensor::from_slice(&[20i64, 21, 22])) + ); + let [ + CacheAction::FreeDeviceKVFullOnly(freed), + CacheAction::SwaRebuild { + node_id, + source_value, + }, + ] = result.cache_actions.as_slice() + else { + panic!("expected FreeDeviceKVFullOnly then SwaRebuild"); + }; + assert!(freed[0].equal(&Tensor::from_slice(&[10i64, 11, 12]))); + assert_eq!(*node_id, tc.arena.node(leaf).id); + assert!(source_value.equal(&Tensor::from_slice(&[20i64, 21, 22]))); + let adopted = result.adopted_ranges.as_ref().unwrap(); + assert_eq!(adopted[&FULL], [(0, 3)]); + assert_eq!(adopted[&SWA], [(0, 3)]); +} + +#[test] +#[should_panic(expected = "tombstone Swa lock_ref should be 0, node")] +fn insert_overlap_panics_on_a_locked_swa_tombstone() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0)); + let root = tc.arena.root(); + let leaf = child_of(&tc, root, &[1]); + // The rebuild is deferred, so the leaf is still an SWA tombstone; a raw + // lock on it breaks the tombstones-are-unlocked contract. + tc.arena + .node_mut(leaf) + .set_lock_ref_(ValueSlotIdx::device(SWA), 1); + tc.insert(&insert_params_swa(&vec![1, 2, 3], &[20, 21, 22], 0, 0)); +} + +#[test] +fn insert_overlap_with_a_locked_full_emits_the_recover_action() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0)); + let root = tc.arena.root(); + let leaf = child_of(&tc, root, &[1]); + tc.arena + .node_mut(leaf) + .set_lock_ref_(ValueSlotIdx::device(FULL), 1); + let result = tc.insert(&InsertParams { + track_adopted_ranges: true, + ..insert_params_swa(&vec![1, 2, 3], &[20, 21, 22], 0, 0) + }); + // The locked Full stays on the node; the cache resolves the recover action. + assert!( + tc.arena + .device_value(leaf, FULL) + .equal(&Tensor::from_slice(&[10i64, 11, 12])) + ); + let [ + CacheAction::RecoverSwaWithLockedFull { + node_id, + kept_full, + incoming_full, + }, + ] = result.cache_actions.as_slice() + else { + panic!("expected one RecoverSwaWithLockedFull action"); + }; + assert_eq!(*node_id, tc.arena.node(leaf).id); + assert!(kept_full.equal(&Tensor::from_slice(&[10i64, 11, 12]))); + assert!(incoming_full.equal(&Tensor::from_slice(&[20i64, 21, 22]))); + let adopted = result.adopted_ranges.as_ref().unwrap(); + assert!(!adopted.contains_key(&FULL)); + assert_eq!(adopted[&SWA], [(0, 3)]); +} + +#[test] +fn insert_overlap_straddling_the_boundary_splits_and_recovers_the_tail() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4], + &[10, 11, 12, 13], + 0, + 0, + )); + let root = tc.arena.root(); + let node = child_of(&tc, root, &[1]); + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4], + &[20, 21, 22, 23], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 2, + )); + assert_eq!(result.prefix_len, 4); + // The node split at the boundary: the parent keeps the old + // out-of-window Full span, the tail adopts the fresh KV. + let parent = child_of(&tc, root, &[1]); + assert_eq!(tc.arena.node(parent).key, vec![1, 2]); + assert_eq!(tc.arena.node(node).key, vec![3, 4]); + assert!( + tc.arena + .device_value(parent, FULL) + .equal(&Tensor::from_slice(&[10i64, 11])) + ); + assert!( + tc.arena + .device_value(node, FULL) + .equal(&Tensor::from_slice(&[22i64, 23])) + ); + let [ + CacheAction::FreeDeviceKVFullOnly(old_tail), + CacheAction::SwaRebuild { + node_id, + source_value, + }, + CacheAction::FreeDeviceKV(duplicates), + ] = result.cache_actions.as_slice() + else { + panic!( + "expected FreeDeviceKVFullOnly, SwaRebuild, FreeDeviceKV, got {:?}", + action_kinds(&result.cache_actions) + ); + }; + assert!(old_tail[0].equal(&Tensor::from_slice(&[12i64, 13]))); + assert_eq!(*node_id, tc.arena.node(node).id); + assert!(source_value.equal(&Tensor::from_slice(&[22i64, 23]))); + // Only the out-of-window head is duplicate; the consumed tail is not re-freed. + assert!(duplicates[0].equal(&Tensor::from_slice(&[20i64, 21]))); +} + +#[test] +fn insert_overlap_straddling_with_a_locked_full_defers_the_tail() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4], + &[10, 11, 12, 13], + 0, + 0, + )); + let root = tc.arena.root(); + let node = child_of(&tc, root, &[1]); + tc.arena + .node_mut(node) + .set_lock_ref_(ValueSlotIdx::device(FULL), 1); + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4], + &[20, 21, 22, 23], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 2, + )); + // The locked tail keeps its Full value; only the recover action crosses. + assert!( + tc.arena + .device_value(node, FULL) + .equal(&Tensor::from_slice(&[12i64, 13])) + ); + let [ + CacheAction::RecoverSwaWithLockedFull { + node_id, + kept_full, + incoming_full, + }, + CacheAction::FreeDeviceKV(duplicates), + ] = result.cache_actions.as_slice() + else { + panic!("expected RecoverSwaWithLockedFull then FreeDeviceKV"); + }; + assert_eq!(*node_id, tc.arena.node(node).id); + assert!(kept_full.equal(&Tensor::from_slice(&[12i64, 13]))); + assert!(incoming_full.equal(&Tensor::from_slice(&[22i64, 23]))); + assert!(duplicates[0].equal(&Tensor::from_slice(&[20i64, 21]))); +} + +#[test] +fn insert_overlap_entirely_outside_the_window_is_all_duplicate() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0)); + // The boundary sits at the node end: nothing consumed, no recovery. + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3], + &[20, 21, 22], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 3, + )); + let [CacheAction::FreeDeviceKV(freed)] = result.cache_actions.as_slice() else { + panic!("expected one FreeDeviceKV action"); + }; + assert!(freed[0].equal(&Tensor::from_slice(&[20i64, 21, 22]))); +} + +#[test] +fn insert_overlap_already_cached_prefix_skips_the_tombstone_recovery() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0)); + // prev_prefix_len covers the node: no recovery and no duplicate free. + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3], + &[20, 21, 22], + /* prev_prefix_len = */ 3, + /* swa_evicted_seqlen = */ 0, + )); + assert!(result.cache_actions.is_empty()); +} + +#[test] +fn insert_overlap_boundary_at_the_node_start_recovers_the_whole_node() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0)); + tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4, 5], + &[30, 31, 32, 13, 14], + 0, + 0, + )); + let root = tc.arena.root(); + let a = child_of(&tc, root, &[1]); + let b = child_of(&tc, a, &[4]); + tc.set_component_device_value( + tc.arena.node(a).id, + SWA, + Tensor::from_slice(&[50i64, 51, 52]), + ); + // The boundary lands exactly on b's start: full recovery, no split. + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4, 5], + &[20, 21, 22, 23, 24], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 3, + )); + assert_eq!(tc.arena.node(b).key, vec![4, 5]); + assert!( + tc.arena + .device_value(b, FULL) + .equal(&Tensor::from_slice(&[23i64, 24])) + ); + let [ + CacheAction::FreeDeviceKV(duplicates), + CacheAction::FreeDeviceKVFullOnly(old_full), + CacheAction::SwaRebuild { + node_id, + source_value, + }, + ] = result.cache_actions.as_slice() + else { + panic!( + "expected FreeDeviceKV, FreeDeviceKVFullOnly, SwaRebuild, got {:?}", + action_kinds(&result.cache_actions) + ); + }; + assert!(duplicates[0].equal(&Tensor::from_slice(&[20i64, 21, 22]))); + assert!(old_full[0].equal(&Tensor::from_slice(&[13i64, 14]))); + assert_eq!(*node_id, tc.arena.node(b).id); + assert!(source_value.equal(&Tensor::from_slice(&[23i64, 24]))); +} + +#[test] +fn insert_overlap_straddling_a_second_level_node_recovers_the_tail() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0)); + tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4, 5], + &[10, 11, 12, 13, 14], + 0, + 0, + )); + let root = tc.arena.root(); + let a = child_of(&tc, root, &[1]); + let b = child_of(&tc, a, &[4]); + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4, 5], + &[20, 21, 22, 23, 24], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 4, + )); + assert_eq!(result.prefix_len, 5); + // The boundary lands one atom into b (total_prefix_len 3): b splits at + // its node-relative offset 1, not at the absolute seqlen. + let p = child_of(&tc, a, &[4]); + assert_eq!(tc.arena.node(p).key, vec![4]); + assert_eq!(tc.arena.node(b).key, vec![5]); + assert!( + tc.arena + .device_value(p, FULL) + .equal(&Tensor::from_slice(&[13i64])) + ); + assert!( + tc.arena + .device_value(b, FULL) + .equal(&Tensor::from_slice(&[24i64])) + ); + let [ + CacheAction::FreeDeviceKV(duplicates_head), + CacheAction::FreeDeviceKVFullOnly(old_tail), + CacheAction::SwaRebuild { + node_id, + source_value, + }, + CacheAction::FreeDeviceKV(duplicates_tail), + ] = result.cache_actions.as_slice() + else { + panic!( + "expected FreeDeviceKV, FreeDeviceKVFullOnly, SwaRebuild, FreeDeviceKV, got {:?}", + action_kinds(&result.cache_actions) + ); + }; + assert!(duplicates_head[0].equal(&Tensor::from_slice(&[20i64, 21, 22]))); + assert!(old_tail[0].equal(&Tensor::from_slice(&[14i64]))); + assert_eq!(*node_id, tc.arena.node(b).id); + assert!(source_value.equal(&Tensor::from_slice(&[24i64]))); + assert!(duplicates_tail[0].equal(&Tensor::from_slice(&[23i64]))); +} + +#[test] +fn insert_overlap_straddling_a_second_level_locked_node_defers_the_tail() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0)); + tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4, 5], + &[10, 11, 12, 13, 14], + 0, + 0, + )); + let root = tc.arena.root(); + let a = child_of(&tc, root, &[1]); + let b = child_of(&tc, a, &[4]); + tc.arena + .node_mut(b) + .set_lock_ref_(ValueSlotIdx::device(FULL), 1); + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4, 5], + &[20, 21, 22, 23, 24], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 4, + )); + // The locked tail keeps its old Full value; recovery is deferred. + let p = child_of(&tc, a, &[4]); + assert_eq!(tc.arena.node(p).key, vec![4]); + assert_eq!(tc.arena.node(b).key, vec![5]); + assert!( + tc.arena + .device_value(b, FULL) + .equal(&Tensor::from_slice(&[14i64])) + ); + let [ + CacheAction::FreeDeviceKV(duplicates_head), + CacheAction::RecoverSwaWithLockedFull { + node_id, + kept_full, + incoming_full, + }, + CacheAction::FreeDeviceKV(duplicates_tail), + ] = result.cache_actions.as_slice() + else { + panic!( + "expected FreeDeviceKV, RecoverSwaWithLockedFull, FreeDeviceKV, got {:?}", + action_kinds(&result.cache_actions) + ); + }; + assert!(duplicates_head[0].equal(&Tensor::from_slice(&[20i64, 21, 22]))); + assert_eq!(*node_id, tc.arena.node(b).id); + assert!(kept_full.equal(&Tensor::from_slice(&[14i64]))); + assert!(incoming_full.equal(&Tensor::from_slice(&[24i64]))); + assert!(duplicates_tail[0].equal(&Tensor::from_slice(&[23i64]))); +} + +#[test] +fn insert_overlap_prev_prefix_strictly_inside_the_node_still_recovers() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4], + &[10, 11, 12, 13], + 0, + 0, + )); + let root = tc.arena.root(); + let node = child_of(&tc, root, &[1]); + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4], + &[20, 21, 22, 23], + /* prev_prefix_len = */ 2, + /* swa_evicted_seqlen = */ 0, + )); + // prev covers only part of the node, so the tombstone recovery runs. + assert!( + tc.arena + .device_value(node, FULL) + .equal(&Tensor::from_slice(&[20i64, 21, 22, 23])) + ); + assert!(result.cache_actions.iter().any(|action| matches!( + action, + CacheAction::SwaRebuild { node_id, .. } if *node_id == tc.arena.node(node).id + ))); +} + +#[test] +#[should_panic(expected = "swa_evicted_seqlen must be page-aligned")] +fn insert_overlap_rejects_a_page_misaligned_boundary() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 2); + tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0)); + tc.insert(&insert_params_swa( + &vec![1, 2], + &[20, 21], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 1, + )); +} + +#[test] +fn reinsert_after_full_eviction_rebuilds_swa_from_the_fresh_kv() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0)); + let root = tc.arena.root(); + let leaf = child_of(&tc, root, &[1]); + evict_full(&mut tc, leaf, /* remaining_size = */ 0); + let result = tc.insert(&insert_params_swa(&vec![1, 2, 3], &[20, 21, 22], 0, 0)); + assert_eq!(result.prefix_len, 3); + assert!( + tc.arena + .device_value(leaf, FULL) + .equal(&Tensor::from_slice(&[20i64, 21, 22])) + ); + let [ + CacheAction::SwaRebuild { + node_id, + source_value, + }, + ] = result.cache_actions.as_slice() + else { + panic!("expected one SwaRebuild action"); + }; + assert_eq!(*node_id, tc.arena.node(leaf).id); + assert!(source_value.equal(&Tensor::from_slice(&[20i64, 21, 22]))); +} + +#[test] +fn reinsert_straddling_the_boundary_splits_before_the_rebuild() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4], + &[10, 11, 12, 13], + 0, + 0, + )); + let root = tc.arena.root(); + let node = child_of(&tc, root, &[1]); + evict_full(&mut tc, node, /* remaining_size = */ 0); + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4], + &[20, 21, 22, 23], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 2, + )); + // The out-of-window head stays an SWA tombstone on the split-off parent. + let parent = child_of(&tc, root, &[1]); + assert_eq!(tc.arena.node(parent).key, vec![1, 2]); + assert_eq!(tc.arena.node(node).key, vec![3, 4]); + assert!( + tc.arena + .device_value(parent, FULL) + .equal(&Tensor::from_slice(&[20i64, 21])) + ); + assert!(!tc.arena.has_device_value(parent, SWA)); + let [ + CacheAction::SwaRebuild { + node_id, + source_value, + }, + ] = result.cache_actions.as_slice() + else { + panic!("expected one SwaRebuild action"); + }; + assert_eq!(*node_id, tc.arena.node(node).id); + assert!(source_value.equal(&Tensor::from_slice(&[22i64, 23]))); +} + +#[test] +fn reinsert_straddling_the_boundary_at_a_second_level_node_splits_before_the_rebuild() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0)); + tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4, 5], + &[10, 11, 12, 13, 14], + 0, + 0, + )); + let root = tc.arena.root(); + let a = child_of(&tc, root, &[1]); + let b = child_of(&tc, a, &[4]); + evict_full(&mut tc, b, /* remaining_size = */ 0); + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4, 5], + &[20, 21, 22, 23, 24], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 4, + )); + // The unevicted b splits at its node-relative offset 1; only the + // in-window tail is rebuilt. + let p = child_of(&tc, a, &[4]); + assert_eq!(tc.arena.node(p).key, vec![4]); + assert_eq!(tc.arena.node(b).key, vec![5]); + assert!( + tc.arena + .device_value(p, FULL) + .equal(&Tensor::from_slice(&[23i64])) + ); + assert!( + tc.arena + .device_value(b, FULL) + .equal(&Tensor::from_slice(&[24i64])) + ); + assert!(!tc.arena.has_device_value(p, SWA)); + assert!(result.cache_actions.iter().any(|action| matches!( + action, + CacheAction::SwaRebuild { node_id, source_value } + if *node_id == tc.arena.node(b).id && source_value.equal(&Tensor::from_slice(&[24i64])) + ))); + assert!(!result.cache_actions.iter().any(|action| matches!( + action, + CacheAction::SwaRebuild { node_id, .. } if *node_id == tc.arena.node(p).id + ))); +} + +#[test] +fn reinsert_entirely_outside_the_window_skips_the_rebuild() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0)); + let root = tc.arena.root(); + let leaf = child_of(&tc, root, &[1]); + evict_full(&mut tc, leaf, /* remaining_size = */ 0); + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3], + &[20, 21, 22], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 3, + )); + assert!(result.cache_actions.is_empty()); + assert!(!tc.arena.has_device_value(leaf, SWA)); +} + +#[test] +fn reinsert_with_live_swa_skips_recovery() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0)); + let root = tc.arena.root(); + let leaf = child_of(&tc, root, &[1]); + tc.set_component_device_value( + tc.arena.node(leaf).id, + SWA, + Tensor::from_slice(&[50i64, 51, 52]), + ); + evict_full(&mut tc, leaf, /* remaining_size = */ 0); + let result = tc.insert(&insert_params_swa(&vec![1, 2, 3], &[20, 21, 22], 0, 0)); + // The SWA value is already live: no rebuild is emitted. + assert!(result.cache_actions.is_empty()); + assert!( + tc.arena + .device_value(leaf, SWA) + .equal(&Tensor::from_slice(&[50i64, 51, 52])) + ); +} + +#[test] +fn reinsert_boundary_at_the_node_start_rebuilds_the_whole_node() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0)); + tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4, 5], + &[30, 31, 32, 13, 14], + 0, + 0, + )); + let root = tc.arena.root(); + let a = child_of(&tc, root, &[1]); + let b = child_of(&tc, a, &[4]); + evict_full(&mut tc, b, /* remaining_size = */ 3); + // a is an out-of-window tombstone (all duplicate); b unevicts and + // rebuilds in full, no split. + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4, 5], + &[20, 21, 22, 23, 24], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 3, + )); + assert_eq!(tc.arena.node(b).key, vec![4, 5]); + let [ + CacheAction::FreeDeviceKV(duplicates), + CacheAction::SwaRebuild { + node_id, + source_value, + }, + ] = result.cache_actions.as_slice() + else { + panic!( + "expected FreeDeviceKV then SwaRebuild, got {:?}", + action_kinds(&result.cache_actions) + ); + }; + assert!(duplicates[0].equal(&Tensor::from_slice(&[20i64, 21, 22]))); + assert_eq!(*node_id, tc.arena.node(b).id); + assert!(source_value.equal(&Tensor::from_slice(&[23i64, 24]))); +} + +#[test] +fn walk_split_redistributes_the_live_swa_value() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4], + &[10, 11, 12, 13], + 0, + 0, + )); + let root = tc.arena.root(); + let node = child_of(&tc, root, &[1]); + tc.set_component_device_value( + tc.arena.node(node).id, + SWA, + Tensor::from_slice(&[50i64, 51, 52, 53]), + ); + let result = tc.insert(&insert_params_swa(&vec![1, 2, 9], &[20, 21, 29], 0, 0)); + assert_eq!(result.prefix_len, 2); + let parent = child_of(&tc, root, &[1]); + let leaf = child_of(&tc, parent, &[9]); + // The split slices the SWA value alongside the Full value; both sides + // stay in the SWA device LRU. + assert!( + tc.arena + .device_value(parent, SWA) + .equal(&Tensor::from_slice(&[50i64, 51])) + ); + assert!( + tc.arena + .device_value(node, SWA) + .equal(&Tensor::from_slice(&[52i64, 53])) + ); + assert!(tc.device_lru_list(SWA).in_list(Some(parent))); + assert!(tc.device_lru_list(SWA).in_list(Some(node))); + let [ + CacheAction::FreeDeviceKV(duplicates), + CacheAction::SwaRebuild { + node_id, + source_value, + }, + ] = result.cache_actions.as_slice() + else { + panic!("expected FreeDeviceKV then SwaRebuild"); + }; + assert!(duplicates[0].equal(&Tensor::from_slice(&[20i64, 21]))); + assert_eq!(*node_id, tc.arena.node(leaf).id); + assert!(source_value.equal(&Tensor::from_slice(&[29i64]))); +} + +#[test] +fn redistribute_on_node_split_slices_host_values_and_parks_tombstones() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0)); + let root = tc.arena.root(); + let node = child_of(&tc, root, &[1]); + tc.arena + .set_host_value(node, SWA, Tensor::from_slice(&[70i64, 71])); + tc.host_lru_list_mut(SWA).insert_mru(node); + let (parent, action) = tc.split_node_(node, /* split_len = */ 1); + assert!(action.is_none()); + assert!( + tc.arena + .host_value(parent, SWA) + .equal(&Tensor::from_slice(&[70i64])) + ); + assert!( + tc.arena + .host_value(node, SWA) + .equal(&Tensor::from_slice(&[71i64])) + ); + // Both sides are device tombstones: the parent joins the host LRU, the + // child stays listed. + assert!(tc.host_lru_list(SWA).in_list(Some(parent))); + assert!(tc.host_lru_list(SWA).in_list(Some(node))); +} + +#[test] +fn redistribute_on_node_split_keeps_device_valued_sides_off_the_host_lru() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0)); + let root = tc.arena.root(); + let node = child_of(&tc, root, &[1]); + tc.set_component_device_value( + tc.arena.node(node).id, + SWA, + Tensor::from_slice(&[50i64, 51]), + ); + tc.arena + .set_host_value(node, SWA, Tensor::from_slice(&[70i64, 71])); + tc.arena + .node_mut(node) + .set_lock_ref_(ValueSlotIdx::device(SWA), 2); + let (parent, _) = tc.split_node_(node, /* split_len = */ 1); + // Device-valued sides slice both tiers and inherit the SWA lock_ref; + // neither enters the host LRU. + assert!( + tc.arena + .device_value(parent, SWA) + .equal(&Tensor::from_slice(&[50i64])) + ); + assert!( + tc.arena + .host_value(parent, SWA) + .equal(&Tensor::from_slice(&[70i64])) + ); + assert_eq!(tc.arena.device_lock_ref(parent, SWA), 2); + assert_eq!(tc.arena.device_lock_ref(node, SWA), 2); + assert_eq!(tc.host_lru_list(SWA).len(), 0); +} + +#[test] +fn redistribute_on_node_split_lists_an_unlisted_tombstone_child() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0)); + let root = tc.arena.root(); + let node = child_of(&tc, root, &[1]); + tc.arena + .set_host_value(node, SWA, Tensor::from_slice(&[70i64, 71])); + let (parent, _) = tc.split_node_(node, /* split_len = */ 1); + // Neither side was in the host LRU; both device tombstones join it. + assert!(tc.host_lru_list(SWA).in_list(Some(parent))); + assert!(tc.host_lru_list(SWA).in_list(Some(node))); +} + +#[test] +fn acquire_lock_walks_until_the_window_fills_and_stamps_the_crossing_node() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + store_swa_device(&mut tc, a); + store_swa_device(&mut tc, b); + store_swa_device(&mut tc, c); + let result = swa_component(2).acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + assert_eq!(tc.arena.device_lock_ref(c, SWA), 1); + assert_eq!(tc.arena.device_lock_ref(b, SWA), 1); + assert_eq!(tc.arena.device_lock_ref(a, SWA), 0); + // The walk fills the 2-atom window at b; b carries the first minted uuid. + assert_eq!(result.swa_uuid_for_lock, Some(2)); + assert_eq!(result.swa_uuid_for_host_lock, None); + assert_eq!(node_swa_uuid(&tc, b), Some(2)); + assert_eq!(node_swa_uuid(&tc, c), None); + assert_eq!(tc.swa_evictable_size(), 1); + assert_eq!(tc.swa_protected_size(), 2); +} + +#[test] +fn acquire_lock_overshooting_the_window_stops_at_the_crossing_node() { + let mut tc = swa_core(/* window = */ 3, /* page_size = */ 1); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let b = tc + .arena + .alloc_child( + a, + /* key = */ vec![3, 4], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let c = tc + .arena + .alloc_child( + b, + /* key = */ vec![5, 6], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + store_swa_device(&mut tc, a); + store_swa_device(&mut tc, b); + store_swa_device(&mut tc, c); + let result = swa_component(3).acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + // The 2-atom nodes overshoot the 3-atom window at b (2 -> 4): the walk + // stops there, stamps b, and leaves a untouched. + assert_eq!(tc.arena.device_lock_ref(c, SWA), 1); + assert_eq!(tc.arena.device_lock_ref(b, SWA), 1); + assert_eq!(tc.arena.device_lock_ref(a, SWA), 0); + assert_eq!(result.swa_uuid_for_lock, Some(2)); + assert_eq!(node_swa_uuid(&tc, b), Some(2)); + assert_eq!(node_swa_uuid(&tc, a), None); + assert_eq!(tc.swa_evictable_size(), 2); + assert_eq!(tc.swa_protected_size(), 4); +} + +#[test] +fn acquire_lock_reuses_the_stamped_uuid_and_shifts_sizes_once() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + store_swa_device(&mut tc, a); + store_swa_device(&mut tc, b); + store_swa_device(&mut tc, c); + let swa = swa_component(2); + let first = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + let second = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + assert_eq!(second.swa_uuid_for_lock, first.swa_uuid_for_lock); + assert_eq!(tc.arena.device_lock_ref(c, SWA), 2); + assert_eq!(tc.arena.device_lock_ref(b, SWA), 2); + assert_eq!(tc.swa_evictable_size(), 1); + assert_eq!(tc.swa_protected_size(), 2); +} + +#[test] +fn acquire_lock_skips_tombstones_and_records_them() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + store_swa_device(&mut tc, a); + store_swa_device(&mut tc, c); + let result = swa_component(2).acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + // The valueless b is recorded and skipped; the window fills at a. + assert_eq!(tc.arena.device_lock_ref(c, SWA), 1); + assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(a, SWA), 1); + assert_eq!(result.skip_lock_node_ids[&SWA].len(), 1); + assert!(result.skip_lock_node_ids[&SWA].contains(&tc.arena.node(b).id)); + assert!(result.swa_uuid_for_lock.is_some()); + assert_eq!(node_swa_uuid(&tc, a), result.swa_uuid_for_lock); +} + +#[test] +fn acquire_lock_under_the_window_reaches_the_root_without_a_uuid() { + let mut tc = swa_core(/* window = */ 100, /* page_size = */ 1); + let [a, b] = chain(&mut tc); + store_swa_device(&mut tc, a); + store_swa_device(&mut tc, b); + let result = swa_component(100).acquire_component_lock( + &mut tc, + b, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + assert_eq!(result.swa_uuid_for_lock, None); + assert_eq!(tc.arena.device_lock_ref(a, SWA), 1); + assert_eq!(tc.arena.device_lock_ref(b, SWA), 1); + assert_eq!(tc.swa_evictable_size(), 0); + assert_eq!(tc.swa_protected_size(), 2); +} + +#[test] +fn inc_lock_ref_runs_full_and_swa_walks_together() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + store_swa_device(&mut tc, a); + store_swa_device(&mut tc, b); + store_swa_device(&mut tc, c); + let result = tc.inc_lock_ref(tc.arena.node(c).id); + // FULL sees a valueless path (skip segment only); SWA locks its window. + assert_eq!(result.delta, Some(0)); + assert_eq!(result.skip_lock_node_ids[&FULL].len(), 3); + assert!(result.swa_uuid_for_lock.is_some()); + assert_eq!(tc.arena.device_lock_ref(c, SWA), 1); + assert_eq!(tc.arena.device_lock_ref(b, SWA), 1); + assert_eq!(tc.arena.device_lock_ref(a, SWA), 0); +} + +#[test] +fn inc_host_lock_ref_runs_full_and_swa_host_arms_together() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + for node in [a, b, c] { + set_swa_host(&mut tc, node); + } + tc.arena + .set_host_value(c, FULL, Tensor::from_slice(&[0i64])); + let result = tc.inc_host_lock_ref(tc.arena.node(c).id); + // FULL pins only the anchor; SWA walks its host window up to b. + assert_eq!(tc.arena.host_lock_ref(c, FULL), 1); + assert_eq!(tc.arena.host_lock_ref(b, FULL), 0); + assert_eq!(tc.arena.host_lock_ref(c, SWA), 1); + assert_eq!(tc.arena.host_lock_ref(b, SWA), 1); + assert_eq!(tc.arena.host_lock_ref(a, SWA), 0); + assert!(result.swa_uuid_for_host_lock.is_some()); + // The release replays the acquire's uuid and unwinds both arms. + let params = DecLockRefParams { + swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, + skip_lock_node_ids: result.skip_lock_node_ids, + ..Default::default() + }; + tc.dec_host_lock_ref(tc.arena.node(c).id, Some(¶ms)); + assert_eq!(tc.arena.host_lock_ref(c, FULL), 0); + assert_eq!(tc.arena.host_lock_ref(c, SWA), 0); + assert_eq!(tc.arena.host_lock_ref(b, SWA), 0); + // Only the dispatcher's final leaf-set pass sees the fully-unlocked anchor. + assert!(tc.evictable_host_leaves.contains(c)); +} + +#[test] +fn dec_host_lock_ref_with_the_inner_uuid_leaves_an_outer_window_pinned() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + for node in [a, b, c] { + set_swa_host(&mut tc, node); + } + // Overlapping host windows: {c, b} stamps its uuid at b, {b, a} at a. + let inner = tc.inc_host_lock_ref(tc.arena.node(c).id); + tc.inc_host_lock_ref(tc.arena.node(b).id); + assert!(inner.swa_uuid_for_host_lock.is_some()); + assert_eq!(tc.arena.host_lock_ref(c, SWA), 1); + assert_eq!(tc.arena.host_lock_ref(b, SWA), 2); + assert_eq!(tc.arena.host_lock_ref(a, SWA), 1); + // Releasing the inner window with its own uuid stops at b; the outer + // window's lock above the boundary survives. + let params = DecLockRefParams { + swa_uuid_for_host_lock: inner.swa_uuid_for_host_lock, + skip_lock_node_ids: inner.skip_lock_node_ids, + ..Default::default() + }; + tc.dec_host_lock_ref(tc.arena.node(c).id, Some(¶ms)); + assert_eq!(tc.arena.host_lock_ref(c, SWA), 0); + assert_eq!(tc.arena.host_lock_ref(b, SWA), 1); + assert_eq!(tc.arena.host_lock_ref(a, SWA), 1); +} + +#[test] +fn acquire_host_lock_walks_until_the_window_fills_and_stamps_the_host_uuid() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + for node in [a, b, c] { + set_swa_host(&mut tc, node); + tc.host_lru_list_mut(SWA).insert_mru(node); + } + let result = swa_component(2).acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ true, + ); + assert_eq!(tc.arena.host_lock_ref(c, SWA), 1); + assert_eq!(tc.arena.host_lock_ref(b, SWA), 1); + assert_eq!(tc.arena.host_lock_ref(a, SWA), 0); + // The window fills at b; b carries the host uuid and leaves the host LRU. + assert_eq!(result.swa_uuid_for_host_lock, Some(2)); + assert_eq!(result.swa_uuid_for_lock, None); + assert_eq!(node_swa_host_uuid(&tc, b), Some(2)); + assert_eq!(node_swa_host_uuid(&tc, c), None); + assert!(!tc.host_lru_list(SWA).in_list(Some(c))); + assert!(!tc.host_lru_list(SWA).in_list(Some(b))); + assert!(tc.host_lru_list(SWA).in_list(Some(a))); + // Host locks never touch the device tier or its sizes. + assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); + assert_eq!(tc.swa_evictable_size(), 0); + assert_eq!(tc.swa_protected_size(), 0); +} + +#[test] +fn acquire_host_lock_reuses_the_stamped_uuid_and_skips_unlisted_nodes() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + for node in [a, b, c] { + set_swa_host(&mut tc, node); + } + let swa = swa_component(2); + let first = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ true, + ); + let second = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ true, + ); + assert_eq!(second.swa_uuid_for_host_lock, first.swa_uuid_for_host_lock); + assert_eq!(tc.arena.host_lock_ref(c, SWA), 2); + assert_eq!(tc.arena.host_lock_ref(b, SWA), 2); + assert_eq!(tc.arena.host_lock_ref(a, SWA), 0); +} + +#[test] +fn acquire_host_lock_skips_host_tombstones_and_records_them() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + set_swa_host(&mut tc, a); + set_swa_host(&mut tc, c); + let result = swa_component(2).acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ true, + ); + assert_eq!(tc.arena.host_lock_ref(c, SWA), 1); + assert_eq!(tc.arena.host_lock_ref(b, SWA), 0); + assert_eq!(tc.arena.host_lock_ref(a, SWA), 1); + assert_eq!(result.skip_lock_node_ids[&SWA].len(), 1); + assert!(result.skip_lock_node_ids[&SWA].contains(&tc.arena.node(b).id)); + assert_eq!(node_swa_host_uuid(&tc, a), result.swa_uuid_for_host_lock); + assert!(result.swa_uuid_for_host_lock.is_some()); +} + +#[test] +fn acquire_host_lock_under_the_window_reaches_the_root_without_a_uuid() { + let mut tc = swa_core(/* window = */ 100, /* page_size = */ 1); + let [a, b] = chain(&mut tc); + set_swa_host(&mut tc, a); + set_swa_host(&mut tc, b); + let result = swa_component(100).acquire_component_lock( + &mut tc, + b, + IncLockRefResult::default(), + /* lock_host = */ true, + ); + assert_eq!(result.swa_uuid_for_host_lock, None); + assert_eq!(tc.arena.host_lock_ref(a, SWA), 1); + assert_eq!(tc.arena.host_lock_ref(b, SWA), 1); +} + +#[test] +fn acquire_host_lock_delists_only_on_the_host_ref_transition() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + for node in [a, b, c] { + store_swa_device(&mut tc, node); + set_swa_host(&mut tc, node); + tc.host_lru_list_mut(SWA).insert_mru(node); + } + let swa = swa_component(2); + // A held device lock must not stand in for the host ref transition. + let _ = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + let _ = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ true, + ); + assert!(!tc.host_lru_list(SWA).in_list(Some(c))); + assert!(!tc.host_lru_list(SWA).in_list(Some(b))); + assert!(tc.host_lru_list(SWA).in_list(Some(a))); + // A re-listed node with a live host lock stays listed on re-acquire. + tc.host_lru_list_mut(SWA).insert_mru(b); + let _ = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ true, + ); + assert!(tc.host_lru_list(SWA).in_list(Some(b))); + assert_eq!(tc.arena.host_lock_ref(c, SWA), 2); + assert_eq!(tc.arena.host_lock_ref(b, SWA), 2); +} + +#[test] +fn acquire_host_lock_stamps_the_host_tier_uuid_field_only() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + for node in [a, b, c] { + set_swa_host(&mut tc, node); + } + let result = swa_component(2).acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ true, + ); + // The boundary uuid lands on the host-tier field; the device field stays clear. + assert_eq!(result.swa_uuid_for_host_lock, Some(2)); + assert_eq!(node_swa_host_uuid(&tc, b), Some(2)); + assert_eq!(node_swa_uuid(&tc, b), None); + assert_eq!(node_swa_host_uuid(&tc, c), None); + assert_eq!(node_swa_uuid(&tc, c), None); + let _ = a; +} + +#[test] +fn device_and_host_lock_walks_mint_independent_uuids() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + for node in [a, b, c] { + store_swa_device(&mut tc, node); + set_swa_host(&mut tc, node); + } + let swa = swa_component(2); + let device = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + let host = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ true, + ); + assert_eq!(device.swa_uuid_for_lock, Some(2)); + assert_eq!(host.swa_uuid_for_host_lock, Some(3)); + assert_eq!(node_swa_uuid(&tc, b), Some(2)); + assert_eq!(node_swa_host_uuid(&tc, b), Some(3)); +} + +#[test] +fn eviction_priority_is_zero_for_leaf_one_for_internal() { + let swa = swa_component(4); + assert_eq!( + >>::eviction_priority(&swa, true), + 0 + ); + assert_eq!( + >>::eviction_priority(&swa, false), + 1 + ); +} + +#[test] +fn evict_component_device_frees_the_full_indices_and_tombstones_swa() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0)); + let root = tc.arena.root(); + let node = child_of(&tc, root, &[1]); + tc.set_component_device_value( + tc.arena.node(node).id, + SWA, + Tensor::from_slice(&[50i64, 51]), + ); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + let (freed, host_freed) = swa_component(4).evict_component( + &mut tc, + node, + &mut device_frees, + &mut host_frees, + EvictLayer::Device, + ); + assert_eq!((freed, host_freed), (2, 0)); + // The freed indices are the FULL slice (SWA slots pair through it); + // the FULL value itself stays on the node. + assert!(device_frees[&SWA][0].equal(&Tensor::from_slice(&[10i64, 11]))); + assert!(!tc.arena.has_device_value(node, SWA)); + assert!( + tc.arena + .device_value(node, FULL) + .equal(&Tensor::from_slice(&[10i64, 11])) + ); + assert_eq!(tc.swa_evictable_size(), 0); +} + +#[test] +fn evict_component_device_parks_a_remaining_host_value() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0)); + let root = tc.arena.root(); + let node = child_of(&tc, root, &[1]); + tc.set_component_device_value( + tc.arena.node(node).id, + SWA, + Tensor::from_slice(&[50i64, 51]), + ); + set_swa_host(&mut tc, node); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + swa_component(4).evict_component( + &mut tc, + node, + &mut device_frees, + &mut host_frees, + EvictLayer::Device, + ); + assert!(tc.arena.has_host_value(node, SWA)); + assert!(tc.host_lru_list(SWA).in_list(Some(node))); + assert!(host_frees.is_empty()); +} + +#[test] +fn evict_component_host_frees_and_delists_the_host_value() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0)); + let root = tc.arena.root(); + let node = child_of(&tc, root, &[1]); + set_swa_host(&mut tc, node); + tc.host_lru_list_mut(SWA).insert_mru(node); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + let (freed, host_freed) = swa_component(4).evict_component( + &mut tc, + node, + &mut device_frees, + &mut host_frees, + EvictLayer::Host, + ); + assert_eq!((freed, host_freed), (0, 2)); + assert!(host_frees[&SWA][0].equal(&Tensor::from_slice(&[0i64, 0]))); + assert!(!tc.arena.has_host_value(node, SWA)); + assert!(!tc.host_lru_list(SWA).in_list(Some(node))); + assert!(device_frees.is_empty()); +} + +#[test] +fn release_lock_returns_the_window_to_evictable() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + store_swa_device(&mut tc, a); + store_swa_device(&mut tc, b); + store_swa_device(&mut tc, c); + let swa = swa_component(2); + let result = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + let params = DecLockRefParams { + swa_uuid_for_lock: result.swa_uuid_for_lock, + swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, + skip_lock_node_ids: result.skip_lock_node_ids, + }; + swa.release_component_lock(&mut tc, c, Some(¶ms), /* lock_host = */ false); + assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(a, SWA), 0); + assert_eq!(tc.swa_evictable_size(), 3); + assert_eq!(tc.swa_protected_size(), 0); +} + +#[test] +fn release_lock_keeps_sizes_while_other_locks_remain() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + store_swa_device(&mut tc, a); + store_swa_device(&mut tc, b); + store_swa_device(&mut tc, c); + let swa = swa_component(2); + let first = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + let _ = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + let params = DecLockRefParams { + swa_uuid_for_lock: first.swa_uuid_for_lock, + swa_uuid_for_host_lock: first.swa_uuid_for_host_lock, + skip_lock_node_ids: first.skip_lock_node_ids, + }; + swa.release_component_lock(&mut tc, c, Some(¶ms), /* lock_host = */ false); + assert_eq!(tc.arena.device_lock_ref(c, SWA), 1); + assert_eq!(tc.arena.device_lock_ref(b, SWA), 1); + assert_eq!(tc.swa_evictable_size(), 1); + assert_eq!(tc.swa_protected_size(), 2); +} + +#[test] +fn release_lock_replays_the_tombstone_skips() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + store_swa_device(&mut tc, a); + store_swa_device(&mut tc, c); + let swa = swa_component(2); + let result = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + // b gained a device value AFTER the acquire recorded it as a tombstone. + store_swa_device(&mut tc, b); + let params = DecLockRefParams { + swa_uuid_for_lock: result.swa_uuid_for_lock, + swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, + skip_lock_node_ids: result.skip_lock_node_ids, + }; + swa.release_component_lock(&mut tc, c, Some(¶ms), /* lock_host = */ false); + assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(a, SWA), 0); + assert_eq!(tc.swa_evictable_size(), 3); + assert_eq!(tc.swa_protected_size(), 0); +} + +#[test] +fn release_lock_stops_at_the_window_uuid() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + store_swa_device(&mut tc, a); + store_swa_device(&mut tc, b); + store_swa_device(&mut tc, c); + let swa = swa_component(2); + let result = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + // A manually locked ancestor above the window must stay untouched. + tc.arena + .node_mut(a) + .set_lock_ref_(ValueSlotIdx::device(SWA), 1); + let params = DecLockRefParams { + swa_uuid_for_lock: result.swa_uuid_for_lock, + swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, + skip_lock_node_ids: result.skip_lock_node_ids, + }; + swa.release_component_lock(&mut tc, c, Some(¶ms), /* lock_host = */ false); + assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(a, SWA), 1); +} + +#[test] +fn release_host_lock_stops_at_the_host_uuid_boundary() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + for node in [a, b, c] { + set_swa_host(&mut tc, node); + tc.host_lru_list_mut(SWA).insert_mru(node); + } + let swa = swa_component(2); + let result = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ true, + ); + // A second request holds its own host lock on a, above the boundary b. + let _ = swa.acquire_component_lock( + &mut tc, + a, + IncLockRefResult::default(), + /* lock_host = */ true, + ); + let params = DecLockRefParams { + swa_uuid_for_lock: result.swa_uuid_for_lock, + swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, + skip_lock_node_ids: result.skip_lock_node_ids, + }; + swa.release_component_lock(&mut tc, c, Some(¶ms), /* lock_host = */ true); + assert_eq!(tc.arena.host_lock_ref(c, SWA), 0); + assert_eq!(tc.arena.host_lock_ref(b, SWA), 0); + assert_eq!(tc.arena.host_lock_ref(a, SWA), 1); + assert!(!tc.host_lru_list(SWA).in_list(Some(a))); +} + +#[test] +fn release_lock_without_params_passes_over_an_unlocked_middle_node() { + let mut tc = swa_core(/* window = */ 1, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + store_swa_device(&mut tc, a); + store_swa_device(&mut tc, b); + store_swa_device(&mut tc, c); + let swa = swa_component(1); + // The 1-atom window locks only the acquired node: c and a, never b. + let _ = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + let _ = swa.acquire_component_lock( + &mut tc, + a, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + swa.release_component_lock( + &mut tc, c, /* params = */ None, /* lock_host = */ false, + ); + assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(a, SWA), 0); + assert_eq!(tc.swa_evictable_size(), 3); + assert_eq!(tc.swa_protected_size(), 0); +} + +#[test] +fn release_host_lock_reparks_tombstoned_host_nodes() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + for node in [a, b, c] { + set_swa_host(&mut tc, node); + tc.host_lru_list_mut(SWA).insert_mru(node); + } + let swa = swa_component(2); + let result = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ true, + ); + let params = DecLockRefParams { + swa_uuid_for_lock: result.swa_uuid_for_lock, + swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, + skip_lock_node_ids: result.skip_lock_node_ids, + }; + swa.release_component_lock(&mut tc, c, Some(¶ms), /* lock_host = */ true); + assert_eq!(tc.arena.host_lock_ref(c, SWA), 0); + assert_eq!(tc.arena.host_lock_ref(b, SWA), 0); + assert!(tc.host_lru_list(SWA).in_list(Some(c))); + assert!(tc.host_lru_list(SWA).in_list(Some(b))); + assert!(tc.host_lru_list(SWA).in_list(Some(a))); +} + +#[test] +fn inc_then_dec_lock_ref_roundtrips_with_dec_params() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + store_swa_device(&mut tc, a); + store_swa_device(&mut tc, b); + store_swa_device(&mut tc, c); + let result = tc.inc_lock_ref(tc.arena.node(c).id); + let params = DecLockRefParams { + swa_uuid_for_lock: result.swa_uuid_for_lock, + swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, + skip_lock_node_ids: result.skip_lock_node_ids, + }; + tc.dec_lock_ref( + tc.arena.node(c).id, + Some(¶ms), + /* skip_swa = */ false, + ); + assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); + assert_eq!(tc.swa_evictable_size(), 3); + assert_eq!(tc.swa_protected_size(), 0); +} + +#[test] +fn dec_swa_lock_only_releases_swa_while_full_stays_locked() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + for node in [a, b, c] { + store_swa_device(&mut tc, node); + let len = tc.arena.node(node).key.atom_len(); + tc.arena + .set_device_value(node, FULL, Tensor::from_slice(&vec![9i64; len])); + } + // Fund FULL's evictable counter for its lock walk (raw slot sets skip it). + tc.component_state_mut(FULL).evictable_size = 3; + let result = tc.inc_lock_ref(tc.arena.node(c).id); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + tc.dec_swa_lock_only( + tc.arena.node(c).id, + result.swa_uuid_for_lock, + &mut device_frees, + &mut host_frees, + ); + // SWA is early-released; the FULL locks on the path stay. + assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(c, FULL), 1); + assert_eq!(tc.arena.device_lock_ref(b, FULL), 1); + // Full still locks the nodes, so nothing is device-leaf evictable. + assert!(device_frees.is_empty()); + assert_eq!(tc.swa_evictable_size(), 3); + assert_eq!(tc.swa_protected_size(), 0); +} + +#[test] +fn dec_swa_lock_only_evicts_a_fully_unlocked_device_leaf() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + for node in [a, b, c] { + store_swa_device(&mut tc, node); + let len = tc.arena.node(node).key.atom_len(); + tc.arena + .set_device_value(node, FULL, Tensor::from_slice(&vec![9i64; len])); + } + let swa = swa_component(2); + let result = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + tc.dec_swa_lock_only( + tc.arena.node(c).id, + result.swa_uuid_for_lock, + &mut device_frees, + &mut host_frees, + ); + // The fully unlocked leaf c is device-evicted on release; b keeps its + // SWA value because its child still holds FULL KV. + assert!(!tc.arena.has_device_value(c, SWA)); + assert!(device_frees[&SWA][0].equal(&Tensor::from_slice(&[9i64]))); + assert!(tc.arena.has_device_value(b, SWA)); + assert_eq!(tc.swa_evictable_size(), 2); + assert_eq!(tc.swa_protected_size(), 0); +} + +#[test] +fn dec_swa_lock_only_is_a_noop_without_the_swa_component() { + let mut tc: UnifiedTreeCore> = + UnifiedTreeCore::new(CacheInitParams::default(), vec![FULL]); + let root = tc.arena.root(); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + tc.dec_swa_lock_only( + tc.arena.node(root).id, + None, + &mut device_frees, + &mut host_frees, + ); + assert!(device_frees.is_empty()); +} + +#[test] +fn release_window_lock_breaks_on_a_tombstone_carrying_the_uuid() { + let mut tc = swa_core(/* window = */ 100, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + store_swa_device(&mut tc, a); + store_swa_device(&mut tc, c); + let swa = swa_component(100); + let _ = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + // A stale uuid on the tombstone b ends the walk before a. + tc.arena.node_mut(b).swa_uuid = Some(99); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + >>::release_window_lock( + &swa, + &mut tc, + c, + Some(99), + &mut device_frees, + &mut host_frees, + ); + assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(a, SWA), 1); +} + +#[test] +#[should_panic(expected = "release_window_lock is SWA-only")] +fn release_window_lock_panics_on_a_non_swa_component() { + use crate::components::full::FullComponent; + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let root = tc.arena.root(); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + >>::release_window_lock( + &FullComponent, + &mut tc, + root, + None, + &mut device_frees, + &mut host_frees, + ); +} + +#[test] +fn release_lock_skip_set_leaves_a_relocked_tombstone_credited() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + store_swa_device(&mut tc, a); + store_swa_device(&mut tc, c); + let swa = swa_component(2); + let first = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + // b regains a value and a second request locks it before the first + // release replays its tombstone skip set. + store_swa_device(&mut tc, b); + let _ = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + let params = DecLockRefParams { + swa_uuid_for_lock: first.swa_uuid_for_lock, + swa_uuid_for_host_lock: first.swa_uuid_for_host_lock, + skip_lock_node_ids: first.skip_lock_node_ids, + }; + swa.release_component_lock(&mut tc, c, Some(¶ms), /* lock_host = */ false); + assert_eq!(tc.arena.device_lock_ref(c, SWA), 1); + assert_eq!(tc.arena.device_lock_ref(b, SWA), 1); + assert_eq!(tc.arena.device_lock_ref(a, SWA), 0); + assert_eq!(tc.swa_evictable_size(), 1); + assert_eq!(tc.swa_protected_size(), 2); +} + +#[test] +fn release_lock_passes_over_uncredited_nodes_without_params() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + store_swa_device(&mut tc, a); + store_swa_device(&mut tc, b); + store_swa_device(&mut tc, c); + let swa = swa_component(2); + let _ = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + // No params: the walk crosses the never-credited a up to the root. + swa.release_component_lock( + &mut tc, c, /* params = */ None, /* lock_host = */ false, + ); + assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(a, SWA), 0); + assert_eq!(tc.swa_evictable_size(), 3); + assert_eq!(tc.swa_protected_size(), 0); +} + +#[test] +fn dec_swa_lock_only_releases_the_window_exactly_once() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + store_swa_device(&mut tc, a); + store_swa_device(&mut tc, b); + store_swa_device(&mut tc, c); + let swa = swa_component(2); + let first = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + let _ = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + tc.dec_swa_lock_only( + tc.arena.node(c).id, + first.swa_uuid_for_lock, + &mut device_frees, + &mut host_frees, + ); + // The second window still holds the lock: refs drop to 1, sizes stay. + assert_eq!(tc.arena.device_lock_ref(c, SWA), 1); + assert_eq!(tc.arena.device_lock_ref(b, SWA), 1); + assert_eq!(tc.swa_evictable_size(), 1); + assert_eq!(tc.swa_protected_size(), 2); + tc.dec_swa_lock_only( + tc.arena.node(c).id, + first.swa_uuid_for_lock, + &mut device_frees, + &mut host_frees, + ); + assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); + assert_eq!(tc.swa_evictable_size(), 3); + assert_eq!(tc.swa_protected_size(), 0); + let _ = a; +} + +#[test] +fn dec_swa_lock_only_leaves_out_of_window_swa_locks_alone() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + store_swa_device(&mut tc, a); + store_swa_device(&mut tc, b); + store_swa_device(&mut tc, c); + let swa = swa_component(2); + let result = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + // A holds a lock beyond the window (e.g. another request's window). + tc.arena + .node_mut(a) + .set_lock_ref_(ValueSlotIdx::device(SWA), 1); + tc.dec_evictable_size(SWA, 1); + tc.inc_protected_size(SWA, 1); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + tc.dec_swa_lock_only( + tc.arena.node(c).id, + result.swa_uuid_for_lock, + &mut device_frees, + &mut host_frees, + ); + // Only the SWA window is released; a's out-of-window lock survives. + assert_eq!(tc.arena.device_lock_ref(a, SWA), 1); + assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); + assert_eq!(tc.swa_protected_size(), 1); +} + +#[test] +fn release_window_lock_passes_over_an_unlocked_valued_node() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + store_swa_device(&mut tc, a); + store_swa_device(&mut tc, b); + store_swa_device(&mut tc, c); + let swa = swa_component(2); + let _ = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + // No uuid bound: the walk crosses the valued-but-unlocked a to the root. + swa.release_window_lock(&mut tc, c, None, &mut device_frees, &mut host_frees); + assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(a, SWA), 0); +} + +#[test] +fn release_window_lock_passes_over_a_mid_chain_tombstone_without_a_uuid() { + let mut tc = swa_core(/* window = */ 100, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + store_swa_device(&mut tc, a); + store_swa_device(&mut tc, c); + let swa = swa_component(100); + let _ = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + // No uuid bound: the walk crosses the mid-chain tombstone b and releases a. + swa.release_window_lock(&mut tc, c, None, &mut device_frees, &mut host_frees); + assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); + assert_eq!(tc.arena.device_lock_ref(a, SWA), 0); + assert_eq!(tc.swa_protected_size(), 0); +} + +#[test] +fn release_host_lock_does_not_repark_a_node_whose_host_value_was_taken() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a] = chain(&mut tc); + set_swa_host(&mut tc, a); + tc.host_lru_list_mut(SWA).insert_mru(a); + let swa = swa_component(2); + let result = swa.acquire_component_lock( + &mut tc, + a, + IncLockRefResult::default(), + /* lock_host = */ true, + ); + // The host value moved out while the lock was held; the node has no + // device value either, so the release has nothing to park. + let _ = tc.arena.take_host_value(a, SWA); + let params = DecLockRefParams { + swa_uuid_for_lock: result.swa_uuid_for_lock, + swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, + skip_lock_node_ids: result.skip_lock_node_ids, + }; + swa.release_component_lock(&mut tc, a, Some(¶ms), /* lock_host = */ true); + assert_eq!(tc.arena.host_lock_ref(a, SWA), 0); + assert!(!tc.host_lru_list(SWA).in_list(Some(a))); +} + +#[test] +fn release_host_lock_skips_reparking_device_valued_nodes() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + for node in [a, b, c] { + store_swa_device(&mut tc, node); + set_swa_host(&mut tc, node); + tc.host_lru_list_mut(SWA).insert_mru(node); + } + let swa = swa_component(2); + let result = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ true, + ); + let params = DecLockRefParams { + swa_uuid_for_lock: result.swa_uuid_for_lock, + swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, + skip_lock_node_ids: result.skip_lock_node_ids, + }; + swa.release_component_lock(&mut tc, c, Some(¶ms), /* lock_host = */ true); + // Device-valued nodes never re-park in the host LRU on host release. + assert!(!tc.host_lru_list(SWA).in_list(Some(c))); + assert!(!tc.host_lru_list(SWA).in_list(Some(b))); +} + +#[test] +fn release_host_lock_leaves_an_already_listed_node_listed() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain(&mut tc); + for node in [a, b, c] { + set_swa_host(&mut tc, node); + tc.host_lru_list_mut(SWA).insert_mru(node); + } + let swa = swa_component(2); + let result = swa.acquire_component_lock( + &mut tc, + c, + IncLockRefResult::default(), + /* lock_host = */ true, + ); + // Something re-listed b while the lock was held (e.g. a split re-park). + tc.host_lru_list_mut(SWA).insert_mru(b); + let params = DecLockRefParams { + swa_uuid_for_lock: result.swa_uuid_for_lock, + swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, + skip_lock_node_ids: result.skip_lock_node_ids, + }; + swa.release_component_lock(&mut tc, c, Some(¶ms), /* lock_host = */ true); + assert!(tc.host_lru_list(SWA).in_list(Some(b))); + assert!(tc.host_lru_list(SWA).in_list(Some(c))); + let _ = a; +} + +// A three-node chain built through insert (leaf sets maintained), with +// SWA device values stored on every node. +fn swa_evict_chain(tc: &mut UnifiedTreeCore>) -> [NodeIdx_; 3] { + tc.insert(&insert_params_swa(&vec![1], &[10], 0, 0)); + tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0)); + tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0)); + let root = tc.arena.root(); + let a = child_of(tc, root, &[1]); + let b = child_of(tc, a, &[2]); + let c = child_of(tc, b, &[3]); + for node in [a, b, c] { + store_swa_device(tc, node); + } + [a, b, c] +} + +fn swa_tracker() -> HashMap { + HashMap::from([(FULL, 0), (SWA, 0)]) +} + +#[test] +fn evict_walk_advances_one_allocator_mutation_per_call() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = swa_evict_chain(&mut tc); + let mut tracker = swa_tracker(); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + tc.evict_device_start(SWA, /* request_cnt = */ 100); + let (first, step) = tc.evict_device_next_node(SWA, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + // Each internal tombstone is its own step so the allocator can observe + // and reuse the freed slice before the walk mutates another node. + assert_eq!(first, None); + assert!(!tc.arena.has_device_value(a, SWA)); + assert!(tc.arena.has_device_value(b, SWA)); + assert!(tc.arena.has_device_value(c, SWA)); + assert_eq!(tracker[&SWA], 1); + assert_eq!(device_frees[&SWA].len(), 1); + assert!(device_frees[&SWA][0].equal(&Tensor::from_slice(&[10i64]))); + + let (second, step) = tc.evict_device_next_node(SWA, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + assert_eq!(second, None); + assert!(!tc.arena.has_device_value(b, SWA)); + assert!(tc.arena.has_device_value(c, SWA)); + assert!(tc.arena.has_device_value(a, FULL)); + assert!(tc.arena.has_device_value(b, FULL)); + assert_eq!(tracker[&SWA], 2); + assert_eq!(device_frees[&SWA].len(), 2); + assert!(device_frees[&SWA][1].equal(&Tensor::from_slice(&[11i64]))); + + let (third, step) = tc.evict_device_next_node(SWA, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + assert_eq!(third, Some(tc.arena.node(c).id)); + assert_eq!(tracker[&SWA], 2); + tc.evict_device_end(SWA); +} + +#[test] +fn evict_walk_stops_at_the_token_budget() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = swa_evict_chain(&mut tc); + let mut tracker = swa_tracker(); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + tc.evict_device_start(SWA, /* request_cnt = */ 1); + let (next, step) = tc.evict_device_next_node(SWA, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + // The first inline tombstone (a) fills the budget; b and c survive. + assert_eq!(next, None); + assert!(!tc.arena.has_device_value(a, SWA)); + assert!(tc.arena.has_device_value(b, SWA)); + assert!(tc.arena.has_device_value(c, SWA)); + assert_eq!(tracker[&SWA], 1); + tc.evict_device_end(SWA); +} + +#[test] +fn evict_walk_step_tracker_carries_only_the_deltas_over_the_baseline() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [_a, _b, _c] = swa_evict_chain(&mut tc); + // A non-zero baseline stands in for prior steps' evictions. + let baseline = HashMap::from([(FULL, 0), (SWA, 3)]); + tc.evict_device_start(SWA, /* request_cnt = */ 100); + let (next, step) = tc.evict_device_next_node(SWA, &baseline); + assert_eq!(next, None); + // One internal node tombstones: the step reports 1, not the running total. + assert_eq!(step.tracker[&SWA], 1); + tc.evict_device_end(SWA); +} + +#[test] +fn evict_walk_skips_locked_nodes() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = swa_evict_chain(&mut tc); + tc.arena + .node_mut(c) + .set_lock_ref_(ValueSlotIdx::device(SWA), 1); + let mut tracker = swa_tracker(); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + tc.evict_device_start(SWA, /* request_cnt = */ 100); + let (first, step) = tc.evict_device_next_node(SWA, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + assert_eq!(first, None); + assert_eq!(tracker[&SWA], 1); + let (second, step) = tc.evict_device_next_node(SWA, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + // a and b tombstone in separate steps; the locked c is invisible. + assert_eq!(second, None); + assert!(tc.arena.has_device_value(c, SWA)); + assert_eq!(tracker[&SWA], 2); + let _ = (a, b); + tc.evict_device_end(SWA); +} + +#[test] +fn evict_walk_revalidates_a_delisted_cursor() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = swa_evict_chain(&mut tc); + tc.evict_device_start(SWA, /* request_cnt = */ 100); + // The cursor (LRU-most a) is delisted before the walk resumes. + tc.device_lru_list_mut(SWA).remove_node(a); + let mut tracker = swa_tracker(); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + let (next, step) = tc.evict_device_next_node(SWA, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + // The walk resets to the list's LRU end (b) and performs one tombstone. + assert_eq!(next, None); + assert!(tc.arena.has_device_value(a, SWA)); + assert!(!tc.arena.has_device_value(b, SWA)); + tc.evict_device_end(SWA); +} + +#[test] +fn evict_walk_second_call_resumes_past_the_returned_leaf() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = swa_evict_chain(&mut tc); + let mut tracker = swa_tracker(); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + tc.evict_device_start(SWA, /* request_cnt = */ 100); + let (first, step) = tc.evict_device_next_node(SWA, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + assert_eq!(first, None); + let (second, step) = tc.evict_device_next_node(SWA, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + assert_eq!(second, None); + let (third, step) = tc.evict_device_next_node(SWA, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + assert_eq!(third, Some(tc.arena.node(c).id)); + // The driver has not delisted c yet: the pre-advanced cursor must not + // hand the same leaf out again. + let (fourth, step) = tc.evict_device_next_node(SWA, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + assert_eq!(fourth, None); + tc.evict_device_end(SWA); + let _ = (a, b); +} + +#[test] +fn evict_walk_start_skips_a_locked_lru_end() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = swa_evict_chain(&mut tc); + tc.arena + .node_mut(a) + .set_lock_ref_(ValueSlotIdx::device(SWA), 1); + let mut tracker = swa_tracker(); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + tc.evict_device_start(SWA, /* request_cnt = */ 100); + let (first, step) = tc.evict_device_next_node(SWA, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + // The locked LRU-end a is invisible; the walk starts at b. + assert_eq!(first, None); + assert!(tc.arena.has_device_value(a, SWA)); + assert!(!tc.arena.has_device_value(b, SWA)); + assert_eq!(tracker[&SWA], 1); + let (second, step) = tc.evict_device_next_node(SWA, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + assert_eq!(second, Some(tc.arena.node(c).id)); + tc.evict_device_end(SWA); +} + +#[test] +fn evict_walk_revalidation_skips_a_locked_lru_end() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = swa_evict_chain(&mut tc); + tc.evict_device_start(SWA, /* request_cnt = */ 100); + // The cursor's node leaves the list and the new LRU end gets locked + // before the walk resumes. + tc.device_lru_list_mut(SWA).remove_node(a); + tc.arena + .node_mut(b) + .set_lock_ref_(ValueSlotIdx::device(SWA), 1); + let mut tracker = swa_tracker(); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + let (next, step) = tc.evict_device_next_node(SWA, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + assert_eq!(next, Some(tc.arena.node(c).id)); + assert!(tc.arena.has_device_value(b, SWA)); + assert_eq!(tracker[&SWA], 0); + tc.evict_device_end(SWA); +} + +#[test] +fn evict_device_end_clears_the_walk_state() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + tc.evict_device_start(SWA, /* request_cnt = */ 1); + tc.evict_device_end(SWA); + // A second walk only starts cleanly when the end cleared the state. + tc.evict_device_start(SWA, /* request_cnt = */ 1); + tc.evict_device_end(SWA); +} + +#[test] +#[should_panic(expected = "valueless node")] +fn evict_walk_asserts_a_valued_cursor_node() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, _b, _c] = swa_evict_chain(&mut tc); + // a stays in the SWA LRU but loses its device value out of band. + let _ = tc.arena.take_device_value(a, SWA); + let tracker = swa_tracker(); + tc.evict_device_start(SWA, /* request_cnt = */ 100); + tc.evict_device_next_node(SWA, &tracker); +} + +#[test] +#[should_panic(expected = "Swa device eviction not started")] +fn evict_walk_requires_a_start() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let tracker = swa_tracker(); + tc.evict_device_next_node(SWA, &tracker); +} + +#[test] +fn try_device_value_and_evictable_size_read_the_swa_slots() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b] = chain(&mut tc); + store_swa_device(&mut tc, a); + assert!( + tc.get_component_device_value(tc.arena.node(a).id, SWA) + .unwrap() + .equal(&Tensor::from_slice(&[0i64])) + ); + assert!( + tc.get_component_device_value(tc.arena.node(b).id, SWA) + .is_none() + ); + assert_eq!(tc.evictable_size_(SWA), 1); +} + +#[test] +fn redistribute_on_node_split_moves_the_swa_uuid_to_the_parent() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0)); + let root = tc.arena.root(); + let node = child_of(&tc, root, &[1]); + tc.arena.node_mut(node).swa_uuid = Some(7); + let (parent, _) = tc.split_node_(node, /* split_len = */ 1); + assert_eq!(node_swa_uuid(&tc, parent), Some(7)); + assert_eq!(node_swa_uuid(&tc, node), None); +} + +#[test] +fn finalize_window_arithmetic_at_page_boundaries() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 2); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let h = tc + .arena + .alloc_child( + a, + /* key = */ vec![3, 4], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let c = tc + .arena + .alloc_child( + h, + /* key = */ vec![5, 6], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + set_swa_host(&mut tc, a); + set_swa_host(&mut tc, h); + set_swa_device(&mut tc, c); + // The window of 4 lands exactly on the c/h page boundary sum: h is + // counted, the page above it is not. + let out = finalize(&tc, &swa_component(4), c, /* prior = */ 0); + assert_eq!(out.swa_host_hit_length, 2); +} + +#[test] +fn new_leaf_after_a_cached_prefix_splits_at_the_leaf_relative_boundary() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0)); + let root = tc.arena.root(); + let a = child_of(&tc, root, &[1]); + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4, 5, 6], + &[20, 21, 22, 23, 24, 25], + /* prev_prefix_len = */ 3, + /* swa_evicted_seqlen = */ 4, + )); + assert_eq!(result.prefix_len, 3); + // The leaf starts at prefix 3, so the boundary at seqlen 4 is + // leaf-relative offset 1: parent [4] tombstone, tail [5, 6] rebuilt. + let p = child_of(&tc, a, &[4]); + let leaf = child_of(&tc, p, &[5]); + assert_eq!(tc.arena.node(p).key, vec![4]); + assert_eq!(tc.arena.node(leaf).key, vec![5, 6]); + assert!( + tc.arena + .device_value(p, FULL) + .equal(&Tensor::from_slice(&[23i64])) + ); + assert!( + tc.arena + .device_value(leaf, FULL) + .equal(&Tensor::from_slice(&[24i64, 25])) + ); + assert!(result.cache_actions.iter().any(|action| matches!( + action, + CacheAction::SwaRebuild { node_id, source_value } + if *node_id == tc.arena.node(leaf).id && source_value.equal(&Tensor::from_slice(&[24i64, 25])) + ))); + assert!(!result.cache_actions.iter().any(|action| matches!( + action, + CacheAction::SwaRebuild { node_id, .. } if *node_id == tc.arena.node(p).id + ))); +} + +#[test] +fn new_leaf_boundary_split_and_window_cap_compose_in_one_commit() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0)); + let root = tc.arena.root(); + let a = child_of(&tc, root, &[1]); + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4, 5, 6, 7], + &[20, 21, 22, 23, 24, 25, 26], + /* prev_prefix_len = */ 2, + /* swa_evicted_seqlen = */ 3, + )); + // Boundary split first ([3] tombstone), then the window cap splits the + // in-window run into [4, 5] + [6, 7], rebuilt older-prefix-first. + let p = child_of(&tc, a, &[3]); + let capped = child_of(&tc, p, &[4]); + let leaf = child_of(&tc, capped, &[6]); + assert_eq!(tc.arena.node(p).key, vec![3]); + assert_eq!(tc.arena.node(capped).key, vec![4, 5]); + assert_eq!(tc.arena.node(leaf).key, vec![6, 7]); + let rebuilds: Vec = result + .cache_actions + .iter() + .filter_map(|action| match action { + CacheAction::SwaRebuild { node_id, .. } => Some(*node_id), + _ => None, + }) + .collect(); + assert_eq!( + rebuilds, + vec![tc.arena.node(capped).id, tc.arena.node(leaf).id] + ); + assert!(result.cache_actions.iter().any(|action| matches!( + action, + CacheAction::SwaRebuild { node_id, source_value } + if *node_id == tc.arena.node(capped).id && source_value.equal(&Tensor::from_slice(&[23i64, 24])) + ))); + assert!(result.cache_actions.iter().any(|action| matches!( + action, + CacheAction::SwaRebuild { node_id, source_value } + if *node_id == tc.arena.node(leaf).id && source_value.equal(&Tensor::from_slice(&[25i64, 26])) + ))); +} + +#[test] +fn finalize_counts_a_dual_tier_node_as_a_device_hit() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b] = chain(&mut tc); + set_swa_device(&mut tc, a); + set_swa_host(&mut tc, a); + set_swa_device(&mut tc, b); + set_swa_host(&mut tc, b); + let out = finalize(&tc, &swa_component(2), b, /* prior = */ 0); + assert_eq!(out.swa_host_hit_length, 0); +} + +#[test] +fn finalize_at_the_root_leaves_the_host_hit_untouched() { + let tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let root = tc.arena.root(); + let out = finalize(&tc, &swa_component(2), root, /* prior = */ 0); + assert_eq!(out.swa_host_hit_length, 0); +} + +#[test] +fn match_prefix_with_an_empty_key_on_a_swa_core_is_a_clean_miss() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0)); + let result = tc.match_prefix(&MatchPrefixParams { + key: &Vec::new(), + namespace: Default::default(), + }); + assert_eq!(result.device_indices.size()[0], 0); + assert_eq!(result.swa_host_hit_length, 0); + let root = tc.arena.root(); + assert_eq!(result.best_match_node_id, tc.arena.node(root).id); +} + +#[test] +fn refresh_window_extends_by_a_full_page_beyond_the_sliding_window() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 4); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2, 3, 4], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let b = tc + .arena + .alloc_child( + a, + /* key = */ vec![5, 6, 7, 8], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let c = tc + .arena + .alloc_child( + b, + /* key = */ vec![9, 10, 11, 12], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + for node in [c, b, a] { + set_swa_device(&mut tc, node); + tc.device_lru_list_mut(SWA).insert_mru(node); + } + // The walk window is sliding_window_size + page_size = 6: c and b + // re-rank deepest first; a stays put. + swa_component(2).refresh_lru(&mut tc, LRURefreshPhase::MatchEnd, c); + assert_eq!(swa_lru_order(&tc), vec![c, b, a]); +} + +#[test] +fn window_cap_skips_a_page_misaligned_leaf() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 2); + let root = tc.arena.root(); + let leaf = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2, 3, 4, 5], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + // split_at = 5 - 2 = 3 is not page-aligned: the cap is skipped. + let capped = swa_component(2).maybe_split_leaf_for_swa_lock_(&mut tc, leaf); + assert_eq!(capped, None); + assert_eq!(tc.arena.node(leaf).key, vec![1, 2, 3, 4, 5]); + assert_eq!(tc.arena.node(leaf).parent(), root); +} + +#[test] +fn new_leaf_with_the_boundary_above_its_start_skips_the_split() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0)); + let root = tc.arena.root(); + let a = child_of(&tc, root, &[1]); + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4, 5], + &[20, 21, 22, 23, 24], + /* prev_prefix_len = */ 3, + /* swa_evicted_seqlen = */ 2, + )); + assert_eq!(result.prefix_len, 3); + // The boundary (2) sits above the leaf start (3): the whole leaf is + // in-window and stays unsplit. + let leaf = child_of(&tc, a, &[4]); + assert_eq!(tc.arena.node(leaf).key, vec![4, 5]); + assert!(result.cache_actions.iter().any(|action| matches!( + action, + CacheAction::SwaRebuild { node_id, source_value } + if *node_id == tc.arena.node(leaf).id && source_value.equal(&Tensor::from_slice(&[23i64, 24])) + ))); +} + +#[test] +fn insert_overlap_straddling_with_a_partial_prev_prefix_recovers_the_tail() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4], + &[10, 11, 12, 13], + 0, + 0, + )); + let root = tc.arena.root(); + let node = child_of(&tc, root, &[1]); + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4], + &[20, 21, 22, 23], + /* prev_prefix_len = */ 1, + /* swa_evicted_seqlen = */ 2, + )); + // prev covers only one atom, so the straddle recovery still runs. + let parent = child_of(&tc, root, &[1]); + assert_eq!(tc.arena.node(parent).key, vec![1, 2]); + assert_eq!(tc.arena.node(node).key, vec![3, 4]); + assert!( + tc.arena + .device_value(parent, FULL) + .equal(&Tensor::from_slice(&[10i64, 11])) + ); + assert!( + tc.arena + .device_value(node, FULL) + .equal(&Tensor::from_slice(&[22i64, 23])) + ); + let [ + CacheAction::FreeDeviceKVFullOnly(old_tail), + CacheAction::SwaRebuild { + node_id, + source_value, + }, + CacheAction::FreeDeviceKV(duplicates), + ] = result.cache_actions.as_slice() + else { + panic!( + "expected FreeDeviceKVFullOnly, SwaRebuild, FreeDeviceKV, got {:?}", + action_kinds(&result.cache_actions) + ); + }; + assert!(old_tail[0].equal(&Tensor::from_slice(&[12i64, 13]))); + assert_eq!(*node_id, tc.arena.node(node).id); + assert!(source_value.equal(&Tensor::from_slice(&[22i64, 23]))); + // The request already owns the first prev_prefix_len token; only the + // stretch between prev and the boundary is duplicate. + assert!(duplicates[0].equal(&Tensor::from_slice(&[21i64]))); +} + +#[test] +fn insert_overlap_straddling_recovers_across_page_two() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 2); + tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4], + &[10, 11, 12, 13], + 0, + 0, + )); + let root = tc.arena.root(); + let node = child_of(&tc, root, &[1, 2]); + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4], + &[20, 21, 22, 23], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 2, + )); + let parent = child_of(&tc, root, &[1, 2]); + assert_eq!(tc.arena.node(parent).key, vec![1, 2]); + assert_eq!(tc.arena.node(node).key, vec![3, 4]); + assert!( + tc.arena + .device_value(node, FULL) + .equal(&Tensor::from_slice(&[22i64, 23])) + ); + assert!(result.cache_actions.iter().any(|action| matches!( + action, + CacheAction::SwaRebuild { node_id, source_value } + if *node_id == tc.arena.node(node).id && source_value.equal(&Tensor::from_slice(&[22i64, 23])) + ))); +} + +#[test] +fn reinsert_straddling_recovers_across_page_two() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 2); + tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4], + &[10, 11, 12, 13], + 0, + 0, + )); + let root = tc.arena.root(); + let node = child_of(&tc, root, &[1, 2]); + evict_full(&mut tc, node, /* remaining_size = */ 0); + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4], + &[20, 21, 22, 23], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 2, + )); + let parent = child_of(&tc, root, &[1, 2]); + assert_eq!(tc.arena.node(parent).key, vec![1, 2]); + assert_eq!(tc.arena.node(node).key, vec![3, 4]); + assert!( + tc.arena + .device_value(parent, FULL) + .equal(&Tensor::from_slice(&[20i64, 21])) + ); + assert!( + tc.arena + .device_value(node, FULL) + .equal(&Tensor::from_slice(&[22i64, 23])) + ); + assert!(result.cache_actions.iter().any(|action| matches!( + action, + CacheAction::SwaRebuild { node_id, source_value } + if *node_id == tc.arena.node(node).id && source_value.equal(&Tensor::from_slice(&[22i64, 23])) + ))); + assert!(!result.cache_actions.iter().any(|action| matches!( + action, + CacheAction::SwaRebuild { node_id, .. } if *node_id == tc.arena.node(parent).id + ))); +} + +#[test] +#[should_panic(expected = "swa_evicted_seqlen must be page-aligned")] +fn reinsert_rejects_a_page_misaligned_boundary() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 2); + tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0)); + let root = tc.arena.root(); + let node = child_of(&tc, root, &[1, 2]); + evict_full(&mut tc, node, /* remaining_size = */ 0); + tc.insert(&insert_params_swa( + &vec![1, 2], + &[20, 21], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 1, + )); +} + +#[test] +#[should_panic(expected = "tombstone Swa lock_ref should be 0 on unevict")] +fn reinsert_rejects_a_locked_tombstone() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0)); + let root = tc.arena.root(); + let node = child_of(&tc, root, &[1]); + evict_full(&mut tc, node, /* remaining_size = */ 0); + tc.arena + .node_mut(node) + .set_lock_ref_(ValueSlotIdx::device(SWA), 1); + tc.insert(&insert_params_swa(&vec![1, 2], &[20, 21], 0, 0)); +} + +fn set_full_host(tc: &mut UnifiedTreeCore>, node: NodeIdx_) { + let len = tc.arena.node(node).key.atom_len(); + tc.arena + .set_host_value(node, FULL, Tensor::from_slice(&vec![0i64; len])); +} + +fn host_drive_state() -> ( + HashMap, + HashMap>, + HashMap>, +) { + ( + HashMap::from([(FULL, 0), (SWA, 0)]), + HashMap::new(), + HashMap::new(), + ) +} + +#[test] +fn host_drive_tombstones_internal_nodes_and_evicts_host_leaves() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [p, c] = chain(&mut tc); + set_full_host(&mut tc, p); + set_swa_host(&mut tc, p); + set_full_host(&mut tc, c); + set_swa_host(&mut tc, c); + // p enters the LRU first so the walk reaches it before c. + tc.host_lru_list_mut(SWA).insert_mru(p); + tc.host_lru_list_mut(SWA).insert_mru(c); + tc.evictable_host_leaves.add(c); + let (mut tr, mut df, mut hf) = host_drive_state(); + accumulate_step( + tc.drive_host_eviction(SWA, /* num_tokens = */ 100), + &mut tr, + &mut df, + &mut hf, + ); + // p: private tombstone (SWA host only); c: atomic H-leaf teardown. + assert_eq!(tr[&SWA], 2); + assert_eq!(tr[&FULL], 1); + assert_eq!(hf[&SWA].len(), 2); + assert_eq!(hf[&FULL].len(), 1); + assert!(tc.arena.has_host_value(p, FULL)); + assert!(!tc.arena.has_host_value(p, SWA)); + assert_eq!(tc.arena.len(), 2); + assert!(tc.evictable_host_leaves.contains(p)); + assert_eq!(tc.host_lru_list(SWA).len(), 0); + tc.sanity_check(&[], &[]); +} + +#[test] +fn host_drive_skips_host_locked_nodes() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let root = tc.arena.root(); + let locked = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let victim = tc + .arena + .alloc_child( + root, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + for node in [locked, victim] { + set_full_host(&mut tc, node); + set_swa_host(&mut tc, node); + tc.host_lru_list_mut(SWA).insert_mru(node); + } + tc.arena + .node_mut(locked) + .set_lock_ref_(ValueSlotIdx::host(SWA), 1); + tc.evictable_host_leaves.add(victim); + let (mut tr, mut df, mut hf) = host_drive_state(); + accumulate_step( + tc.drive_host_eviction(SWA, /* num_tokens = */ 100), + &mut tr, + &mut df, + &mut hf, + ); + // The locked LRU-most node is passed over; the unlocked one is evicted. + assert!(tc.arena.has_host_value(locked, SWA)); + assert!(tc.host_lru_list(SWA).in_list(Some(locked))); + assert_eq!(tr[&SWA], 1); + assert_eq!(tr[&FULL], 1); + assert_eq!(tc.arena.len(), 2); + tc.sanity_check(&[], &[]); +} + +#[test] +fn host_drive_stops_at_the_token_budget_consuming_lru_first() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let root = tc.arena.root(); + let h1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let h2 = tc + .arena + .alloc_child( + root, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + for node in [h1, h2] { + set_full_host(&mut tc, node); + set_swa_host(&mut tc, node); + tc.host_lru_list_mut(SWA).insert_mru(node); + tc.evictable_host_leaves.add(node); + } + let (mut tr, mut df, mut hf) = host_drive_state(); + accumulate_step( + tc.drive_host_eviction(SWA, /* num_tokens = */ 1), + &mut tr, + &mut df, + &mut hf, + ); + // The LRU-most h1 fills the budget; h2 survives untouched. + assert_eq!(tr[&SWA], 1); + assert!(tc.evictable_host_leaves.contains(h2)); + assert!(tc.host_lru_list(SWA).in_list(Some(h2))); + assert_eq!(tc.arena.len(), 2); + tc.sanity_check(&[], &[]); +} + +#[test] +fn host_drive_ends_when_the_next_candidate_left_the_lru() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let root = tc.arena.root(); + let p = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let c = tc + .arena + .alloc_child( + p, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let s = tc + .arena + .alloc_child( + root, + /* key = */ vec![9], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + // p holds only a transient SWA host chunk: c's teardown tombstone-walks + // p away, so the captured next candidate leaves the list mid-drive. + set_swa_host(&mut tc, p); + set_full_host(&mut tc, c); + set_swa_host(&mut tc, c); + set_full_host(&mut tc, s); + set_swa_host(&mut tc, s); + tc.host_lru_list_mut(SWA).insert_mru(c); + tc.host_lru_list_mut(SWA).insert_mru(p); + tc.host_lru_list_mut(SWA).insert_mru(s); + tc.evictable_host_leaves.add(c); + tc.evictable_host_leaves.add(s); + let (mut tr, mut df, mut hf) = host_drive_state(); + accumulate_step( + tc.drive_host_eviction(SWA, /* num_tokens = */ 100), + &mut tr, + &mut df, + &mut hf, + ); + // The walk stops at the vanished p; the more-recent s is never reached. + assert!(tc.arena.has_host_value(s, SWA)); + assert!(tc.host_lru_list(SWA).in_list(Some(s))); + assert_eq!(tr[&SWA], 2); + assert_eq!(tr[&FULL], 1); + assert_eq!(tc.arena.len(), 2); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "has no host value")] +fn host_drive_panics_on_an_lru_member_without_a_swa_host_value() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let root = tc.arena.root(); + let n = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + set_full_host(&mut tc, n); + tc.host_lru_list_mut(SWA).insert_mru(n); + tc.drive_host_eviction(SWA, /* num_tokens = */ 100); +} + +#[test] +fn backup_storage_transfers_carry_trailing_page_keys() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + tc.arena + .set_host_value(a, SWA, Tensor::from_slice(&[21i64])); + tc.arena.node_mut(a).hash_value = Some(vec!["h0".to_string()]); + let transfers = swa_component(4) + .build_hicache_transfers( + &tc, + a, + CacheTransferPhase::BackupStorage, + /* mamba_pool_idx = */ None, + /* host_indices = */ None, + /* token_ids = */ None, + /* prefetch_tokens = */ 0, + /* last_hash = */ None, + ) + .unwrap() + .unwrap(); + assert_eq!(transfers.len(), 1); + assert_eq!(transfers[0].name, PoolName::Swa); + assert!( + transfers[0] + .host_indices + .as_ref() + .unwrap() + .equal(&Tensor::from_slice(&[21i64])) + ); + assert_eq!(transfers[0].keys, Some(vec!["h0".to_string()])); + assert_eq!(transfers[0].hit_policy, PoolHitPolicy::TrailingPages); +} + +#[test] +fn backup_storage_is_none_without_host_value_or_hashes() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + let build = |tc: &UnifiedTreeCore>| { + swa_component(4).build_hicache_transfers( + tc, + a, + CacheTransferPhase::BackupStorage, + /* mamba_pool_idx = */ None, + /* host_indices = */ None, + /* token_ids = */ None, + /* prefetch_tokens = */ 0, + /* last_hash = */ None, + ) + }; + assert!(build(&tc).unwrap().is_none()); + tc.arena + .set_host_value(a, SWA, Tensor::from_slice(&[21i64])); + tc.arena.node_mut(a).hash_value = None; + assert!(build(&tc).unwrap().is_none()); +} + +#[test] +fn build_transfers_are_gated_off_until_the_swa_host_pool_is_wired() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + tc.set_hicache_enabled(); + let [a] = chain::<1>(&mut tc); + set_swa_device(&mut tc, a); + for phase in [CacheTransferPhase::BackupHost, CacheTransferPhase::LoadBack] { + let transfers = swa_component(4) + .build_hicache_transfers( + &tc, a, phase, /* mamba_pool_idx = */ None, /* host_indices = */ None, + /* token_ids = */ None, /* prefetch_tokens = */ 0, + /* last_hash = */ None, + ) + .unwrap(); + assert!(transfers.is_none()); + } + // Wiring the pool opens the gate. + tc.set_has_swa_host_pool(); + let transfers = swa_component(4) + .build_hicache_transfers( + &tc, + a, + CacheTransferPhase::BackupHost, + /* mamba_pool_idx = */ None, + /* host_indices = */ None, + /* token_ids = */ None, + /* prefetch_tokens = */ 0, + /* last_hash = */ None, + ) + .unwrap(); + assert!(transfers.is_some()); +} + +#[test] +fn backup_host_build_wraps_the_device_value_as_int64() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + tc.arena + .set_device_value(a, SWA, Tensor::from_slice(&[5i32])); + let transfers = swa_component(4) + .build_hicache_transfers( + &tc, + a, + CacheTransferPhase::BackupHost, + /* mamba_pool_idx = */ None, + /* host_indices = */ None, + /* token_ids = */ None, + /* prefetch_tokens = */ 0, + /* last_hash = */ None, + ) + .unwrap() + .unwrap(); + assert_eq!(transfers.len(), 1); + let xfer = &transfers[0]; + assert_eq!(xfer.name, PoolName::Swa); + let device_indices = xfer.device_indices.as_ref().unwrap(); + assert_eq!(device_indices.kind(), Kind::Int64); + assert!(device_indices.equal(&Tensor::from_slice(&[5i64]))); + assert!(xfer.host_indices.is_none()); + assert!(xfer.nodes_to_load.is_none()); +} + +#[test] +fn backup_host_build_returns_none_for_a_tombstone() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + let transfers = swa_component(4) + .build_hicache_transfers( + &tc, + a, + CacheTransferPhase::BackupHost, + /* mamba_pool_idx = */ None, + /* host_indices = */ None, + /* token_ids = */ None, + /* prefetch_tokens = */ 0, + /* last_hash = */ None, + ) + .unwrap(); + assert!(transfers.is_none()); +} + +#[test] +fn backup_spec_reads_the_swa_value_recovered_by_an_earlier_action() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + tc.arena + .set_device_value(a, FULL, Tensor::from_slice(&[9i64])); + let a_id = tc.arena.node(a).id; + // The tombstone carries no SWA transfer into the backup spec. + let (_, xfers) = tc.build_backup_spec(a_id); + assert!(xfers.is_empty()); + // The cache resolves the recover/rebuild action, then rebuilds the spec: + // the deferred read now captures the freshly stored SWA value. + tc.set_component_device_value(a_id, SWA, Tensor::from_slice(&[50i64])); + let (_, xfers) = tc.build_backup_spec(a_id); + let swa_xfer = &xfers[&SWA][0]; + assert!( + swa_xfer + .device_indices + .as_ref() + .unwrap() + .equal(&Tensor::from_slice(&[50i64])) + ); +} + +#[test] +fn load_back_build_collects_host_only_nodes_within_the_window() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a, b, c] = chain::<3>(&mut tc); + set_swa_device(&mut tc, a); + tc.arena + .set_host_value(b, SWA, Tensor::from_slice(&[21i64])); + tc.arena + .set_host_value(c, SWA, Tensor::from_slice(&[22i64])); + let transfers = swa_component(4) + .build_hicache_transfers( + &tc, + c, + CacheTransferPhase::LoadBack, + /* mamba_pool_idx = */ None, + /* host_indices = */ None, + /* token_ids = */ None, + /* prefetch_tokens = */ 0, + /* last_hash = */ None, + ) + .unwrap() + .unwrap(); + assert_eq!(transfers.len(), 1); + let xfer = &transfers[0]; + assert_eq!(xfer.name, PoolName::Swa); + // Ancestor-first; a's device value is skipped, not collected. + assert!( + xfer.host_indices + .as_ref() + .unwrap() + .equal(&Tensor::from_slice(&[21i64, 22])) + ); + assert!(xfer.device_indices.is_none()); + assert_eq!( + xfer.nodes_to_load, + Some(vec![tc.arena.node(b).id, tc.arena.node(c).id]) + ); +} + +#[test] +fn load_back_build_stops_at_the_window_boundary() { + let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); + let [a, b, c] = chain::<3>(&mut tc); + tc.arena + .set_host_value(a, SWA, Tensor::from_slice(&[20i64])); + tc.arena + .set_host_value(b, SWA, Tensor::from_slice(&[21i64])); + tc.arena + .set_host_value(c, SWA, Tensor::from_slice(&[22i64])); + let transfers = swa_component(2) + .build_hicache_transfers( + &tc, + c, + CacheTransferPhase::LoadBack, + /* mamba_pool_idx = */ None, + /* host_indices = */ None, + /* token_ids = */ None, + /* prefetch_tokens = */ 0, + /* last_hash = */ None, + ) + .unwrap() + .unwrap(); + // The two-token window covers c and b; a stays out of the transfer. + let xfer = &transfers[0]; + assert!( + xfer.host_indices + .as_ref() + .unwrap() + .equal(&Tensor::from_slice(&[21i64, 22])) + ); + assert_eq!( + xfer.nodes_to_load, + Some(vec![tc.arena.node(b).id, tc.arena.node(c).id]) + ); +} + +#[test] +fn load_back_build_returns_none_when_the_window_is_on_device() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a, b] = chain::<2>(&mut tc); + set_swa_device(&mut tc, a); + set_swa_device(&mut tc, b); + let transfers = swa_component(4) + .build_hicache_transfers( + &tc, + b, + CacheTransferPhase::LoadBack, + /* mamba_pool_idx = */ None, + /* host_indices = */ None, + /* token_ids = */ None, + /* prefetch_tokens = */ 0, + /* last_hash = */ None, + ) + .unwrap(); + assert!(transfers.is_none()); +} + +#[test] +fn load_back_build_rejects_a_bare_window_node() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + assert!(matches!( + swa_component(4).build_hicache_transfers( + &tc, + a, + CacheTransferPhase::LoadBack, + /* mamba_pool_idx = */ None, + /* host_indices = */ None, + /* token_ids = */ None, + /* prefetch_tokens = */ 0, + /* last_hash = */ None, + ), + Err(TreeCoreRuntimeError::SwaLoadBackMissingValue { node_id }) + if node_id == tc.arena.node(a).id + )); +} + +#[test] +fn fallible_load_back_boundaries_reject_a_bare_window_node() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + tc.arena + .set_device_value(a, FULL, Tensor::from_slice(&[10i64])); + let node_id = tc.arena.node(a).id; + + assert!(matches!( + tc.try_build_hicache_transfers( + SWA, + node_id, + CacheTransferPhase::LoadBack, + /* host_indices = */ None, + /* token_ids = */ None, + /* prefetch_tokens = */ 0, + /* last_hash = */ None, + ), + Err(TreeCoreRuntimeError::SwaLoadBackMissingValue { node_id: missing }) + if missing == node_id + )); + assert!(matches!( + tc.try_build_load_back_spec(node_id, /* req = */ None), + Err(TreeCoreRuntimeError::SwaLoadBackMissingValue { node_id: missing }) + if missing == node_id + )); +} + +#[test] +fn load_back_commit_attaches_chunks_and_emits_the_rebuild_action() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a, b] = chain::<2>(&mut tc); + tc.arena + .set_device_value(a, FULL, Tensor::from_slice(&[40i64])); + tc.arena + .set_device_value(b, FULL, Tensor::from_slice(&[41i64])); + tc.arena + .set_host_value(a, SWA, Tensor::from_slice(&[21i64])); + tc.arena + .set_host_value(b, SWA, Tensor::from_slice(&[22i64])); + let mut cache_actions = Vec::new(); + let transfer = PoolTransfer { + name: PoolName::Swa, + host_indices: Some(Tensor::from_slice(&[21i64, 22])), + device_indices: Some(Tensor::from_slice(&[50i64, 51])), + nodes_to_load: Some(vec![tc.arena.node(a).id, tc.arena.node(b).id]), + ..Default::default() + }; + swa_component(4).commit_hicache_transfer( + &mut tc, + b, + CacheTransferPhase::LoadBack, + vec![transfer], + &mut cache_actions, + /* insert_result = */ None, + /* pool_storage_result = */ None, + ); + assert!( + tc.arena + .device_value(a, SWA) + .equal(&Tensor::from_slice(&[50i64])) + ); + assert!( + tc.arena + .device_value(b, SWA) + .equal(&Tensor::from_slice(&[51i64])) + ); + // The attach path restamps the SWA device LRU and evictable size. + assert!(tc.device_lru_list(SWA).in_list(Some(a))); + assert!(tc.device_lru_list(SWA).in_list(Some(b))); + assert_eq!(tc.swa_evictable_size(), 2); + assert_eq!(cache_actions.len(), 1); + let CacheAction::RebuildFullToSwaMapping { + full_indices, + swa_indices, + } = &cache_actions[0] + else { + panic!("expected a RebuildFullToSwaMapping action"); + }; + assert_eq!(full_indices.len(), 2); + assert!(full_indices[0].equal(&Tensor::from_slice(&[40i64]))); + assert!(full_indices[1].equal(&Tensor::from_slice(&[41i64]))); + assert_eq!(swa_indices.len(), 2); + assert!(swa_indices[0].equal(&Tensor::from_slice(&[50i64]))); + assert!(swa_indices[1].equal(&Tensor::from_slice(&[51i64]))); +} + +#[test] +#[should_panic(expected = "SWA LOAD_BACK commit requires device indices")] +fn load_back_commit_panics_without_device_indices() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + let transfer = PoolTransfer { + name: PoolName::Swa, + host_indices: Some(Tensor::from_slice(&[21i64])), + nodes_to_load: Some(vec![tc.arena.node(a).id]), + ..Default::default() + }; + swa_component(4).commit_hicache_transfer( + &mut tc, + a, + CacheTransferPhase::LoadBack, + vec![transfer], + &mut Vec::new(), + /* insert_result = */ None, + /* pool_storage_result = */ None, + ); +} + +#[test] +#[should_panic(expected = "left == right")] +fn load_back_commit_asserts_the_loaded_length_matches_the_host_indices() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + tc.arena + .set_device_value(a, FULL, Tensor::from_slice(&[40i64])); + tc.arena + .set_host_value(a, SWA, Tensor::from_slice(&[21i64])); + // Two host indices but only one loaded token: the commit must fail loudly. + let transfer = PoolTransfer { + name: PoolName::Swa, + host_indices: Some(Tensor::from_slice(&[21i64, 22])), + device_indices: Some(Tensor::from_slice(&[50i64, 51])), + nodes_to_load: Some(vec![tc.arena.node(a).id]), + ..Default::default() + }; + swa_component(4).commit_hicache_transfer( + &mut tc, + a, + CacheTransferPhase::LoadBack, + vec![transfer], + &mut Vec::new(), + /* insert_result = */ None, + /* pool_storage_result = */ None, + ); +} + +#[test] +fn backup_host_commit_sets_the_host_value_once() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + for host in [30i64, 31] { + let transfer = PoolTransfer { + name: PoolName::Swa, + host_indices: Some(Tensor::from_slice(&[host])), + ..Default::default() + }; + swa_component(4).commit_hicache_transfer( + &mut tc, + a, + CacheTransferPhase::BackupHost, + vec![transfer], + &mut Vec::new(), + /* insert_result = */ None, + /* pool_storage_result = */ None, + ); + } + // The second commit is a no-op: the first host value sticks. + assert!( + tc.arena + .host_value(a, SWA) + .equal(&Tensor::from_slice(&[30i64])) + ); +} + +#[test] +fn backup_host_commit_ignores_transfers_without_host_indices() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + for transfers in [vec![], vec![PoolTransfer::default()]] { + swa_component(4).commit_hicache_transfer( + &mut tc, + a, + CacheTransferPhase::BackupHost, + transfers, + &mut Vec::new(), + /* insert_result = */ None, + /* pool_storage_result = */ None, + ); + } + assert!(!tc.arena.has_host_value(a, SWA)); +} + +#[test] +fn backup_storage_commit_is_a_noop() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + let mut cache_actions = Vec::new(); + swa_component(4).commit_hicache_transfer( + &mut tc, + a, + CacheTransferPhase::BackupStorage, + vec![], + &mut cache_actions, + /* insert_result = */ None, + /* pool_storage_result = */ None, + ); + assert!(cache_actions.is_empty()); +} + +#[test] +fn commit_hicache_transfers_routes_to_the_component() { + // Underloaded prefetch through the core dispatcher releases the buffer. + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [_a] = chain::<1>(&mut tc); + let mut cache_actions = Vec::new(); + let root = tc.arena.root(); + tc.commit_hicache_transfers( + tc.arena.node(root).id, + CacheTransferPhase::Prefetch, + HashMap::from([( + SWA, + vec![PoolTransfer { + name: PoolName::Swa, + host_indices: Some(Tensor::from_slice(&[30i64])), + ..Default::default() + }], + )]), + &mut cache_actions, + /* insert_result = */ None, + /* pool_storage_result = */ None, + ); + assert_eq!(cache_actions.len(), 1); +} + +#[test] +fn prefetch_build_wraps_the_host_buffer_with_placeholder_keys() { + let tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let transfers = swa_component(4) + .build_hicache_transfers( + &tc, + tc.arena.root(), + CacheTransferPhase::Prefetch, + /* mamba_pool_idx = */ None, + /* host_indices = */ Some(Tensor::from_slice(&[30i64, 31])), + /* token_ids = */ None, + /* prefetch_tokens = */ 0, + /* last_hash = */ None, + ) + .unwrap() + .unwrap(); + assert_eq!(transfers.len(), 1); + assert_eq!( + transfers[0].keys, + Some(vec![ + "__placeholder__".to_string(), + "__placeholder__".to_string() + ]) + ); + assert_eq!(transfers[0].hit_policy, PoolHitPolicy::TrailingPages); + assert!( + transfers[0] + .host_indices + .as_ref() + .unwrap() + .equal(&Tensor::from_slice(&[30i64, 31])) + ); +} + +#[test] +fn prefetch_commit_drops_the_whole_window_when_underloaded() { + // loaded_pages (1) < window_require_pages (2): all-or-nothing releases + // the full buffer and attaches nothing. + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + let mut cache_actions = Vec::new(); + let mut insert_result = InsertResult { + total_len: 2, + inserted_host_node: Some(tc.arena.node(a).id), + ..InsertResult::default() + }; + let storage_result = PoolTransferResult { + kv_hit_pages: 2, + extra_pool_hit_pages: HashMap::from([(PoolName::Swa, 1)]), + }; + let root = tc.arena.root(); + swa_component(4).commit_hicache_transfer( + &mut tc, + root, + CacheTransferPhase::Prefetch, + vec![PoolTransfer { + name: PoolName::Swa, + host_indices: Some(Tensor::from_slice(&[30i64, 31])), + ..Default::default() + }], + &mut cache_actions, + Some(&mut insert_result), + Some(&storage_result), + ); + assert!(!tc.arena.node(a).has_host_value(SWA)); + assert_eq!(cache_actions.len(), 1); + let CacheAction::FreeComponentHostSlot { host_indices, .. } = &cache_actions[0] else { + panic!("expected a host free"); + }; + assert!(host_indices[0].equal(&Tensor::from_slice(&[30i64, 31]))); +} + +#[test] +fn prefetch_commit_without_a_target_releases_the_whole_buffer() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + let storage_result = PoolTransferResult { + kv_hit_pages: 2, + extra_pool_hit_pages: HashMap::from([(PoolName::Swa, 2)]), + }; + let root = tc.arena.root(); + // No insert result at all: the buffer has no anchor and fully releases. + let mut cache_actions = Vec::new(); + swa_component(4).commit_hicache_transfer( + &mut tc, + root, + CacheTransferPhase::Prefetch, + vec![PoolTransfer { + name: PoolName::Swa, + host_indices: Some(Tensor::from_slice(&[30i64, 31])), + ..Default::default() + }], + &mut cache_actions, + /* insert_result = */ None, + Some(&storage_result), + ); + assert_eq!(cache_actions.len(), 1); + let CacheAction::FreeComponentHostSlot { host_indices, .. } = &cache_actions[0] else { + panic!("expected a host free"); + }; + assert!(host_indices[0].equal(&Tensor::from_slice(&[30i64, 31]))); + + // An insert result without an inserted host node releases the same way. + let mut insert_result = InsertResult { + total_len: 2, + inserted_host_node: None, + ..InsertResult::default() + }; + let mut cache_actions = Vec::new(); + swa_component(4).commit_hicache_transfer( + &mut tc, + root, + CacheTransferPhase::Prefetch, + vec![PoolTransfer { + name: PoolName::Swa, + host_indices: Some(Tensor::from_slice(&[30i64, 31])), + ..Default::default() + }], + &mut cache_actions, + Some(&mut insert_result), + Some(&storage_result), + ); + assert_eq!(cache_actions.len(), 1); + let CacheAction::FreeComponentHostSlot { host_indices, .. } = &cache_actions[0] else { + panic!("expected a host free"); + }; + assert!(host_indices[0].equal(&Tensor::from_slice(&[30i64, 31]))); + assert!(!tc.arena.node(a).has_host_value(SWA)); +} + +#[test] +fn prefetch_commit_releases_the_out_of_path_prefix() { + // root -> a -> b -> c, one token each; anchor b, target c: the loaded + // window spans two tokens but the leaf->anchor path covers only c's one. + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a, b, c] = chain::<3>(&mut tc); + let mut cache_actions = Vec::new(); + let mut insert_result = InsertResult { + total_len: 3, + inserted_host_node: Some(tc.arena.node(c).id), + ..InsertResult::default() + }; + let storage_result = PoolTransferResult { + kv_hit_pages: 3, + extra_pool_hit_pages: HashMap::from([(PoolName::Swa, 2)]), + }; + swa_component(4).commit_hicache_transfer( + &mut tc, + b, + CacheTransferPhase::Prefetch, + vec![PoolTransfer { + name: PoolName::Swa, + host_indices: Some(Tensor::from_slice(&[30i64, 31])), + ..Default::default() + }], + &mut cache_actions, + Some(&mut insert_result), + Some(&storage_result), + ); + // c (on path) fills with the buffer tail; the out-of-path prefix releases. + assert!( + tc.arena + .node(c) + .host_value(SWA) + .equal(&Tensor::from_slice(&[31i64])) + ); + assert!(!tc.arena.node(b).has_host_value(SWA)); + assert!(!tc.arena.node(a).has_host_value(SWA)); + assert!(tc.host_lru_list(SWA).in_list(Some(c))); + assert_eq!(cache_actions.len(), 1); + let CacheAction::FreeComponentHostSlot { host_indices, .. } = &cache_actions[0] else { + panic!("expected a host free"); + }; + assert!(host_indices[0].equal(&Tensor::from_slice(&[30i64]))); +} + +#[test] +fn prefetch_commit_fills_tombstoned_nodes_and_releases_covered_ones() { + // Chain root -> a -> b (one token each); b holds SWA host already, a is + // a tombstone: b's slice releases, a's fills. + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a, b] = chain::<2>(&mut tc); + tc.arena + .set_host_value(b, SWA, Tensor::from_slice(&[40i64])); + let mut cache_actions = Vec::new(); + let mut insert_result = InsertResult { + total_len: 2, + inserted_host_node: Some(tc.arena.node(b).id), + ..InsertResult::default() + }; + let storage_result = PoolTransferResult { + kv_hit_pages: 2, + extra_pool_hit_pages: HashMap::from([(PoolName::Swa, 2)]), + }; + let root = tc.arena.root(); + swa_component(4).commit_hicache_transfer( + &mut tc, + root, + CacheTransferPhase::Prefetch, + vec![PoolTransfer { + name: PoolName::Swa, + host_indices: Some(Tensor::from_slice(&[30i64, 31])), + ..Default::default() + }], + &mut cache_actions, + Some(&mut insert_result), + Some(&storage_result), + ); + // b (already hosted) released its slice [31]; a filled with [30]. + assert!( + tc.arena + .node(a) + .host_value(SWA) + .equal(&Tensor::from_slice(&[30i64])) + ); + assert!( + tc.arena + .node(b) + .host_value(SWA) + .equal(&Tensor::from_slice(&[40i64])) + ); + assert_eq!(cache_actions.len(), 1); + let CacheAction::FreeComponentHostSlot { host_indices, .. } = &cache_actions[0] else { + panic!("expected a host free"); + }; + assert!(host_indices[0].equal(&Tensor::from_slice(&[31i64]))); +} + +#[test] +fn prefetch_commit_splits_a_partially_covered_tombstone() { + // One two-token tombstone node; the buffer covers only its tail token, + // so the node splits and the tail attaches. + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(a, FULL, Tensor::from_slice(&[10i64, 11])); + tc.update_evictable_leaf_sets_(a); + let mut cache_actions = Vec::new(); + let mut insert_result = InsertResult { + total_len: 2, + inserted_host_node: Some(tc.arena.node(a).id), + ..InsertResult::default() + }; + let storage_result = PoolTransferResult { + kv_hit_pages: 1, + extra_pool_hit_pages: HashMap::from([(PoolName::Swa, 1)]), + }; + swa_component(4).commit_hicache_transfer( + &mut tc, + root, + CacheTransferPhase::Prefetch, + vec![PoolTransfer { + name: PoolName::Swa, + host_indices: Some(Tensor::from_slice(&[31i64])), + ..Default::default() + }], + &mut cache_actions, + Some(&mut insert_result), + Some(&storage_result), + ); + // The node split at token 1; its tail (still node a) got the slice. + let node = tc.arena.node(a); + assert_eq!(node.key.atom_len(), 1); + assert!(node.host_value(SWA).equal(&Tensor::from_slice(&[31i64]))); + assert!(cache_actions.is_empty()); +} + +#[test] +fn release_swa_host_queues_a_free_action_for_non_empty_indices() { + let mut cache_actions = Vec::new(); + swa_component(4).release_swa_host_(Tensor::from_slice(&[7i64]), &mut cache_actions); + assert_eq!(cache_actions.len(), 1); + let CacheAction::FreeComponentHostSlot { + component_type, + host_indices, + } = &cache_actions[0] + else { + panic!("expected a FreeComponentHostSlot action"); + }; + assert_eq!(*component_type, SWA); + assert_eq!(host_indices.len(), 1); + assert!(host_indices[0].equal(&Tensor::from_slice(&[7i64]))); + // Empty indices queue nothing. + let empty: [i64; 0] = []; + swa_component(4).release_swa_host_(Tensor::from_slice(&empty), &mut cache_actions); + assert_eq!(cache_actions.len(), 1); +} + +#[test] +fn attach_swa_host_value_inserts_tombstones_into_the_host_lru() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a, b] = chain::<2>(&mut tc); + set_swa_device(&mut tc, a); + let swa = swa_component(4); + // A tombstone lands in the host LRU; a device-backed node does not. + swa.attach_swa_host_value_(&mut tc, b, Tensor::from_slice(&[9i64])); + assert!( + tc.arena + .host_value(b, SWA) + .equal(&Tensor::from_slice(&[9i64])) + ); + assert!(tc.host_lru_list(SWA).in_list(Some(b))); + swa.attach_swa_host_value_(&mut tc, a, Tensor::from_slice(&[8i64])); + assert!(tc.arena.has_host_value(a, SWA)); + assert!(!tc.host_lru_list(SWA).in_list(Some(a))); +} + +#[test] +fn attach_swa_host_value_skips_reinsertion_when_already_listed() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a] = chain::<1>(&mut tc); + tc.host_lru_list_mut(SWA).insert_mru(a); + swa_component(4).attach_swa_host_value_(&mut tc, a, Tensor::from_slice(&[9i64])); + assert!(tc.host_lru_list(SWA).in_list(Some(a))); +} + +#[test] +fn build_load_back_spec_includes_the_swa_transfers() { + let mut tc: UnifiedTreeCore> = UnifiedTreeCore::new( + CacheInitParams { + enable_hicache: true, + has_swa_host_pool: true, + ..swa_params_with_window(4) + }, + vec![FULL, SWA], + ); + let [n] = chain::<1>(&mut tc); + set_full_host(&mut tc, n); + tc.arena + .set_host_value(n, SWA, Tensor::from_slice(&[30i64])); + let (kv_xfer, mut comp_xfers) = + tc.build_load_back_spec(tc.arena.node(n).id, /* req = */ None); + assert_eq!(kv_xfer.nodes_to_load, Some(vec![tc.arena.node(n).id])); + let swa_xfers = comp_xfers.get_mut(&SWA).unwrap(); + assert_eq!(swa_xfers.len(), 1); + assert!( + swa_xfers[0] + .host_indices + .as_ref() + .unwrap() + .equal(&Tensor::from_slice(&[30i64])) + ); + assert_eq!(swa_xfers[0].nodes_to_load, Some(vec![tc.arena.node(n).id])); + // The orchestrator fills each transfer's device side from the pool load. + swa_xfers[0].device_indices = Some(Tensor::from_slice(&[60i64])); + let actions = tc.commit_load_back( + tc.arena.node(n).id, + Tensor::from_slice(&[50i64]), + kv_xfer, + comp_xfers, + ); + assert!( + tc.arena + .device_value(n, FULL) + .equal(&Tensor::from_slice(&[50i64])) + ); + assert!( + tc.arena + .device_value(n, SWA) + .equal(&Tensor::from_slice(&[60i64])) + ); + assert_eq!(actions.len(), 1); + assert!(matches!( + actions[0], + CacheAction::RebuildFullToSwaMapping { .. } + )); +} + +#[test] +fn auxiliary_load_does_not_reuse_a_full_pending_pin() { + let mut tc: UnifiedTreeCore> = UnifiedTreeCore::new( + CacheInitParams { + is_write_back: true, + enable_hicache: true, + has_swa_host_pool: true, + ..swa_params_with_window(4) + }, + vec![FULL, SWA], + ); + let [shared, anchor_b] = chain::<2>(&mut tc); + set_full_host(&mut tc, shared); + set_full_host(&mut tc, anchor_b); + set_swa_host(&mut tc, shared); + + let shared_id = tc.arena.node(shared).id; + let anchor_b_id = tc.arena.node(anchor_b).id; + tc.commit_load_back( + shared_id, + Tensor::from_slice(&[10i64]), + PoolTransfer { + name: PoolName::Kv, + host_indices: Some(Tensor::from_slice(&[1i64])), + nodes_to_load: Some(vec![shared_id]), + ..Default::default() + }, + HashMap::new(), + ); + assert_eq!(tc.arena.node(shared).load_back_pending_id, Some(shared_id)); + + let swa_xfer = PoolTransfer { + name: PoolName::Swa, + host_indices: Some(Tensor::from_slice(&[2i64])), + device_indices: Some(Tensor::from_slice(&[20i64])), + nodes_to_load: Some(vec![shared_id]), + ..Default::default() + }; + tc.commit_load_back( + anchor_b_id, + Tensor::from_slice(&[30i64]), + PoolTransfer { + name: PoolName::Kv, + host_indices: Some(Tensor::from_slice(&[3i64])), + nodes_to_load: Some(vec![anchor_b_id]), + ..Default::default() + }, + HashMap::from([(SWA, vec![swa_xfer])]), + ); + + assert_eq!(tc.arena.node(shared).load_back_pending_id, Some(shared_id)); + assert_eq!( + tc.arena.node(anchor_b).load_back_pending_id, + Some(anchor_b_id) + ); +} + +#[test] +fn swa_device_eviction_skips_a_load_back_pinned_node() { + let mut tc: UnifiedTreeCore> = UnifiedTreeCore::new( + CacheInitParams { + is_write_back: true, + enable_hicache: true, + has_swa_host_pool: true, + ..swa_params_with_window(4) + }, + vec![FULL, SWA], + ); + let [n] = chain::<1>(&mut tc); + set_full_host(&mut tc, n); + set_swa_host(&mut tc, n); + let (kv_xfer, mut comp_xfers) = + tc.build_load_back_spec(tc.arena.node(n).id, /* req = */ None); + comp_xfers.get_mut(&SWA).unwrap()[0].device_indices = Some(Tensor::from_slice(&[60i64])); + tc.commit_load_back( + tc.arena.node(n).id, + Tensor::from_slice(&[50i64]), + kv_xfer, + comp_xfers, + ); + // The pin alone keeps the in-flight SWA slice out of every eviction branch. + tc.evict_device_start(SWA, 4); + let (next, _) = tc.evict_device_next_node(SWA, &HashMap::new()); + assert_eq!(next, None); + tc.evict_device_end(SWA); + assert!(tc.arena.has_device_value(n, SWA)); + tc.finish_load_back(tc.arena.node(n).id); + tc.evict_device_start(SWA, 4); + let (next, _) = tc.evict_device_next_node(SWA, &HashMap::new()); + assert_eq!(next, Some(tc.arena.node(n).id)); + tc.evict_device_end(SWA); +} + +#[test] +fn swa_host_eviction_skips_a_load_back_pinned_node() { + let mut tc: UnifiedTreeCore> = UnifiedTreeCore::new( + CacheInitParams { + is_write_back: true, + enable_hicache: true, + has_swa_host_pool: true, + ..swa_params_with_window(1) + }, + vec![FULL, SWA], + ); + let [a, b] = chain::<2>(&mut tc); + set_full_host(&mut tc, a); + set_swa_host(&mut tc, a); + set_full_host(&mut tc, b); + set_swa_host(&mut tc, b); + tc.host_lru_list_mut(SWA).insert_mru(a); + let (kv_xfer, mut comp_xfers) = + tc.build_load_back_spec(tc.arena.node(b).id, /* req = */ None); + assert_eq!( + comp_xfers.get(&SWA).unwrap()[0].nodes_to_load, + Some(vec![tc.arena.node(b).id]) + ); + comp_xfers.get_mut(&SWA).unwrap()[0].device_indices = Some(Tensor::from_slice(&[60i64])); + tc.commit_load_back( + tc.arena.node(b).id, + Tensor::from_slice(&[50i64, 51]), + kv_xfer, + comp_xfers, + ); + + let result = tc.drive_host_eviction(SWA, /* num_tokens = */ 1); + assert_eq!(result.tracker[&SWA], 0); + assert!(result.host_frees.is_empty()); + assert!(tc.arena.has_host_value(a, SWA)); + + tc.finish_load_back(tc.arena.node(b).id); + let result = tc.drive_host_eviction(SWA, /* num_tokens = */ 1); + assert_eq!(result.tracker[&SWA], 1); + assert_eq!(result.host_frees[&SWA].len(), 1); + // Write-back reclaims the loaded node's coexisting host duplicate first. + assert!(tc.arena.has_host_value(a, SWA)); + assert!(!tc.arena.has_host_value(b, SWA)); + tc.sanity_check(&[], &[]); +} + +#[test] +fn build_load_back_spec_degrades_to_empty_on_a_foreign_pin() { + let mut tc: UnifiedTreeCore> = UnifiedTreeCore::new( + CacheInitParams { + is_write_back: true, + enable_hicache: true, + has_swa_host_pool: true, + ..swa_params_with_window(4) + }, + vec![FULL, SWA], + ); + let [a, b] = chain::<2>(&mut tc); + set_full_host(&mut tc, a); + set_swa_host(&mut tc, a); + set_full_host(&mut tc, b); + set_swa_host(&mut tc, b); + // Anchor `a` models a Full-only load whose SWA slice remains host-only. + let (kv_xfer, _comp_xfers) = + tc.build_load_back_spec(tc.arena.node(a).id, /* req = */ None); + tc.commit_load_back( + tc.arena.node(a).id, + Tensor::from_slice(&[50i64]), + kv_xfer, + HashMap::new(), + ); + // Anchor `b` must reject its SWA window because `a` has a foreign pin. + let (kv_xfer, comp_xfers) = tc.build_load_back_spec(tc.arena.node(b).id, /* req = */ None); + assert_eq!(kv_xfer.host_indices.unwrap().numel(), 0); + assert_eq!(kv_xfer.nodes_to_load, Some(vec![])); + assert!(comp_xfers.is_empty()); + tc.finish_load_back(tc.arena.node(a).id); + let (kv_xfer, comp_xfers) = tc.build_load_back_spec(tc.arena.node(b).id, /* req = */ None); + assert_eq!(kv_xfer.nodes_to_load, Some(vec![tc.arena.node(b).id])); + assert_eq!( + comp_xfers.get(&SWA).unwrap()[0].nodes_to_load, + Some(vec![tc.arena.node(a).id, tc.arena.node(b).id]) + ); +} + +#[test] +fn host_drive_reclaims_swa_coexisting_host_values_when_the_host_lru_is_empty() { + let mut tc: UnifiedTreeCore> = UnifiedTreeCore::new( + CacheInitParams { + is_write_back: true, + enable_hicache: true, + has_swa_host_pool: true, + ..swa_params_with_window(4) + }, + vec![FULL, SWA], + ); + tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0)); + tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0)); + let root = tc.arena.root(); + let parent_idx = child_of(&tc, root, &[1]); + let leaf_idx = child_of(&tc, parent_idx, &[3]); + let (parent, leaf) = (tc.arena.node(parent_idx).id, tc.arena.node(leaf_idx).id); + for (handle, slots) in [(parent, vec![30i64, 31]), (leaf, vec![32i64])] { + tc.set_component_device_value(handle, SWA, Tensor::from_slice(&slots)); + } + for (handle, host) in [(parent, vec![20i64, 21]), (leaf, vec![22i64])] { + let swa_xfer = PoolTransfer { + name: PoolName::Swa, + host_indices: Some(Tensor::from_slice(&host)), + ..Default::default() + }; + tc.commit_backup( + handle, + Tensor::from_slice(&host), + HashMap::from([(SWA, vec![swa_xfer])]), + ); + } + assert_eq!(tc.host_lru_list(SWA).len(), 0); + + let mut tracker = swa_tracker(); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + accumulate_step( + tc.drive_host_eviction(SWA, /* num_tokens = */ 2), + &mut tracker, + &mut df, + &mut hf, + ); + assert_eq!(tracker[&SWA], 2); + assert!(!tc.arena.node(parent_idx).has_host_value(SWA)); + assert!(tc.arena.node(parent_idx).has_device_value(SWA)); + assert!(tc.arena.node(parent_idx).has_host_value(FULL)); + assert!(tc.arena.node(leaf_idx).has_host_value(SWA)); + tc.sanity_check(&[], &[]); +} + +fn match_params(key: &Vec) -> MatchPrefixParams<'_, Vec> { + MatchPrefixParams { + key, + namespace: Default::default(), + } +} + +#[test] +fn swa_evict_on_a_full_locked_leaf_tombstones_only_the_swa_slot() { + let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1); + let [a] = chain(&mut tc); + tc.arena + .set_device_value(a, FULL, Tensor::from_slice(&[9i64])); + store_swa_device(&mut tc, a); + // The held Full lock keeps the leaf out of the D-leaf set. + tc.arena + .node_mut(a) + .set_lock_ref_(ValueSlotIdx::device(FULL), 1); + tc.update_evictable_leaf_sets_(a); + assert!(!tc.evictable_device_leaves.contains(a)); + let mut tracker = swa_tracker(); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + tc.evict_device_start(SWA, /* request_cnt = */ 10); + let (next, step) = tc.evict_device_next_node(SWA, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + assert_eq!(next, None); + tc.evict_device_end(SWA); + // The SWA slot tombstoned inline; its frees report the node's Full indices. + assert!(!tc.arena.has_device_value(a, SWA)); + assert!(device_frees[&SWA][0].equal(&Tensor::from_slice(&[9i64]))); + assert_eq!(tracker[&SWA], 1); + assert_eq!(tc.swa_evictable_size(), 0); + assert!(host_frees.is_empty()); + // The higher-tier locked Full is spared and pins the node in the tree. + assert!( + tc.arena + .device_value(a, FULL) + .equal(&Tensor::from_slice(&[9i64])) + ); + assert_eq!(tc.arena.device_lock_ref(a, FULL), 1); + assert_eq!(tc.arena.len(), 2); + assert!(!tc.device_lru_list(SWA).in_list(Some(a))); + assert!(!tc.evictable_device_leaves.contains(a)); +} + +#[test] +fn write_through_offloads_a_boundary_split_leaf() { + let mut tc: UnifiedTreeCore> = UnifiedTreeCore::new( + CacheInitParams { + write_through_threshold: 1, + enable_hicache: true, + ..swa_params_with_window(8) + }, + vec![FULL, SWA], + ); + let result = tc.insert(&insert_params_swa( + &vec![1, 2, 3, 4], + &[10, 11, 12, 13], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 2, + )); + let root = tc.arena.root(); + let parent = child_of(&tc, root, &[1]); + let leaf = child_of(&tc, parent, &[3]); + assert_eq!(tc.arena.node(parent).key, vec![1, 2]); + assert_eq!(tc.arena.node(leaf).key, vec![3, 4]); + // The threshold crossing backs up both split fragments ancestors-first. + let backups: Vec<_> = result + .cache_actions + .iter() + .filter_map(|action| match action { + CacheAction::BackupKV(backup) => Some(backup.node_ids.clone()), + _ => None, + }) + .collect(); + assert_eq!( + backups, + vec![vec![tc.arena.node(parent).id, tc.arena.node(leaf).id]] + ); + tc.commit_backup( + tc.arena.node(parent).id, + Tensor::from_slice(&[100i64, 101]), + HashMap::new(), + ); + tc.commit_backup( + tc.arena.node(leaf).id, + Tensor::from_slice(&[102i64, 103]), + HashMap::new(), + ); + let mut tracker = swa_tracker(); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.evict_device_start(FULL, /* request_cnt = */ 100); + loop { + let (next, step) = tc.evict_device_next_node(FULL, &tracker); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + let Some(next) = next else { break }; + let (backup, step) = tc.evict_device_leaf(next, /* is_write_back = */ false); + assert!(backup.is_none()); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + } + tc.evict_device_end(FULL); + // The split leaf demoted like any backuped leaf and awaits host eviction. + let leaf_node = tc.arena.node(leaf); + assert!(leaf_node.evicted() && leaf_node.backuped()); + assert!(tc.evictable_host_leaves.contains(leaf)); + tc.sanity_check(&[], &[]); +} + +#[test] +fn deep_swa_tree_survives_backup_evict_and_load_back_rounds() { + let mut tc: UnifiedTreeCore> = UnifiedTreeCore::new( + CacheInitParams { + enable_hicache: true, + ..swa_params_with_window(4) + }, + vec![FULL, SWA], + ); + // Varied topology: capped long leaf, decode-evicted chain, depth + // extension, and two branches off the shared prefix. + let inserts: [(Vec, Vec, usize); 5] = [ + (vec![1, 2, 3, 4, 5, 6], vec![10, 11, 12, 13, 14, 15], 0), + ( + vec![21, 22, 23, 24, 25, 26], + vec![30, 31, 32, 33, 34, 35], + 1, + ), + ( + vec![1, 2, 3, 4, 5, 6, 7, 8], + vec![10, 11, 12, 13, 14, 15, 16, 17], + 0, + ), + (vec![1, 2, 41, 42, 43], vec![10, 11, 50, 51, 52], 0), + (vec![1, 2, 61, 62, 63], vec![10, 11, 70, 71, 72], 0), + ]; + for (key, kv, swa_evicted_seqlen) in inserts { + let result = tc.insert(&insert_params_swa(&key, &kv, 0, swa_evicted_seqlen)); + // Apply the emitted rebuilds the way the cache would. + for action in &result.cache_actions { + if let CacheAction::SwaRebuild { + node_id, + source_value, + } = action + { + tc.set_component_device_value(*node_id, SWA, source_value.copy()); + } + } + tc.sanity_check(&[], &[]); + } + let nodes = tc.collect_all_nodes_(); + assert!(nodes.len() >= 6); + // Back up every non-root node so device eviction demotes instead of deleting. + for node in nodes { + if tc.arena.node(node).is_root() { + continue; + } + let len = tc.arena.node(node).key.atom_len(); + tc.commit_backup( + tc.arena.node(node).id, + Tensor::from_slice(&vec![0i64; len]), + HashMap::new(), + ); + } + tc.sanity_check(&[], &[]); + // Stepwise eviction rounds: half the Full budget, then the whole SWA budget. + for _ in 0..4 { + let full_budget = (tc.full_evictable_size() / 2).max(1); + let mut tracker = swa_tracker(); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.evict_device_start(FULL, full_budget); + loop { + let (leaf, step) = tc.evict_device_next_node(FULL, &tracker); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + let Some(leaf) = leaf else { break }; + let (_, step) = tc.evict_device_leaf(leaf, /* is_write_back = */ false); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + } + tc.evict_device_end(FULL); + let swa_budget = tc.swa_evictable_size(); + if swa_budget > 0 { + let mut tracker = swa_tracker(); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.evict_device_start(SWA, swa_budget); + loop { + let (leaf, step) = tc.evict_device_next_node(SWA, &tracker); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + let Some(leaf) = leaf else { break }; + let (_, step) = tc.evict_device_leaf(leaf, /* is_write_back = */ false); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + } + tc.evict_device_end(SWA); + } + tc.sanity_check(&[], &[]); + } + // Load evicted prefixes back from host, mirroring the orchestrator's + // commit-then-lock sequence. + for key in [vec![1i64, 2, 3, 4, 5, 6], vec![1i64, 2]] { + let anchor = tc.match_prefix(&match_params(&key)).best_match_node_id; + if !tc.is_root(anchor) && tc.is_full_device_evicted(anchor) { + let (kv_xfer, comp_xfers) = tc.build_load_back_spec(anchor, /* req = */ None); + let loaded = kv_xfer.host_indices.as_ref().unwrap().numel(); + let actions = tc.commit_load_back( + anchor, + Tensor::from_slice(&vec![0i64; loaded]), + kv_xfer, + comp_xfers, + ); + assert!(actions.is_empty()); + let lock = tc.inc_lock_ref(anchor); + let params = DecLockRefParams { + swa_uuid_for_lock: lock.swa_uuid_for_lock, + swa_uuid_for_host_lock: lock.swa_uuid_for_host_lock, + skip_lock_node_ids: lock.skip_lock_node_ids, + }; + tc.dec_lock_ref(anchor, Some(¶ms), /* skip_swa = */ false); + tc.finish_load_back(anchor); + } + tc.sanity_check(&[], &[]); + } + let matched = tc.match_prefix(&match_params(&vec![1, 2, 3, 4, 5, 6])); + assert_eq!(matched.device_indices.numel(), 6); + tc.sanity_check(&[], &[]); +} + +#[test] +fn recovered_swa_span_evicts_before_the_window_leaf() { + let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); + let key: Vec = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]; + let result = tc.insert(&insert_params_swa( + &key, + &[ + 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, + ], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 8, + )); + let root = tc.arena.root(); + let prefix = child_of(&tc, root, &[1]); + let leaf = child_of(&tc, prefix, &[9]); + assert_eq!(tc.arena.node(prefix).key, vec![1, 2, 3, 4, 5, 6, 7, 8]); + let [ + CacheAction::SwaRebuild { + node_id, + source_value, + }, + ] = result.cache_actions.as_slice() + else { + panic!( + "expected one SwaRebuild action, got {:?}", + action_kinds(&result.cache_actions) + ); + }; + assert_eq!(*node_id, tc.arena.node(leaf).id); + tc.set_component_device_value(*node_id, SWA, source_value.copy()); + assert!(!tc.arena.has_device_value(prefix, SWA)); + + // The fully-in-window re-insert recovers the prefix at its walk barrier. + let step = tc.begin_insert(&insert_params_swa( + &key, + &[ + 200, 201, 202, 203, 204, 205, 206, 207, 208, 209, 210, 211, 212, 213, 214, 215, + ], + /* prev_prefix_len = */ 0, + /* swa_evicted_seqlen = */ 0, + )); + assert!(step.result.is_none()); + let [ + CacheAction::FreeDeviceKVFullOnly(old_full), + CacheAction::SwaRebuild { + node_id, + source_value, + }, + ] = step.actions.as_slice() + else { + panic!( + "expected FreeDeviceKVFullOnly then SwaRebuild, got {:?}", + action_kinds(&step.actions) + ); + }; + assert!(old_full[0].equal(&Tensor::from_slice(&[ + 100i64, 101, 102, 103, 104, 105, 106, 107 + ]))); + assert_eq!(*node_id, tc.arena.node(prefix).id); + tc.set_component_device_value(*node_id, SWA, source_value.copy()); + let done = tc.resume_insert(); + assert_eq!( + done.result.expect("the resumed walk completes").prefix_len, + 16 + ); + // The insert-end window refresh parks the recovered prefix below the leaf. + assert_eq!(swa_lru_order(&tc), vec![leaf, prefix]); + + let mut tracker = swa_tracker(); + let (mut device_frees, mut host_frees) = (HashMap::new(), HashMap::new()); + tc.evict_device_start(SWA, /* request_cnt = */ 8); + let (next, step) = tc.evict_device_next_node(SWA, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + assert_eq!(next, None); + tc.evict_device_end(SWA); + // The recovered span is retaken first; the window leaf survives whole. + assert!(!tc.arena.has_device_value(prefix, SWA)); + assert!(device_frees[&SWA][0].equal(&Tensor::from_slice(&[ + 200i64, 201, 202, 203, 204, 205, 206, 207 + ]))); + assert_eq!(tracker[&SWA], 8); + assert!(tc.arena.has_device_value(leaf, SWA)); + assert!(tc.arena.has_device_value(leaf, FULL)); + tc.sanity_check(&[], &[]); +} diff --git a/rust/mem-cache/src/tests/node.rs b/rust/mem-cache/src/tests/node.rs new file mode 100644 index 000000000..1fd6dd39c --- /dev/null +++ b/rust/mem-cache/src/tests/node.rs @@ -0,0 +1,1942 @@ +use std::collections::HashMap; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use tch::Tensor; + +use super::*; +use crate::components::{FULL, MAMBA, SWA}; +use crate::node::TreeCoreRuntimeError; + +static COUNTED_KEY_CONSTRUCTIONS: AtomicUsize = AtomicUsize::new(0); + +#[derive(Default)] +struct ByteCountingHasher { + bytes_written: usize, +} + +impl std::hash::Hasher for ByteCountingHasher { + fn finish(&self) -> u64 { + 0 + } + + fn write(&mut self, bytes: &[u8]) { + self.bytes_written += bytes.len(); + } +} + +#[derive(Clone, Debug, Default, Eq, Hash, PartialEq)] +struct CountedKey(Vec); + +impl AsRef<[i64]> for CountedKey { + fn as_ref(&self) -> &[i64] { + &self.0 + } +} + +impl Borrow<[i64]> for CountedKey { + fn borrow(&self) -> &[i64] { + &self.0 + } +} + +impl From> for CountedKey { + fn from(value: Vec) -> Self { + COUNTED_KEY_CONSTRUCTIONS.fetch_add(1, Ordering::Relaxed); + Self(value) + } +} + +impl ChildKeyType for CountedKey { + type Atom = i64; + const IS_BIGRAM: bool = false; + + fn key_from(token_ids: Cow<'_, Vec>) -> Cow<'_, Self> { + Cow::Owned(Self::from(token_ids.into_owned())) + } + + fn hash_words(atom: &Self::Atom) -> impl Iterator { + std::iter::once(*atom as u32) + } + + fn raw_token_ids(atoms: &[Self::Atom]) -> Cow<'_, [i64]> { + Cow::Borrowed(atoms) + } +} + +#[test] +fn evicted_only_for_attached_nodes_without_full_device_value() { + let mut parent: Node> = Node::new_root(/* id = */ 0); + assert!(!parent.evicted()); + let mut child = Node::new_child( + /* id = */ 1, + /* key = */ vec![7], + /* priority = */ 0, + ); + // Detached (parent-less) nodes are root-like: never evicted. + assert!(!child.evicted()); + parent + .attach_child(&mut child, /* page_size = */ 1) + .unwrap(); + assert!(child.evicted()); + child.set_device_value(FULL, Tensor::from_slice(&[1i64])); + assert!(!child.evicted()); +} + +#[test] +fn get_last_hash_value_returns_the_final_page_hash() { + let mut node: Node> = Node::new_child( + /* id = */ 1, + /* key = */ vec![7], + /* priority = */ 0, + ); + assert_eq!(node.get_last_hash_value(), None); + node.hash_value = Some(Vec::new()); + assert_eq!(node.get_last_hash_value(), None); + node.hash_value = Some(vec!["h0".to_string(), "h1".to_string()]); + assert_eq!(node.get_last_hash_value(), Some("h1")); +} + +#[test] +fn backuped_tracks_the_full_host_value() { + let mut node: Node> = Node::new_child( + /* id = */ 1, + /* key = */ vec![7], + /* priority = */ 0, + ); + assert!(!node.backuped()); + node.set_host_value(FULL, Tensor::from_slice(&[1i64])); + assert!(node.backuped()); +} + +#[test] +fn attach_child_links_both_sides() { + let mut parent: Node> = Node::new_root(/* id = */ 0); + let mut child = Node::new_child( + /* id = */ 1, + /* key = */ vec![7], + /* priority = */ 0, + ); + child.namespace = KeyNamespace::new(Some("ns"), None); + parent + .attach_child(&mut child, /* page_size = */ 1) + .unwrap(); + assert_eq!(child.parent, Some(NodeIdx_(0))); + // The edge key mirrors the child's namespace. + assert_eq!( + parent + .children + .get(&(KeyNamespace::new(Some("ns"), None), vec![7])), + Some(&NodeIdx_(1)) + ); +} + +#[test] +fn parent_accessors_resolve_the_link() { + let mut parent: Node> = Node::new_root(/* id = */ 0); + let mut child = Node::new_child( + /* id = */ 1, + /* key = */ vec![7], + /* priority = */ 0, + ); + parent + .attach_child(&mut child, /* page_size = */ 1) + .unwrap(); + assert_eq!(child.parent(), NodeIdx_(0)); + assert_eq!(child.try_parent(), Some(NodeIdx_(0))); + assert_eq!(parent.try_parent(), None); +} + +#[test] +#[should_panic(expected = "node 0 is a root and has no parent")] +fn parent_panics_on_a_root() { + let root: Node> = Node::new_root(/* id = */ 0); + root.parent(); +} + +#[test] +#[should_panic(expected = "already attached")] +fn attach_child_panics_on_already_attached() { + let mut parent: Node> = Node::new_root(/* id = */ 0); + let mut child = Node::new_child( + /* id = */ 1, + /* key = */ vec![7], + /* priority = */ 0, + ); + parent + .attach_child(&mut child, /* page_size = */ 1) + .unwrap(); + // Re-attaching an already-attached node is an internal invariant violation. + let mut other: Node> = Node::new_root(/* id = */ 2); + let _ = other.attach_child(&mut child, /* page_size = */ 1); +} + +#[test] +fn attach_child_rejects_duplicate_key() { + let mut parent: Node> = Node::new_root(/* id = */ 0); + let mut a = Node::new_child( + /* id = */ 1, + /* key = */ vec![7], + /* priority = */ 0, + ); + let mut b = Node::new_child( + /* id = */ 2, + /* key = */ vec![7], + /* priority = */ 0, + ); + parent.attach_child(&mut a, /* page_size = */ 1).unwrap(); + assert!(matches!( + parent.attach_child(&mut b, /* page_size = */ 1), + Err(TreeCoreRuntimeError::DuplicateChildKey { parent: p, .. }) if p == 0 + )); + // b was rejected without mutation; a still holds key 7. + assert_eq!( + parent.children.get(&(KeyNamespace::default(), vec![7])), + Some(&NodeIdx_(1)) + ); + assert_eq!(b.parent, None); +} + +#[test] +fn detach_from_parent_unlinks_both_sides() { + let mut parent: Node> = Node::new_root(/* id = */ 0); + let mut child = Node::new_child( + /* id = */ 1, + /* key = */ vec![7], + /* priority = */ 0, + ); + parent + .attach_child(&mut child, /* page_size = */ 1) + .unwrap(); + child.detach_from_parent(&mut parent, /* page_size = */ 1); + assert_eq!(child.parent, None); + assert!(parent.children.is_empty()); +} + +#[test] +#[should_panic(expected = "no child")] +fn detach_from_parent_panics_on_broken_link() { + let mut parent: Node> = Node::new_root(/* id = */ 0); + // `orphan` is not registered under `parent` (never attached) — a broken link. + let mut orphan = Node::new_child( + /* id = */ 1, + /* key = */ vec![7], + /* priority = */ 0, + ); + orphan.detach_from_parent(&mut parent, /* page_size = */ 1); +} + +#[test] +#[should_panic(expected = "(found Some(NodeIdx_(1)))")] +fn detach_from_parent_panics_on_a_mismatched_child_id() { + let mut parent: Node> = Node::new_root(/* id = */ 0); + let mut a = Node::new_child( + /* id = */ 1, + /* key = */ vec![7], + /* priority = */ 0, + ); + parent.attach_child(&mut a, /* page_size = */ 1).unwrap(); + // `b` claims the same page key, but the parent's entry points at `a`. + let mut b = Node::new_child( + /* id = */ 2, + /* key = */ vec![7], + /* priority = */ 0, + ); + b.parent = Some(NodeIdx_(0)); + b.detach_from_parent(&mut parent, /* page_size = */ 1); +} + +#[test] +fn value_slot_idx_addresses_device_then_host() { + assert_eq!(ValueSlotIdx::device(FULL).idx(), 0); + assert_eq!(ValueSlotIdx::device(MAMBA).idx(), 2); + assert_eq!(ValueSlotIdx::host(FULL).idx(), 3); + assert_eq!(ValueSlotIdx::host(MAMBA).idx(), 5); + assert!(!ValueSlotIdx::device(SWA).is_host()); + assert!(ValueSlotIdx::host(SWA).is_host()); + assert_eq!(ValueSlotIdx::host(SWA).component_type(), SWA); +} + +#[test] +fn device_and_host_accessors_address_their_own_tiers() { + let mut node: Node> = Node::new_child( + /* id = */ 1, + /* key = */ vec![7, 8], + /* priority = */ 0, + ); + assert!(!node.has_device_value(SWA)); + node.set_device_value(SWA, Tensor::from_slice(&[1i64, 2])); + assert!(node.has_device_value(SWA)); + assert!(!node.has_host_value(SWA)); + assert_eq!(node.device_value_len(SWA), 2); + assert_eq!(node.host_value_len(SWA), 0); + + node.set_host_value(SWA, Tensor::from_slice(&[3i64, 4])); + assert_eq!(node.host_value(SWA).size()[0], 2); + let taken = node.take_host_value(SWA); + assert_eq!(taken.size()[0], 2); + assert!(node.has_device_value(SWA)); +} + +#[test] +fn lock_predicates_scan_their_own_tier() { + let mut node: Node> = Node::new_child( + /* id = */ 1, + /* key = */ vec![7], + /* priority = */ 0, + ); + assert!(!node.is_device_locked() && !node.is_host_locked()); + node.inc_device_lock_ref(MAMBA); + assert!(node.is_device_locked() && !node.is_host_locked()); + node.dec_device_lock_ref(MAMBA); + node.inc_host_lock_ref(FULL); + assert!(!node.is_device_locked() && node.is_host_locked()); + assert_eq!(node.host_lock_ref(FULL), 1); + assert_eq!(node.device_lock_ref(FULL), 0); +} + +#[test] +fn load_back_pending_predicate_tracks_the_anchor() { + let mut node: Node> = Node::new_child( + /* id = */ 1, + /* key = */ vec![7], + /* priority = */ 0, + ); + assert!(!node.is_load_back_pending()); + node.load_back_pending_id = Some(2); + assert!(node.is_load_back_pending()); +} + +#[test] +#[should_panic(expected = "expects a single state slot")] +fn set_value_rejects_multi_row_mamba_states() { + let mut node: Node> = Node::new_child( + /* id = */ 1, + /* key = */ vec![7], + /* priority = */ 0, + ); + node.set_device_value(MAMBA, Tensor::from_slice(&[1i64, 2])); +} + +#[test] +#[should_panic(expected = "dec_lock_ref: Full/device lock_ref underflow on node 1")] +fn dec_lock_ref_panics_when_unlocked() { + let mut node: Node> = Node::new_child( + /* id = */ 1, + /* key = */ vec![7], + /* priority = */ 0, + ); + node.dec_device_lock_ref(FULL); +} + +#[test] +#[should_panic(expected = "value length differs from the key")] +fn set_value_rejects_a_length_mismatch() { + let mut node: Node> = Node::new_child( + /* id = */ 1, + /* key = */ vec![7, 8], + /* priority = */ 0, + ); + node.set_device_value(FULL, Tensor::from_slice(&[1i64, 2, 3])); +} + +#[test] +#[should_panic(expected = "slot already set")] +fn set_value_rejects_an_occupied_slot() { + let mut node: Node> = Node::new_child( + /* id = */ 1, + /* key = */ vec![7], + /* priority = */ 0, + ); + node.set_device_value(FULL, Tensor::from_slice(&[1i64])); + node.set_device_value(FULL, Tensor::from_slice(&[2i64])); +} + +#[test] +#[should_panic(expected = "has no value")] +fn value_panics_on_an_empty_slot() { + let node: Node> = Node::new_child( + /* id = */ 1, + /* key = */ vec![7], + /* priority = */ 0, + ); + node.device_value(FULL); +} + +#[test] +#[should_panic(expected = "has no value")] +fn take_value_panics_on_an_empty_slot() { + let mut node: Node> = Node::new_child( + /* id = */ 1, + /* key = */ vec![7], + /* priority = */ 0, + ); + let _ = node.take_device_value(FULL); +} + +#[test] +fn try_accessors_return_none_when_unset() { + let mut node: Node> = Node::new_child( + /* id = */ 1, + /* key = */ vec![7], + /* priority = */ 0, + ); + assert!(node.try_device_value(FULL).is_none()); + assert!(node.try_host_value(FULL).is_none()); + node.set_host_value(FULL, Tensor::from_slice(&[9i64])); + assert!(node.try_device_value(FULL).is_none()); + assert_eq!( + Vec::::try_from(node.try_host_value(FULL).unwrap()).unwrap(), + vec![9] + ); +} + +#[test] +fn redistribute_child_value_splits_head_and_tail() { + let mut parent: Node> = Node::new_child( + /* id = */ 1, + /* key = */ vec![7, 8], + /* priority = */ 0, + ); + let mut child: Node> = Node::new_child( + /* id = */ 2, + /* key = */ vec![9], + /* priority = */ 0, + ); + // A mid-split child briefly holds its full pre-split value under the tail key. + child.state_mut_(ValueSlotIdx::device(FULL)).value = Some(Tensor::from_slice(&[1i64, 2, 3])); + Node::redistribute_child_device_value(&mut parent, &mut child, FULL, /* split_len = */ 2); + assert_eq!( + Vec::::try_from(parent.device_value(FULL)).unwrap(), + vec![1, 2] + ); + assert_eq!( + Vec::::try_from(child.device_value(FULL)).unwrap(), + vec![3] + ); +} + +#[test] +#[should_panic(expected = "out of range")] +fn redistribute_child_value_rejects_a_boundary_split() { + let mut parent: Node> = Node::new_child( + /* id = */ 1, + /* key = */ vec![7, 8, 9], + /* priority = */ 0, + ); + let mut child: Node> = Node::new_child( + /* id = */ 2, + /* key = */ vec![10], + /* priority = */ 0, + ); + child.state_mut_(ValueSlotIdx::device(FULL)).value = Some(Tensor::from_slice(&[1i64, 2, 3])); + Node::redistribute_child_device_value(&mut parent, &mut child, FULL, /* split_len = */ 3); +} + +#[test] +fn copy_device_lock_ref_copies_between_nodes() { + let mut src: Node> = Node::new_child( + /* id = */ 1, + /* key = */ vec![7], + /* priority = */ 0, + ); + let mut dst: Node> = Node::new_child( + /* id = */ 2, + /* key = */ vec![8], + /* priority = */ 0, + ); + src.set_lock_ref_(ValueSlotIdx::device(SWA), 3); + dst.copy_device_lock_ref(SWA, &src); + assert_eq!(dst.device_lock_ref(SWA), 3); + assert_eq!(dst.host_lock_ref(SWA), 0); +} + +#[test] +fn mamba_single_state_skips_the_key_length_check() { + let mut node: Node> = Node::new_child( + /* id = */ 1, + /* key = */ vec![7, 8, 9], + /* priority = */ 0, + ); + // One state slot regardless of the key length. + node.set_device_value(MAMBA, Tensor::from_slice(&[42i64])); + assert_eq!(node.device_value_len(MAMBA), 1); +} + +#[test] +fn present_but_empty_tensor_has_value_true_len_zero() { + // A zero-length tensor is still "present" (matches Python `value is not None`). + let mut node: Node> = Node::new_root(/* id = */ 0); + let empty: [i64; 0] = []; + node.state_mut_(ValueSlotIdx::device(FULL)).value = Some(Tensor::from_slice(&empty)); + assert!(node.has_device_value(FULL)); + assert_eq!(node.device_value_len(FULL), 0); +} + +#[test] +#[should_panic(expected = "out of range")] +fn redistribute_child_value_rejects_a_zero_split() { + let mut parent: Node> = Node::new_child( + /* id = */ 1, + /* key = */ vec![7, 8, 9], + /* priority = */ 0, + ); + let mut child: Node> = Node::new_child( + /* id = */ 2, + /* key = */ vec![10], + /* priority = */ 0, + ); + child.state_mut_(ValueSlotIdx::device(FULL)).value = Some(Tensor::from_slice(&[1i64, 2, 3])); + Node::redistribute_child_device_value(&mut parent, &mut child, FULL, /* split_len = */ 0); +} + +#[test] +#[should_panic(expected = "no value")] +fn redistribute_child_value_panics_when_child_is_value_less() { + let mut parent: Node> = Node::new_child( + /* id = */ 1, + /* key = */ vec![7, 8, 9], + /* priority = */ 0, + ); + let mut child: Node> = Node::new_child( + /* id = */ 2, + /* key = */ vec![10], + /* priority = */ 0, + ); + Node::redistribute_child_device_value(&mut parent, &mut child, FULL, /* split_len = */ 1); +} + +// Node handles and per-slot value state. + +#[test] +fn from_idx_round_trips_every_slot() { + assert_eq!(ValueSlotIdx::from_idx(0), ValueSlotIdx::device(FULL)); + assert_eq!(ValueSlotIdx::from_idx(1), ValueSlotIdx::device(SWA)); + assert_eq!(ValueSlotIdx::from_idx(2), ValueSlotIdx::device(MAMBA)); + assert_eq!(ValueSlotIdx::from_idx(3), ValueSlotIdx::host(FULL)); + assert_eq!(ValueSlotIdx::from_idx(4), ValueSlotIdx::host(SWA)); + assert_eq!(ValueSlotIdx::from_idx(5), ValueSlotIdx::host(MAMBA)); +} + +#[test] +#[should_panic(expected = "from_idx: 6 is not a value-slot index")] +fn from_idx_panics_out_of_range() { + ValueSlotIdx::from_idx(NUM_VALUE_SLOTS); +} + +// Unigram and bigram child keys. + +#[test] +fn raw_token_ids_borrow_unigram_atoms() { + let raw = as ChildKeyType>::raw_token_ids(&[1, 2, 3]); + assert_eq!(raw.as_ref(), &[1, 2, 3]); + assert!(matches!(raw, Cow::Borrowed(_))); +} + +#[test] +fn raw_token_ids_unzip_overlapping_bigram_atoms() { + assert_eq!( + as ChildKeyType>::raw_token_ids(&[(1, 2), (2, 3), (3, 4)]).as_ref(), + &[1, 2, 3, 4] + ); + assert_eq!( + as ChildKeyType>::raw_token_ids(&[]).as_ref(), + &[] as &[i64] + ); +} + +#[test] +fn child_key_takes_the_first_page() { + let key: Vec = vec![1, 2, 3]; + assert_eq!(key.child_key(/* page_size = */ 1), vec![1]); + assert_eq!(key.child_key(/* page_size = */ 2), vec![1, 2]); + let bigram: Vec<(i64, i64)> = vec![(1, 2), (2, 3)]; + assert_eq!(bigram.child_key(/* page_size = */ 1), vec![(1, 2)]); +} + +#[test] +fn match_len_rounds_down_to_page_multiples() { + let key: Vec = vec![1, 2, 3, 4]; + assert_eq!( + key.match_len( + /* start = */ 0, + &vec![1, 2, 3, 9], + /* page_size = */ 1 + ), + 3 + ); + assert_eq!( + key.match_len( + /* start = */ 0, + &vec![1, 2, 3, 9], + /* page_size = */ 2 + ), + 2 + ); + assert_eq!( + key.match_len( + /* start = */ 0, + &vec![9, 2, 3, 4], + /* page_size = */ 1 + ), + 0 + ); + assert_eq!( + key.match_len( + /* start = */ 0, + &vec![1, 2, 3, 4], + /* page_size = */ 1 + ), + 4 + ); + // The shorter key bounds the comparison. + assert_eq!( + key.match_len(/* start = */ 0, &vec![1, 2], /* page_size = */ 1), + 2 + ); +} + +#[test] +fn match_len_from_compares_the_tail_at_start() { + let key: Vec = vec![1, 2, 3, 4]; + assert_eq!( + key.match_len( + /* start = */ 0, + &vec![1, 2, 9], + /* page_size = */ 1 + ), + 2 + ); + assert_eq!( + key.match_len( + /* start = */ 1, + &vec![2, 3, 9], + /* page_size = */ 1 + ), + 2 + ); + assert_eq!( + key.match_len(/* start = */ 3, &vec![4], /* page_size = */ 1), + 1 + ); + assert_eq!( + key.match_len(/* start = */ 1, &vec![9], /* page_size = */ 1), + 0 + ); +} + +#[test] +fn match_len_from_rounds_down_to_page_multiples() { + let key: Vec = vec![1, 2, 3, 4, 5]; + assert_eq!( + key.match_len( + /* start = */ 1, + &vec![2, 3, 4, 9], + /* page_size = */ 2 + ), + 2 + ); + assert_eq!( + key.match_len(/* start = */ 1, &vec![2, 9], /* page_size = */ 2), + 0 + ); + assert_eq!( + key.match_len( + /* start = */ 1, + &vec![2, 3, 4, 5], + /* page_size = */ 2 + ), + 4 + ); +} + +#[test] +fn match_len_from_is_empty_at_the_tail_boundary() { + let key: Vec = vec![1, 2, 3]; + assert_eq!( + key.match_len( + /* start = */ 3, + &vec![1, 2, 3], + /* page_size = */ 1 + ), + 0 + ); +} + +#[test] +#[should_panic(expected = "match_len: start 4 beyond the key length 3")] +fn match_len_from_panics_beyond_the_key() { + let key: Vec = vec![1, 2, 3]; + key.match_len(/* start = */ 4, &vec![1], /* page_size = */ 1); +} + +#[test] +fn match_len_from_compares_bigram_atoms() { + let key: Vec<(i64, i64)> = vec![(1, 2), (2, 3), (3, 4)]; + assert_eq!( + key.match_len( + /* start = */ 1, + &vec![(2, 3), (3, 9)], + /* page_size = */ 1 + ), + 1 + ); + assert_eq!( + key.match_len( + /* start = */ 1, + &vec![(2, 9)], + /* page_size = */ 1 + ), + 0 + ); +} + +#[test] +fn page_at_borrows_the_page_at_start() { + let key: Vec = vec![1, 2, 3, 4]; + assert_eq!(key.page_at(/* start = */ 0, /* page_size = */ 2), &[1, 2]); + assert_eq!(key.page_at(/* start = */ 2, /* page_size = */ 2), &[3, 4]); + assert_eq!(key.page_at(/* start = */ 3, /* page_size = */ 1), &[4]); + let bigram: Vec<(i64, i64)> = vec![(1, 2), (2, 3)]; + assert_eq!( + bigram.page_at(/* start = */ 1, /* page_size = */ 1), + &[(2, 3)] + ); +} + +#[test] +#[should_panic(expected = "page_at: page [2, 4) reaches beyond the key length 3")] +fn page_at_panics_past_the_end() { + let key: Vec = vec![1, 2, 3]; + key.page_at(/* start = */ 2, /* page_size = */ 2); +} + +#[test] +fn page_at_keys_a_child_map_lookup() { + let mut children: HashMap, usize> = HashMap::new(); + children.insert(vec![3, 4], 7); + let key: Vec = vec![1, 2, 3, 4]; + assert_eq!( + children.get(key.page_at(/* start = */ 2, /* page_size = */ 2)), + Some(&7) + ); + assert_eq!( + children.get(key.page_at(/* start = */ 0, /* page_size = */ 2)), + None + ); +} + +#[test] +fn suffix_takes_the_tail_and_allows_the_boundary() { + let key: Vec = vec![1, 2, 3]; + assert_eq!(key.suffix(1), vec![2, 3]); + assert_eq!(key.suffix(3), Vec::::new()); + assert_eq!(key.suffix(0), vec![1, 2, 3]); +} + +#[test] +#[should_panic(expected = "suffix: start 4 beyond the key length 3")] +fn suffix_panics_beyond_the_key() { + let key: Vec = vec![1, 2, 3]; + key.suffix(4); +} + +#[test] +fn page_aligned_truncates_to_whole_pages() { + let key: Vec = vec![1, 2, 3]; + assert_eq!(key.page_aligned(/* page_size = */ 1), vec![1, 2, 3]); + assert_eq!(key.page_aligned(/* page_size = */ 2), vec![1, 2]); + assert_eq!(key.page_aligned(/* page_size = */ 4), Vec::::new()); +} + +#[test] +#[should_panic(expected = "child_key: key of 3 atoms is shorter than a page (4)")] +fn child_key_panics_on_a_key_shorter_than_a_page() { + let key: Vec = vec![1, 2, 3]; + key.child_key(/* page_size = */ 4); +} + +#[test] +fn split_at_partitions_the_key() { + let key: Vec = vec![1, 2, 3]; + assert_eq!(key.split_at(/* split_idx = */ 1), (vec![1], vec![2, 3])); + assert_eq!(key.split_at(/* split_idx = */ 2), (vec![1, 2], vec![3])); +} + +#[test] +#[should_panic(expected = "split_at: split_idx 3 out of range (0, 3)")] +fn split_at_panics_on_the_tail_boundary() { + let key: Vec = vec![1, 2, 3]; + key.split_at(/* split_idx = */ 3); +} + +#[test] +#[should_panic(expected = "split_at: split_idx 0 out of range (0, 3)")] +fn split_at_panics_on_the_head_boundary() { + let key: Vec = vec![1, 2, 3]; + key.split_at(/* split_idx = */ 0); +} + +#[test] +#[should_panic(expected = "child_key: key of 0 atoms is shorter than a page (1)")] +fn child_key_panics_on_an_empty_key() { + let key: Vec = vec![]; + key.child_key(/* page_size = */ 1); +} + +#[test] +fn unigram_key_from_passes_ownership_through() { + let ids = vec![1i64, 2, 3]; + assert!(matches!( + >::key_from(Cow::Borrowed(&ids)), + Cow::Borrowed(_) + )); + assert_eq!( + >::key_from(Cow::Owned(vec![1, 2, 3])).into_owned(), + vec![1, 2, 3] + ); +} + +#[test] +fn bigram_key_from_pairs_overlapping_tokens() { + assert_eq!( + >::key_from(Cow::Owned(vec![1, 2, 3, 4])).into_owned(), + vec![(1, 2), (2, 3), (3, 4)] + ); + assert_eq!( + >::key_from(Cow::Owned(vec![5, 6])).into_owned(), + vec![(5, 6)] + ); +} + +#[test] +fn bigram_key_from_is_empty_below_one_pair() { + assert_eq!( + >::key_from(Cow::Owned(vec![])).into_owned(), + Vec::<(i64, i64)>::new() + ); + assert_eq!( + >::key_from(Cow::Owned(vec![7])).into_owned(), + Vec::<(i64, i64)>::new() + ); +} + +// Per-page hash chains. + +// Expected values are literals produced by the python native hash +// (mem_cache/utils.py::get_hash_str over cpp_utils/hash_binding.cpp). + +#[test] +fn unigram_pages_chain_within_the_node() { + assert_eq!( + get_hash_str::>(&[1, 2, 3, 4, 5], None, 2), + vec![ + "34fb5c825de7ca4aea6e712f19d439c1da0c92c37b423936c5f618545ca4fa1f", + "c57b445f90651b9a650e516ab2238c965b21af35608a31c303e6d9e407f2915c", + "e1fd781b60f933e64fe17100c521f3130659b4761e25f7b79abcf81ae5aa23cf", + ] + ); +} + +#[test] +fn single_page_covers_the_whole_key() { + assert_eq!( + get_hash_str::>(&[1, 2, 3], None, 4), + vec!["4636993d3e1da4e9d6b8f87b79e8f7c6d018580d52661950eabc3845c5897a4d"] + ); +} + +#[test] +fn parent_hash_chains_across_nodes() { + let parent = get_hash_str::>(&[1, 2, 3], None, 2); + assert_eq!( + parent, + vec![ + "34fb5c825de7ca4aea6e712f19d439c1da0c92c37b423936c5f618545ca4fa1f", + "3ac5d352be720e428633fe34fc74591b2b80060a7764e195d8f34068203ae98f", + ] + ); + assert_eq!( + get_hash_str::>(&[7, 8], parent.last().map(String::as_str), 2), + vec!["5d13a4cf14ad5f9dbd0da79d004bd2a339f633be6100067dc3987c1880bcf0dc"] + ); +} + +#[test] +fn bigram_atoms_hash_as_word_pairs() { + // Raw ids [1, 2, 3, 4, 5] as overlapping pairs, two pairs per page. + let atoms: Vec<(i64, i64)> = vec![(1, 2), (2, 3), (3, 4), (4, 5)]; + assert_eq!( + get_hash_str::>(&atoms, None, 2), + vec![ + "ede8ef26a097f0fd889f9f39f6f5af921370630b164c4a4fa28eb716b2df9269", + "678c5307f2a02a2aaf8edbeb80ebac13663f70e14a4fc8d5d62b77e702841e4b", + ] + ); +} + +#[test] +fn empty_key_yields_no_pages() { + assert_eq!(get_hash_str::>(&[], None, 2), Vec::::new()); +} + +#[test] +#[should_panic(expected = "token id does not fit in uint32")] +fn oversized_token_id_is_rejected() { + get_hash_str::>(&[1 << 33], None, 2); +} + +#[test] +#[should_panic(expected = "token id does not fit in uint32")] +fn negative_token_id_is_rejected() { + get_hash_str::>(&[-1], None, 2); +} + +#[test] +fn empty_prior_hash_chains_nothing() { + assert_eq!( + get_hash_str::>(&[1, 2, 3], Some(""), 4), + get_hash_str::>(&[1, 2, 3], None, 4) + ); +} + +#[test] +fn hash_str_to_int64_takes_the_first_sixteen_hex_chars_signed() { + let hex = format!("34fb5c825de7ca4a{}", "0".repeat(48)); + assert_eq!(hash_str_to_int64(&hex), 0x34fb5c825de7ca4a_i64); + assert_eq!(hash_str_to_int64(&"f".repeat(64)), -1); +} + +#[test] +#[should_panic(expected = "byte index 16 is out of bounds")] +fn hash_str_to_int64_panics_on_a_short_string() { + hash_str_to_int64("abc"); +} + +#[test] +#[should_panic(expected = "hash must be a hex digest")] +fn hash_str_to_int64_panics_on_a_non_hex_string() { + hash_str_to_int64(&"z".repeat(64)); +} + +#[test] +#[should_panic(expected = "prior hash contains a non-hex character")] +fn get_hash_str_panics_on_a_non_hex_prior() { + get_hash_str::>(&[1, 2], Some(&"z".repeat(64)), 2); +} + +#[test] +#[should_panic(expected = "prior hash must be a 64-char hex digest")] +fn get_hash_str_panics_on_a_wrong_length_prior() { + get_hash_str::>(&[1, 2], Some("abcd"), 2); +} + +#[test] +#[should_panic(expected = "page_size must be positive")] +fn get_hash_str_panics_on_a_zero_page_size() { + get_hash_str::>(&[1, 2], None, 0); +} + +#[test] +fn split_redistributes_page_hashes() { + let hashes = vec!["a".to_string(), "b".to_string(), "c".to_string()]; + let (head, tail) = split_node_hash_value(Some(hashes), 4, 2); + assert_eq!(head, Some(vec!["a".to_string(), "b".to_string()])); + assert_eq!(tail, Some(vec!["c".to_string()])); +} + +#[test] +fn split_of_an_unhashed_node_stays_none() { + assert_eq!(split_node_hash_value(None, 4, 2), (None, None)); +} + +// Node arena storage. + +fn arena() -> NodeArena> { + NodeArena::new(vec![FULL], /* page_size = */ 1) +} + +#[test] +fn reset_installs_protected_valueless_root() -> Result<(), TreeCoreRuntimeError> { + let arena = arena(); + let root_id = arena.root(); + let node = arena.node(root_id); + assert!(node.is_root()); + assert!(node.is_leaf()); + assert!(node.values[FULL.idx()].value.is_none()); + assert_eq!(node.values[FULL.idx()].lock_ref, 1); + assert_eq!(node.priority, i64::MIN); + assert_eq!(arena.len(), 1); + Ok(()) +} + +#[test] +fn reset_seeds_lock_ref_for_each_enabled_component() -> Result<(), TreeCoreRuntimeError> { + let arena: NodeArena> = NodeArena::new(vec![FULL, SWA], /* page_size = */ 1); + let root_id = arena.root(); + let root = arena.node(root_id); + assert_eq!(root.values[FULL.idx()].lock_ref, 1); + assert_eq!(root.values[SWA.idx()].lock_ref, 1); + assert_eq!(root.values[MAMBA.idx()].lock_ref, 0); + Ok(()) +} + +#[test] +fn alloc_free_recycles_slots() -> Result<(), TreeCoreRuntimeError> { + let mut arena = arena(); + let root = arena.root(); + let a = arena.alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 0, + /* extra_key = */ None, + )?; + assert_eq!(arena.len(), 2); + arena.free_leaf(a)?; + assert_eq!(arena.len(), 1); + let b = arena.alloc_child( + root, + /* key = */ vec![3, 4], + /* priority = */ 0, + /* extra_key = */ None, + )?; + assert_eq!(a, b); + assert_eq!(arena.node(b).key, vec![3, 4]); + Ok(()) +} + +#[test] +fn alloc_detached_reuses_a_freed_slot() -> Result<(), TreeCoreRuntimeError> { + let mut arena = arena(); + let root = arena.root(); + let a = arena.alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + )?; + arena.free_leaf(a)?; + let b = arena.alloc_detached(/* priority = */ 0); + assert_eq!(b, a); + assert_eq!(arena.len(), 2); + Ok(()) +} + +#[test] +fn reset_clears_then_reinstalls_root() -> Result<(), TreeCoreRuntimeError> { + let mut arena = arena(); + let root = arena.root(); + arena.alloc_child( + root, + /* key = */ vec![9], + /* priority = */ 0, + /* extra_key = */ None, + )?; + assert_eq!(arena.len(), 2); + arena.reset(); + assert_eq!(arena.len(), 1); + let root_id = arena.root(); + assert!(arena.node(root_id).is_root()); + Ok(()) +} + +#[test] +fn get_and_bump_access_counter_is_monotonic() { + let mut arena = arena(); + // The root's construction consumed tick 1. + assert_eq!(arena.get_and_bump_access_counter(), 2); + assert_eq!(arena.get_and_bump_access_counter(), 3); +} + +#[test] +fn salted_children_file_under_their_namespace() -> Result<(), TreeCoreRuntimeError> { + let mut arena = arena(); + let root = arena.root(); + let a = arena.alloc_child( + root, + /* key = */ vec![7], + /* priority = */ 0, + Some("lora-1"), + )?; + assert_eq!(arena.node(a).namespace.extra_key(), Some("lora-1")); + // The same page key resolves independently per namespace. + let b = arena.alloc_child( + root, + /* key = */ vec![7], + /* priority = */ 0, + Some("lora-2"), + )?; + let c = arena.alloc_child( + root, + /* key = */ vec![7], + /* priority = */ 0, + None, + )?; + assert_eq!(arena.root_child(Some("lora-1"), &[7]), Some(a)); + assert_eq!(arena.root_child(Some("lora-2"), &[7]), Some(b)); + assert_eq!(arena.root_child(None, &[7]), Some(c)); + assert_eq!(arena.root_child(Some("ghost"), &[7]), None); + Ok(()) +} + +#[test] +fn cache_salt_is_a_distinct_child_namespace_dimension() -> Result<(), TreeCoreRuntimeError> { + let mut arena = arena(); + let root = arena.root(); + let first_namespace = KeyNamespaceRef::new(Some("bc"), Some("a")); + let second_namespace = KeyNamespaceRef::new(Some("c"), Some("ab")); + let first = arena.alloc_child_in_namespace(root, vec![7], 0, first_namespace)?; + let second = arena.alloc_child_in_namespace(root, vec![7], 0, second_namespace)?; + + assert_eq!( + arena.root_child_in_namespace(first_namespace, &[7]), + Some(first) + ); + assert_eq!( + arena.root_child_in_namespace(second_namespace, &[7]), + Some(second) + ); + assert_eq!(arena.root_child(None, &[7]), None); + assert_eq!(arena.node(first).namespace.as_ref(), first_namespace); + assert_eq!(arena.node(second).namespace.as_ref(), second_namespace); + assert_eq!( + KeyNamespaceRef::new(None, Some("")).to_owned(), + KeyNamespace::default() + ); + Ok(()) +} + +#[test] +fn namespace_hashing_uses_the_cached_digest_but_equality_checks_strings() { + let long_extra_key = "x".repeat(64 * 1024); + let long_cache_salt = "y".repeat(64 * 1024); + let namespace = KeyNamespace::new(Some(&long_extra_key), Some(&long_cache_salt)); + + let mut owned_hasher = ByteCountingHasher::default(); + std::hash::Hash::hash(&namespace, &mut owned_hasher); + assert_eq!(owned_hasher.bytes_written, size_of::()); + + let mut borrowed_hasher = ByteCountingHasher::default(); + std::hash::Hash::hash(&namespace.as_ref(), &mut borrowed_hasher); + assert_eq!(borrowed_hasher.bytes_written, size_of::()); + + let first = KeyNamespaceRef { + extra_key: Some("adapter-a"), + cache_salt: Some("tenant-a"), + hash: 7, + }; + let second = KeyNamespaceRef { + extra_key: Some("adapter-b"), + cache_salt: Some("tenant-b"), + hash: 7, + }; + assert_ne!(first, second); + + let page = vec![1]; + let mut children: ChildMap> = ChildMap::with_hasher(RandomState::new()); + children.insert((first.to_owned(), page.clone()), NodeIdx_(1)); + children.insert((second.to_owned(), page.clone()), NodeIdx_(2)); + assert_eq!( + children.get(&ChildEdgeRef::> { + namespace: first, + page: &page, + }), + Some(&NodeIdx_(1)) + ); + assert_eq!( + children.get(&ChildEdgeRef::> { + namespace: second, + page: &page, + }), + Some(&NodeIdx_(2)) + ); +} + +#[test] +fn child_page_lookup_borrows_namespace_and_page_key() -> Result<(), TreeCoreRuntimeError> { + let mut arena: NodeArena = NodeArena::new(vec![FULL], /* page_size = */ 2); + let root = arena.root(); + let default_child = arena.alloc_child( + root, + CountedKey(vec![1, 2, 3]), + /* priority = */ 0, + /* extra_key = */ None, + )?; + let namespaced_child = arena.alloc_child( + root, + CountedKey(vec![1, 2, 4]), + /* priority = */ 0, + Some("adapter-a"), + )?; + + COUNTED_KEY_CONSTRUCTIONS.store(0, Ordering::Relaxed); + for _ in 0..100 { + assert_eq!(arena.root_child(None, &[1, 2]), Some(default_child)); + assert_eq!( + arena.root_child(Some("adapter-a"), &[1, 2]), + Some(namespaced_child) + ); + assert_eq!(arena.root_child(Some("adapter-b"), &[1, 2]), None); + assert_eq!(arena.root_child(None, &[2, 3]), None); + } + assert_eq!(COUNTED_KEY_CONSTRUCTIONS.load(Ordering::Relaxed), 0); + Ok(()) +} + +#[test] +fn node_extra_key_propagates_down_the_chain() -> Result<(), TreeCoreRuntimeError> { + let mut arena = arena(); + let root = arena.root(); + let child = arena.alloc_child( + root, + /* key = */ vec![7], + /* priority = */ 0, + Some("chat"), + )?; + let grandchild = arena.alloc_child( + child, + /* key = */ vec![8], + /* priority = */ 0, + /* extra_key = */ None, + )?; + assert_eq!(arena.node_extra_key(grandchild), Some("chat")); + assert_eq!(arena.node_extra_key(arena.root()), None); + let detached = arena.alloc_detached(/* priority = */ 0); + assert_eq!(arena.node_extra_key(detached), None); + Ok(()) +} + +#[test] +fn reset_clears_namespace_edges() { + let mut arena = arena(); + let root = arena.root(); + arena + .alloc_child( + root, + /* key = */ vec![7], + /* priority = */ 0, + Some("lora-1"), + ) + .unwrap(); + assert_eq!(arena.len(), 2); + arena.reset(); + assert_eq!(arena.len(), 1); + assert!(!arena.namespace_exists(Some("lora-1"))); +} + +#[test] +fn alloc_child_sets_child_contract() -> Result<(), TreeCoreRuntimeError> { + let mut arena = arena(); + let root = arena.root(); + let c = arena.alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 7, + /* extra_key = */ None, + )?; + let node = arena.node(c); + assert!(!node.is_root()); + assert!(node.is_leaf()); + assert_eq!(node.parent, Some(root)); + assert_eq!(node.priority, 7); + assert_eq!(node.key, vec![1, 2]); + assert!(node.values[FULL.idx()].value.is_none()); + assert_eq!(node.values[FULL.idx()].lock_ref, 0); + Ok(()) +} + +#[test] +fn alloc_stamps_self_id_and_a_fresh_access_tick() -> Result<(), TreeCoreRuntimeError> { + let mut arena = arena(); + let root = arena.root(); + // Handles mint monotonically and the map resolves each back to its slot. + assert_eq!(arena.node(root).id, 0); + let a = arena.alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 0, + /* extra_key = */ None, + )?; + let b = arena.alloc_child( + root, + /* key = */ vec![3, 4], + /* priority = */ 0, + /* extra_key = */ None, + )?; + assert_eq!(arena.node(a).id, 1); + assert_eq!(arena.node(b).id, 2); + assert_eq!(arena.resolve(arena.node(a).id), a); + // Construction stamps strictly increasing ticks: root, then a, then b; + // both stamps share the node's single construction tick. + let root_tick = arena.node(root).last_access_counter; + let a_tick = arena.node(a).last_access_counter; + let b_tick = arena.node(b).last_access_counter; + assert!(root_tick > 0); + assert!(root_tick < a_tick); + assert!(a_tick < b_tick); + assert_eq!(arena.node(root).creation_counter, root_tick); + assert_eq!(arena.node(a).creation_counter, a_tick); + assert_eq!(arena.node(b).creation_counter, b_tick); + Ok(()) +} + +#[test] +fn reset_zeroes_access_counter() { + let mut arena = arena(); + arena.get_and_bump_access_counter(); + arena.get_and_bump_access_counter(); + arena.reset(); + // The counter restarts and the fresh root's stamp consumes tick 1. + assert_eq!(arena.get_and_bump_access_counter(), 2); +} + +#[test] +fn node_pair_mut_returns_a_live_parent_child_pair() -> Result<(), TreeCoreRuntimeError> { + let mut arena = arena(); + let root = arena.root(); + let parent = arena.alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + )?; + let child = arena.alloc_child( + parent, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + )?; + let parent_ext = arena.node(parent).id; + let child_ext = arena.node(child).id; + let (parent_node, child_node) = arena.node_pair_mut(parent, child); + assert_eq!(parent_node.id, parent_ext); + assert_eq!(child_node.id, child_ext); + Ok(()) +} + +#[test] +#[should_panic(expected = "distinct nodes required")] +fn node_pair_mut_panics_on_same_id() { + let mut arena = arena(); + let root = arena.root(); + let a = arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let _ = arena.node_pair_mut(a, a); +} + +#[test] +#[should_panic(expected = "is not a child of")] +fn node_pair_mut_panics_when_not_parent_and_child() { + let mut arena = arena(); + let root = arena.root(); + let a = arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let b = arena + .alloc_child( + root, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let _ = arena.node_pair_mut(a, b); +} + +#[test] +#[should_panic(expected = "live node")] +fn node_pair_mut_panics_on_freed_id() { + let mut arena = arena(); + let root = arena.root(); + let parent = arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let child = arena + .alloc_child( + parent, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + arena.free_leaf(child).unwrap(); + let _ = arena.node_pair_mut(parent, child); +} + +#[test] +#[should_panic(expected = "out of bounds")] +fn node_pair_mut_panics_on_out_of_bounds_id() { + let mut arena = arena(); + let root = arena.root(); + let parent = arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let _ = arena.node_pair_mut(parent, NodeIdx_(999)); +} + +#[test] +#[should_panic(expected = "node 1 is not allocated")] +fn node_panics_on_a_freed_id() { + let mut arena = arena(); + let root = arena.root(); + let a = arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + arena.free_leaf(a).unwrap(); + arena.node(a); +} + +#[test] +#[should_panic(expected = "node 1 is not allocated")] +fn node_mut_panics_on_a_freed_id() { + let mut arena = arena(); + let root = arena.root(); + let a = arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + arena.free_leaf(a).unwrap(); + arena.node_mut(a); +} + +#[test] +#[should_panic(expected = "id 999 not in [0, 1)")] +fn node_panics_on_an_out_of_bounds_id() { + let arena = arena(); // one slot: the root + arena.node(NodeIdx_(999)); +} + +#[test] +#[should_panic(expected = "id 999 not in [0, 1)")] +fn node_mut_panics_on_an_out_of_bounds_id() { + let mut arena = arena(); // one slot: the root + arena.node_mut(NodeIdx_(999)); +} + +#[test] +fn out_of_bounds_id_returns_out_of_bound_err() { + let mut arena = arena(); // one slot: the root + let bogus = NodeIdx_(999); + assert!(matches!( + arena.free_leaf(bogus), + Err(TreeCoreRuntimeError::NodeAccessOutOfBound { id, .. }) if id == bogus + )); + assert!(matches!( + arena.alloc_child(bogus, /* key = */ vec![1], /* priority = */ 0, /* extra_key = */ None), + Err(TreeCoreRuntimeError::NodeAccessOutOfBound { id, .. }) if id == bogus + )); +} + +#[test] +fn double_free_returns_err() -> Result<(), TreeCoreRuntimeError> { + let mut arena = arena(); + let root = arena.root(); + let a = arena.alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + )?; + arena.free_leaf(a)?; + assert!(matches!( + arena.free_leaf(a), + Err(TreeCoreRuntimeError::NodeDoubleFree { id }) if id == a + )); + // The rejected free did not re-push the slot onto the freelist. + assert_eq!(arena.len(), 1); + Ok(()) +} + +#[test] +fn free_root_returns_err() -> Result<(), TreeCoreRuntimeError> { + let mut arena = arena(); + let root = arena.root(); + assert!(matches!( + arena.free_leaf(root), + Err(TreeCoreRuntimeError::RootNotFreeable { id }) if id == root + )); + // The root survives and stays accessible. + assert!(arena.node(root).is_root()); + assert_eq!(arena.len(), 1); + Ok(()) +} + +#[test] +fn free_the_root_returns_err() { + let mut arena = arena(); + let r = arena.root(); + assert!(matches!( + arena.free_leaf(r), + Err(TreeCoreRuntimeError::RootNotFreeable { id }) if id == r + )); +} + +#[test] +fn freeing_the_last_salted_child_drops_the_namespace() -> Result<(), TreeCoreRuntimeError> { + let mut arena = arena(); + let root = arena.root(); + let a = arena.alloc_child( + root, + /* key = */ vec![7], + /* priority = */ 0, + Some("chat"), + )?; + assert!(arena.namespace_exists(Some("chat"))); + arena.free_leaf(a)?; + // An emptied salted namespace leaves nothing behind. + assert!(!arena.namespace_exists(Some("chat"))); + assert_eq!(arena.len(), 1); + Ok(()) +} + +#[test] +fn free_node_with_children_returns_err() -> Result<(), TreeCoreRuntimeError> { + let mut arena = arena(); + let root = arena.root(); + let parent = arena.alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + )?; + // alloc_child wires this under `parent`, so `parent` is no longer a leaf. + arena.alloc_child( + parent, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + )?; + assert!(matches!( + arena.free_leaf(parent), + Err(TreeCoreRuntimeError::FreeNonLeafNode { id, num_children }) + if id == parent && num_children == 1 + )); + Ok(()) +} + +#[test] +fn alloc_child_wires_into_parent() -> Result<(), TreeCoreRuntimeError> { + let mut arena = arena(); + let root = arena.root(); + let a = arena.alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 0, + /* extra_key = */ None, + )?; + // Both link sides are set: parent.children[first page] -> a, and a.parent -> root. + assert_eq!(arena.root_child(None, &[1]).as_ref(), Some(&a)); + assert_eq!(arena.node(a).parent, Some(root)); + Ok(()) +} + +#[test] +fn free_leaf_detaches_from_parent() -> Result<(), TreeCoreRuntimeError> { + let mut arena = arena(); + let root = arena.root(); + let a = arena.alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + )?; + assert!(arena.root_child(None, &[1]).is_some()); + arena.free_leaf(a)?; + // The freed leaf is gone from its parent's children. + assert!(arena.root_child(None, &[1]).is_none()); + assert_eq!(arena.len(), 1); + Ok(()) +} + +#[test] +fn children_are_keyed_by_the_first_radix_page() -> Result<(), TreeCoreRuntimeError> { + let mut arena: NodeArena> = NodeArena::new(vec![FULL], /* page_size = */ 1); + let root = arena.root(); + let a = arena.alloc_child( + root, + /* key = */ vec![1, 2, 3], + /* priority = */ 0, + /* extra_key = */ None, + )?; + assert_eq!(arena.root_child(None, &[1]).as_ref(), Some(&a)); + assert!(arena.root_child(None, &[1, 2, 3]).is_none()); + // The freed leaf unlinks through the same page key. + arena.free_leaf(a)?; + assert!(arena.root_child(None, &[1]).is_none()); + Ok(()) +} + +#[test] +fn siblings_sharing_a_first_page_collide() -> Result<(), TreeCoreRuntimeError> { + let mut arena: NodeArena> = NodeArena::new(vec![FULL], /* page_size = */ 1); + let root = arena.root(); + arena.alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 0, + /* extra_key = */ None, + )?; + // A radix tree admits one child per page: same first page is a duplicate. + assert!( + arena + .alloc_child( + root, + /* key = */ vec![1, 9], + /* priority = */ 0, + /* extra_key = */ None + ) + .is_err() + ); + Ok(()) +} + +#[test] +fn page_size_two_keys_children_by_two_atoms() -> Result<(), TreeCoreRuntimeError> { + let mut arena: NodeArena> = NodeArena::new(vec![FULL], /* page_size = */ 2); + let root = arena.root(); + let a = arena.alloc_child( + root, + /* key = */ vec![1, 2, 3, 4], + /* priority = */ 0, + /* extra_key = */ None, + )?; + assert_eq!(arena.root_child(None, &[1, 2]).as_ref(), Some(&a)); + Ok(()) +} + +#[test] +#[should_panic(expected = "get_and_batch_bump_access_counter: delta 0 must be positive")] +fn batch_bump_rejects_a_non_positive_delta() { + let mut arena: NodeArena> = NodeArena::new(vec![FULL], /* page_size = */ 1); + arena.get_and_batch_bump_access_counter(/* delta = */ 0); +} + +#[test] +fn batch_bump_reserves_a_tick_range() { + let mut arena: NodeArena> = NodeArena::new(vec![FULL], /* page_size = */ 1); + let single = arena.get_and_bump_access_counter(); + let newest = arena.get_and_batch_bump_access_counter(/* delta = */ 3); + assert_eq!(newest, single + 3); + assert_eq!(arena.get_and_bump_access_counter(), newest + 1); +} + +#[test] +fn alloc_child_rejects_duplicate_key() -> Result<(), TreeCoreRuntimeError> { + let mut arena = arena(); + let root = arena.root(); + arena.alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + )?; + assert!(matches!( + arena.alloc_child(root, /* key = */ vec![1], /* priority = */ 0, /* extra_key = */ None), + Err(TreeCoreRuntimeError::DuplicateChildKey { parent, .. }) if parent == arena.node(root).id + )); + // The rejected add reserved no slot: root + the first child only. + assert_eq!(arena.len(), 2); + Ok(()) +} + +#[test] +fn failed_alloc_child_mints_no_id_and_keeps_the_freelist() -> Result<(), TreeCoreRuntimeError> { + let mut arena = arena(); + let root = arena.root(); + arena.alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + )?; + let b = arena.alloc_child( + root, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + )?; + let last_minted = arena.node(b).id; + arena.free_leaf(b)?; + assert_eq!(arena.len(), 2); + assert!( + arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None + ) + .is_err() + ); + assert_eq!(arena.len(), 2); + // The failure leaked neither the peeked freelist slot nor a handle. + let c = arena.alloc_child( + root, + /* key = */ vec![3], + /* priority = */ 0, + /* extra_key = */ None, + )?; + assert_eq!(c, b); + assert_eq!(arena.node(c).id, last_minted + 1); + Ok(()) +} + +#[test] +#[should_panic(expected = "is not allocated")] +fn resolve_panics_on_a_never_minted_handle() { + let arena = arena(); + arena.root(); + arena.resolve(1_000_000); +} + +#[test] +fn alloc_child_under_freed_parent_returns_err() -> Result<(), TreeCoreRuntimeError> { + let mut arena = arena(); + let root = arena.root(); + let a = arena.alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + )?; + arena.free_leaf(a)?; + assert!(matches!( + arena.alloc_child(a, /* key = */ vec![2], /* priority = */ 0, /* extra_key = */ None), + Err(TreeCoreRuntimeError::ParentNotAllocated { id }) if id == a + )); + // The rejected alloc_child consumed no slot. + assert_eq!(arena.len(), 1); + Ok(()) +} + +#[test] +fn free_reuses_slots_lifo() -> Result<(), TreeCoreRuntimeError> { + let mut arena = arena(); + let root = arena.root(); + let a = arena.alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + )?; + let b = arena.alloc_child( + root, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + )?; + assert_eq!(arena.len(), 3); + arena.free_leaf(a)?; + arena.free_leaf(b)?; + assert_eq!(arena.len(), 1); + // Last freed is reused first. + assert_eq!( + arena.alloc_child( + root, + /* key = */ vec![3], + /* priority = */ 0, + /* extra_key = */ None + )?, + b + ); + assert_eq!( + arena.alloc_child( + root, + /* key = */ vec![4], + /* priority = */ 0, + /* extra_key = */ None + )?, + a + ); + assert_eq!(arena.len(), 3); + Ok(()) +} + +#[test] +fn arena_supports_bigram_key_type() -> Result<(), TreeCoreRuntimeError> { + let mut arena: NodeArena> = + NodeArena::new(vec![FULL], /* page_size = */ 1); + let root = arena.root(); + let a = arena.alloc_child( + root, + /* key = */ vec![(1, 2), (3, 4)], + /* priority = */ 0, + /* extra_key = */ None, + )?; + assert_eq!(arena.len(), 2); + assert_eq!(arena.node(a).key, vec![(1, 2), (3, 4)]); + Ok(()) +} + +#[test] +fn id_map_stays_consistent_across_free_and_realloc() -> Result<(), TreeCoreRuntimeError> { + let mut arena = arena(); + let root = arena.root(); + let a = arena.alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + )?; + let b = arena.alloc_child( + root, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + )?; + let b_id = arena.node(b).id; + arena.free_leaf(b)?; + assert!(arena.try_resolve(b_id).is_none()); + // The freed slot is recycled with a fresh handle; the old one stays dead. + let c = arena.alloc_child( + root, + /* key = */ vec![3], + /* priority = */ 0, + /* extra_key = */ None, + )?; + assert_eq!(c, b); + assert_ne!(arena.node(c).id, b_id); + assert!(arena.try_resolve(b_id).is_none()); + // Every live slot resolves back from its own handle. + for idx in arena.live_ids().collect::>() { + assert_eq!(arena.resolve(arena.node(idx).id), idx); + } + let _ = a; + Ok(()) +} + +// Eviction-eligible node set. + +#[test] +fn add_then_contains_and_len() { + let mut set = EvictableNodeSet::new(); + assert!(!set.contains(NodeIdx_(3))); + assert!(set.is_empty()); + set.add(NodeIdx_(3)); + assert!(set.contains(NodeIdx_(3))); + assert_eq!(set.len(), 1); + assert!(!set.is_empty()); +} + +#[test] +fn add_is_idempotent() { + let mut set = EvictableNodeSet::new(); + set.add(NodeIdx_(5)); + set.add(NodeIdx_(5)); + assert_eq!(set.len(), 1); + assert_eq!(set.iter().collect::>(), vec![NodeIdx_(5)]); +} + +#[test] +fn discard_removes_the_member() { + let mut set = EvictableNodeSet::new(); + set.add(NodeIdx_(2)); + set.discard(NodeIdx_(2)); + assert!(!set.contains(NodeIdx_(2))); + assert_eq!(set.len(), 0); + assert!(set.is_empty()); +} + +#[test] +fn add_sequential_ids_grows_the_slot_table_one_by_one() { + // Arena NodeIds are sequential, so add(node_id == slots.len()) is the common case. + let mut set = EvictableNodeSet::new(); + set.add(NodeIdx_(0)); + set.add(NodeIdx_(1)); + set.add(NodeIdx_(2)); + assert!(set.contains(NodeIdx_(0))); + assert!(set.contains(NodeIdx_(1))); + assert!(set.contains(NodeIdx_(2))); + assert_eq!(set.len(), 3); +} + +#[test] +fn discard_absent_is_noop() { + let mut set = EvictableNodeSet::new(); + set.add(NodeIdx_(1)); + set.discard(NodeIdx_(9)); + set.discard(NodeIdx_(0)); + assert!(set.contains(NodeIdx_(1))); + assert_eq!(set.len(), 1); +} + +#[test] +fn discard_fixes_up_the_swapped_member_slot() { + let mut set = EvictableNodeSet::new(); + set.add(NodeIdx_(10)); + set.add(NodeIdx_(20)); + set.add(NodeIdx_(30)); + // Removing the first member swap-moves the tail (30) into its slot. + set.discard(NodeIdx_(10)); + assert!(!set.contains(NodeIdx_(10))); + assert!(set.contains(NodeIdx_(20))); + assert!(set.contains(NodeIdx_(30))); + assert_eq!(set.len(), 2); + // The moved member's slot stays consistent for a follow-up discard. + set.discard(NodeIdx_(30)); + assert!(!set.contains(NodeIdx_(30))); + assert!(set.contains(NodeIdx_(20))); + assert_eq!(set.len(), 1); +} + +#[test] +fn discard_the_tail_member() { + let mut set = EvictableNodeSet::new(); + set.add(NodeIdx_(10)); + set.add(NodeIdx_(20)); + set.discard(NodeIdx_(20)); + assert!(set.contains(NodeIdx_(10))); + assert!(!set.contains(NodeIdx_(20))); + assert_eq!(set.len(), 1); +} + +#[test] +fn re_add_after_discard() { + let mut set = EvictableNodeSet::new(); + set.add(NodeIdx_(4)); + set.discard(NodeIdx_(4)); + set.add(NodeIdx_(4)); + assert!(set.contains(NodeIdx_(4))); + assert_eq!(set.len(), 1); +} + +#[test] +fn iter_yields_all_members() { + let mut set = EvictableNodeSet::new(); + set.add(NodeIdx_(10)); + set.add(NodeIdx_(20)); + set.add(NodeIdx_(30)); + set.discard(NodeIdx_(20)); + let mut members = set.iter().collect::>(); + members.sort_unstable(); + assert_eq!(members, vec![NodeIdx_(10), NodeIdx_(30)]); +} diff --git a/rust/mem-cache/src/tests/test_utils.rs b/rust/mem-cache/src/tests/test_utils.rs new file mode 100644 index 000000000..258e8bb7c --- /dev/null +++ b/rust/mem-cache/src/tests/test_utils.rs @@ -0,0 +1,46 @@ +//! Shared helpers for the crate's unit tests. + +use std::collections::HashMap; + +use tch::Tensor; + +use crate::components::ComponentType; +use crate::unified_tree_core::{CacheAction, EvictionStepResult}; + +/// Fold an eviction step into a caller's running accumulators (the Controller +/// consumption contract: deltas add, freed tensors append). +pub(crate) fn accumulate_step( + step: EvictionStepResult, + tracker: &mut HashMap, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, +) { + for (ct, delta) in step.tracker { + *tracker.entry(ct).or_insert(0) += delta; + } + for (ct, tensors) in step.device_frees { + device_frees.entry(ct).or_default().extend(tensors); + } + for (ct, tensors) in step.host_frees { + host_frees.entry(ct).or_default().extend(tensors); + } +} + +/// Short variant names for diagnosing an action sequence's shape. +pub(crate) fn action_kinds(actions: &[CacheAction]) -> Vec<&'static str> { + actions + .iter() + .map(|action| match action { + CacheAction::FreeDeviceKV(_) => "FreeDeviceKV", + CacheAction::FreeDeviceKVFullOnly(_) => "FreeDeviceKVFullOnly", + CacheAction::BackupKV(_) => "BackupKV", + CacheAction::ReplaceWriteThroughOnNodeSplit { .. } => "ReplaceWriteThroughOnNodeSplit", + CacheAction::MambaEvictExcessPathStates { .. } => "MambaEvictExcessPathStates", + CacheAction::FreeComponentDeviceSlot { .. } => "FreeComponentDeviceSlot", + CacheAction::FreeComponentHostSlot { .. } => "FreeComponentHostSlot", + CacheAction::RebuildFullToSwaMapping { .. } => "RebuildFullToSwaMapping", + CacheAction::RecoverSwaWithLockedFull { .. } => "RecoverSwaWithLockedFull", + CacheAction::SwaRebuild { .. } => "SwaRebuild", + }) + .collect() +} diff --git a/rust/mem-cache/src/tests/unified_lru_list.rs b/rust/mem-cache/src/tests/unified_lru_list.rs new file mode 100644 index 000000000..94a471fa4 --- /dev/null +++ b/rust/mem-cache/src/tests/unified_lru_list.rs @@ -0,0 +1,708 @@ +use super::*; +use crate::components::FULL; +use crate::node::{NodeArena, NodeIdx_, ValueSlotIdx}; + +fn order(list: &UnifiedLRUList) -> Vec { + list.iter().collect() +} + +#[test] +fn fresh_list_reads_are_empty() { + let list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + assert_eq!(list.get_lru_where(|_| true), None); + assert_eq!(list.iter().count(), 0); + assert_eq!(list.len(), 0); + list.validate(); +} + +#[test] +fn insert_mru_orders_most_recent_first() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(NodeIdx_(10)); + list.insert_mru(NodeIdx_(20)); + list.insert_mru(NodeIdx_(30)); + assert_eq!(order(&list), vec![NodeIdx_(30), NodeIdx_(20), NodeIdx_(10)]); + assert_eq!(list.len(), 3); + assert!(list.in_list(Some(NodeIdx_(10)))); + list.validate(); +} + +#[test] +#[should_panic(expected = "already in the LRU list")] +fn insert_mru_panics_when_already_a_member() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(NodeIdx_(1)); + list.insert_mru(NodeIdx_(1)); +} + +#[test] +fn remove_node_updates_membership_immediately() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(NodeIdx_(10)); + list.insert_mru(NodeIdx_(20)); + list.insert_mru(NodeIdx_(30)); + list.remove_node(NodeIdx_(20)); + assert_eq!(order(&list), vec![NodeIdx_(30), NodeIdx_(10)]); + assert!(!list.in_list(Some(NodeIdx_(20)))); + assert_eq!(list.len(), 2); + list.validate(); +} + +#[test] +#[should_panic(expected = "index out of bounds")] +fn remove_node_panics_when_absent() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.remove_node(NodeIdx_(1)); +} + +#[test] +#[should_panic(expected = "not in the LRU list")] +fn remove_node_panics_on_a_node_removed_earlier() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(NodeIdx_(10)); + list.insert_mru(NodeIdx_(20)); + list.remove_node(NodeIdx_(10)); + // The cell is still allocated but reset; membership must gate the removal. + list.remove_node(NodeIdx_(10)); +} + +#[test] +#[should_panic(expected = "not in the LRU list")] +fn remove_node_panics_on_an_unlisted_cell() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(NodeIdx_(10)); + list.remove_node(NodeIdx_(10)); + list.remove_node_(UnifiedLRUList::cell_of_(NodeIdx_(10))); +} + +#[test] +#[should_panic(expected = "already in the LRU list")] +fn add_node_panics_on_a_linked_cell() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(NodeIdx_(10)); + list.add_node_(UnifiedLRUList::cell_of_(NodeIdx_(10))); +} + +#[test] +fn removed_nodes_can_be_reinserted() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(NodeIdx_(10)); + list.insert_mru(NodeIdx_(20)); + list.remove_node(NodeIdx_(10)); + list.insert_mru(NodeIdx_(10)); + assert_eq!(order(&list), vec![NodeIdx_(10), NodeIdx_(20)]); + list.validate(); +} + +#[test] +fn reset_node_mru_moves_a_member_to_the_front() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(NodeIdx_(10)); + list.insert_mru(NodeIdx_(20)); + list.insert_mru(NodeIdx_(30)); + list.reset_node_mru(NodeIdx_(10)); + assert_eq!(order(&list), vec![NodeIdx_(10), NodeIdx_(30), NodeIdx_(20)]); + list.validate(); +} + +#[test] +#[should_panic(expected = "index out of bounds")] +fn reset_node_mru_panics_on_a_non_member() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.reset_node_mru(NodeIdx_(7)); +} + +// Arena chain root -> a -> b -> c plus a sibling, two atoms per key. +fn arena_chain() -> ( + NodeArena>, + NodeIdx_, + NodeIdx_, + NodeIdx_, + NodeIdx_, + NodeIdx_, +) { + let mut arena = NodeArena::new(vec![crate::components::FULL], /* page_size = */ 1); + let root = arena.root(); + let a = arena + .alloc_child( + root, + /* key = */ vec![1, 11], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let b = arena + .alloc_child( + a, + /* key = */ vec![2, 22], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let c = arena + .alloc_child( + b, + /* key = */ vec![3, 33], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let other = arena + .alloc_child( + root, + /* key = */ vec![9, 99], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + (arena, root, a, b, c, other) +} + +#[test] +fn reset_parents_mru_reranks_included_nodes_deepest_first() { + let (arena, _root, a, b, c, other) = arena_chain(); + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(a); + list.insert_mru(c); + list.insert_mru(other); + // b is excluded and skipped; c then a become the MRU run. + list.reset_node_and_parents_mru(c, &arena, |node| node.idx != b); + assert_eq!(order(&list), vec![c, a, other]); + list.validate(); +} + +#[test] +fn reset_parents_mru_reranks_ancestors_when_the_deepest_is_excluded() { + let (arena, _root, a, _b, c, other) = arena_chain(); + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(a); + list.insert_mru(other); + // c and b are excluded; a alone becomes the new MRU head. + list.reset_node_and_parents_mru(c, &arena, |node| node.idx == a); + assert_eq!(order(&list), vec![a, other]); + list.validate(); +} + +#[test] +fn reset_walks_are_noops_when_node_is_the_root() { + let (arena, root, a, _b, _c, _other) = arena_chain(); + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(a); + list.reset_node_and_parents_mru(root, &arena, |_| true); + list.reset_node_and_window_ancestors_mru(root, 4, &arena, |_| true); + assert_eq!(order(&list), vec![a]); + list.validate(); +} + +#[test] +#[should_panic(expected = "not in the LRU list")] +fn reset_parents_mru_panics_on_an_unlisted_included_node() { + let (arena, _root, _a, _b, c, _other) = arena_chain(); + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(c); + list.remove_node(c); + list.reset_node_and_parents_mru(c, &arena, |_| true); +} + +#[test] +fn reset_window_ancestors_mru_stops_at_the_window() { + let (arena, _root, a, b, c, other) = arena_chain(); + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(b); + list.insert_mru(c); + list.insert_mru(a); + list.insert_mru(other); + // A window of 4 atoms covers c and b; a stays put beyond it. + list.reset_node_and_window_ancestors_mru(c, 4, &arena, |_| true); + assert_eq!(order(&list), vec![c, b, other, a]); + list.validate(); +} + +#[test] +fn reset_window_ancestors_mru_includes_the_straddling_ancestor() { + let (arena, _root, a, b, c, other) = arena_chain(); + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(a); + list.insert_mru(b); + list.insert_mru(c); + list.insert_mru(other); + // A window of 5 atoms ends mid-a: the straddling ancestor is still included. + list.reset_node_and_window_ancestors_mru(c, 5, &arena, |_| true); + assert_eq!(order(&list), vec![c, b, a, other]); + list.validate(); +} + +#[test] +fn reset_walks_stop_at_the_salted_chains_root() { + let mut arena: NodeArena> = + NodeArena::new(vec![crate::components::FULL], /* page_size = */ 1); + let named = arena.root(); + let a = arena + .alloc_child( + named, + /* key = */ vec![1, 11], + /* priority = */ 0, + Some("lora-1"), + ) + .unwrap(); + let b = arena + .alloc_child( + a, + /* key = */ vec![2, 22], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(a); + list.insert_mru(b); + // Both walks terminate at the root without visiting it. + list.reset_node_and_parents_mru(b, &arena, |_| true); + assert_eq!(order(&list), vec![b, a]); + list.reset_node_and_window_ancestors_mru(b, 100, &arena, |_| true); + assert_eq!(order(&list), vec![b, a]); + list.validate(); +} + +#[test] +fn get_lru_no_lock_returns_the_lru_most_unlocked_member() { + let (mut arena, _root, a, b, c, _other) = arena_chain(); + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(a); + list.insert_mru(b); + list.insert_mru(c); + assert_eq!(list.get_lru_no_lock(&arena), Some(a)); + // A lock on the list's own slot hides the LRU end from the walker. + arena + .node_mut(a) + .set_lock_ref_(ValueSlotIdx::device(FULL), 1); + assert_eq!(list.get_lru_no_lock(&arena), Some(b)); + arena + .node_mut(b) + .set_lock_ref_(ValueSlotIdx::device(FULL), 1); + arena + .node_mut(c) + .set_lock_ref_(ValueSlotIdx::device(FULL), 1); + assert_eq!(list.get_lru_no_lock(&arena), None); +} + +#[test] +fn get_prev_no_lock_skips_locked_members_toward_the_mru_end() { + let (mut arena, _root, a, b, c, other) = arena_chain(); + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(a); + list.insert_mru(b); + list.insert_mru(c); + list.insert_mru(other); + assert_eq!(list.get_prev_no_lock(a, &arena), Some(b)); + // The locked b is skipped; from the MRU end there is no predecessor left. + arena + .node_mut(b) + .set_lock_ref_(ValueSlotIdx::device(FULL), 1); + assert_eq!(list.get_prev_no_lock(a, &arena), Some(c)); + assert_eq!(list.get_prev_no_lock(other, &arena), None); + // A lock on a different slot does not gate this list's walker. + arena.node_mut(c).set_lock_ref_(ValueSlotIdx::host(FULL), 1); + assert_eq!(list.get_prev_no_lock(a, &arena), Some(c)); +} + +#[test] +fn reset_window_accumulation_counts_excluded_nodes() { + let (arena, _root, a, b, c, other) = arena_chain(); + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(c); + list.insert_mru(a); + list.insert_mru(other); + // b is excluded but its atoms still consume the window, keeping a out of reach. + list.reset_node_and_window_ancestors_mru(c, 4, &arena, |node| node.idx != b); + assert_eq!(order(&list), vec![c, other, a]); + list.validate(); +} + +#[test] +#[should_panic(expected = "not in the LRU list")] +fn reset_node_mru_panics_on_a_node_removed_earlier() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(NodeIdx_(10)); + list.remove_node(NodeIdx_(10)); + // The cell is still allocated but unlisted; the gated read must reject it. + list.reset_node_mru(NodeIdx_(10)); +} + +#[test] +fn in_list_is_false_for_none_and_non_members() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + assert!(!list.in_list(None)); + assert!(!list.in_list(Some(NodeIdx_(5)))); + list.insert_mru(NodeIdx_(5)); + assert!(list.in_list(Some(NodeIdx_(5)))); +} + +#[test] +fn get_lru_where_walks_from_the_tail() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(NodeIdx_(10)); + list.insert_mru(NodeIdx_(20)); + list.insert_mru(NodeIdx_(30)); + assert_eq!(list.get_lru_where(|_| true), Some(NodeIdx_(10))); + assert_eq!( + list.get_lru_where(|id| id != NodeIdx_(10)), + Some(NodeIdx_(20)) + ); + assert_eq!(list.get_lru_where(|_| false), None); +} + +#[test] +fn get_prev_where_walks_toward_the_head_from_a_member() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(NodeIdx_(10)); + list.insert_mru(NodeIdx_(20)); + list.insert_mru(NodeIdx_(30)); + // Order is [30, 20, 10]; 10's predecessors are 20 then 30. + assert_eq!( + list.get_prev_where(NodeIdx_(10), |_| true), + Some(NodeIdx_(20)) + ); + assert_eq!( + list.get_prev_where(NodeIdx_(10), |id| id != NodeIdx_(20)), + Some(NodeIdx_(30)) + ); + assert_eq!(list.get_prev_where(NodeIdx_(30), |_| true), None); +} + +#[test] +#[should_panic(expected = "index out of bounds")] +fn get_prev_where_panics_on_a_non_member() { + let list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.get_prev_where(NodeIdx_(7), |_| true); +} + +#[test] +#[should_panic(expected = "not in the LRU list")] +fn get_prev_where_panics_on_a_node_removed_earlier() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(NodeIdx_(10)); + list.insert_mru(NodeIdx_(20)); + list.remove_node(NodeIdx_(10)); + // The cell is still allocated but unlisted; the gated read must reject it. + list.get_prev_where(NodeIdx_(10), |_| true); +} + +#[test] +fn get_prev_before_remove_keeps_the_walk_consistent() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(NodeIdx_(10)); + list.insert_mru(NodeIdx_(20)); + list.insert_mru(NodeIdx_(30)); + // The eviction-cursor contract: compute the predecessor, then remove. + let next = list.get_prev_where(NodeIdx_(10), |_| true); + list.remove_node(NodeIdx_(10)); + assert_eq!(next, Some(NodeIdx_(20))); + assert!(list.in_list(next)); + assert_eq!( + list.get_prev_where(NodeIdx_(20), |_| true), + Some(NodeIdx_(30)) + ); + list.validate(); +} + +#[test] +fn insert_mru_grows_the_cell_table_one_by_one() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(NodeIdx_(0)); + list.insert_mru(NodeIdx_(1)); + list.insert_mru(NodeIdx_(2)); + assert_eq!(order(&list), vec![NodeIdx_(2), NodeIdx_(1), NodeIdx_(0)]); + list.validate(); +} + +#[test] +#[should_panic(expected = "broken prev link")] +fn validate_panics_on_a_corrupted_prev_link() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(NodeIdx_(10)); + list.insert_mru(NodeIdx_(20)); + let cell = UnifiedLRUList::cell_of_(NodeIdx_(10)); + list.cells[cell.0].prev = cell; + list.validate(); +} + +#[test] +#[should_panic(expected = "membership mismatch")] +fn validate_panics_on_a_linked_cell_without_the_flag() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(NodeIdx_(10)); + list.cells[UnifiedLRUList::cell_of_(NodeIdx_(10)).0].in_list = false; + list.validate(); +} + +#[test] +#[should_panic(expected = "membership mismatch")] +fn validate_panics_on_a_flagged_unlinked_cell() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(NodeIdx_(10)); + list.insert_mru(NodeIdx_(20)); + list.remove_node(NodeIdx_(20)); + // The cell is reset but a stray flag claims membership. + list.cells[UnifiedLRUList::cell_of_(NodeIdx_(20)).0].in_list = true; + list.validate(); +} + +#[test] +fn reset_window_ancestors_mru_is_a_noop_on_a_zero_window() { + let (arena, _root, a, _b, c, _other) = arena_chain(); + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(a); + list.insert_mru(c); + list.reset_node_and_window_ancestors_mru(c, 0, &arena, |_| true); + assert_eq!(order(&list), vec![c, a]); + list.validate(); +} + +#[test] +#[should_panic(expected = "out of bounds")] +fn validate_panics_on_an_out_of_range_link() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(NodeIdx_(10)); + list.cells[UnifiedLRUList::cell_of_(NodeIdx_(10)).0].next = CellId(99); + list.validate(); +} + +#[test] +#[should_panic(expected = "length mismatch")] +fn validate_panics_on_a_desynced_member_counter() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(NodeIdx_(10)); + list.len = 2; + list.validate(); +} + +#[test] +fn len_drops_to_zero_after_all_members_removed() { + let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + list.insert_mru(NodeIdx_(1)); + list.remove_node(NodeIdx_(1)); + assert_eq!(list.len(), 0); + assert_eq!(list.iter().count(), 0); + list.validate(); +} + +#[test] +fn check_linked_list_accepts_a_clean_list() { + let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + lru.insert_mru(NodeIdx_(0)); + lru.insert_mru(NodeIdx_(1)); + lru.insert_mru(NodeIdx_(2)); + let mut errors = Vec::new(); + lru.check_linked_list_("[t]", &mut errors); + assert!(errors.is_empty()); +} + +#[test] +fn check_linked_list_reports_a_broken_prev() { + let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + lru.insert_mru(NodeIdx_(0)); + lru.insert_mru(NodeIdx_(1)); + lru.cells[UnifiedLRUList::cell_of_(NodeIdx_(0)).0].prev = UnifiedLRUList::cell_of_(NodeIdx_(0)); + let mut errors = Vec::new(); + lru.check_linked_list_("[t]", &mut errors); + assert_eq!(errors.len(), 1); + assert!(errors[0].contains("broken prev at node 0")); +} + +#[test] +fn check_linked_list_reports_an_unflagged_member() { + let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + lru.insert_mru(NodeIdx_(0)); + lru.insert_mru(NodeIdx_(1)); + lru.cells[UnifiedLRUList::cell_of_(NodeIdx_(0)).0].in_list = false; + let mut errors = Vec::new(); + lru.check_linked_list_("[t]", &mut errors); + assert_eq!(errors.len(), 1); + assert!(errors[0].contains("node 0 in list not flagged")); +} + +#[test] +fn check_linked_list_reports_a_cycle() { + let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + lru.insert_mru(NodeIdx_(0)); + lru.insert_mru(NodeIdx_(1)); + // 0's next loops back to 1 instead of reaching the tail. + lru.cells[UnifiedLRUList::cell_of_(NodeIdx_(0)).0].next = UnifiedLRUList::cell_of_(NodeIdx_(1)); + let mut errors = Vec::new(); + lru.check_linked_list_("[t]", &mut errors); + assert!(errors.iter().any(|e| e.contains("cycle at node 1"))); +} + +#[test] +fn check_linked_list_reports_a_count_mismatch() { + let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + lru.insert_mru(NodeIdx_(0)); + lru.len = 2; + let mut errors = Vec::new(); + lru.check_linked_list_("[t]", &mut errors); + assert_eq!(errors.len(), 1); + assert!(errors[0].contains("list=1 != len=2")); +} + +#[test] +fn check_linked_list_reports_an_out_of_bounds_link() { + let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + lru.insert_mru(NodeIdx_(0)); + lru.cells[UnifiedLRUList::cell_of_(NodeIdx_(0)).0].next = CellId(999); + let mut errors = Vec::new(); + lru.check_linked_list_("[t]", &mut errors); + assert!(errors.iter().any(|e| e.contains("cell 999 out of bounds"))); +} + +#[test] +fn check_linked_list_reports_a_broken_tail_backlink() { + let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + lru.insert_mru(NodeIdx_(0)); + lru.cells[TAIL.0].prev = HEAD; + let mut errors = Vec::new(); + lru.check_linked_list_("[t]", &mut errors); + assert_eq!(errors.len(), 1); + assert!(errors[0].contains("broken tail backlink")); +} + +#[test] +fn check_linked_list_reports_a_flagged_unreachable_cell() { + let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL)); + lru.insert_mru(NodeIdx_(0)); + lru.insert_mru(NodeIdx_(1)); + lru.remove_node(NodeIdx_(0)); + // Re-flag the unlinked cell without relinking it. + lru.cells[UnifiedLRUList::cell_of_(NodeIdx_(0)).0].in_list = true; + let mut errors = Vec::new(); + lru.check_linked_list_("[t]", &mut errors); + assert!( + errors + .iter() + .any(|e| e.contains("node 0 flagged but unreachable")) + ); +} + +// Eviction priority keys. + +// A node with distinct field values: last_access 5, creation 7, hits 3, priority 9. +fn arena_with_node() -> (NodeArena>, NodeIdx_) { + let mut arena: NodeArena> = NodeArena::new(vec![FULL], /* page_size = */ 1); + let root = arena.root(); + let a = arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 9, + /* extra_key = */ None, + ) + .unwrap(); + let node = arena.node_mut(a); + node.last_access_counter = 5; + node.creation_counter = 7; + node.hit_count = 3; + (arena, NodeIdx_(a.0)) +} + +#[test] +fn each_strategy_maps_its_node_fields_into_the_key() { + let (arena, a) = arena_with_node(); + let node = arena.node(NodeIdx_(a.0)); + assert_eq!(LruStrategy.get_priority(node), PriorityKey(5, 0)); + assert_eq!(LfuStrategy.get_priority(node), PriorityKey(3, 5)); + assert_eq!(FifoStrategy.get_priority(node), PriorityKey(7, 0)); + assert_eq!(MruStrategy.get_priority(node), PriorityKey(-5, 0)); + assert_eq!(FiloStrategy.get_priority(node), PriorityKey(-7, 0)); + assert_eq!(PriorityStrategy.get_priority(node), PriorityKey(9, 5)); +} + +#[test] +fn slru_segments_on_the_protected_threshold() { + let (mut arena, a) = arena_with_node(); + let slru = SlruStrategy { + protected_threshold: 2, + }; + // 3 hits >= threshold 2: protected segment. + assert_eq!( + slru.get_priority(arena.node(NodeIdx_(a.0))), + PriorityKey(1, 5) + ); + // Exactly at the threshold counts as protected. + arena.node_mut(NodeIdx_(a.0)).hit_count = 2; + assert_eq!( + slru.get_priority(arena.node(NodeIdx_(a.0))), + PriorityKey(1, 5) + ); + arena.node_mut(NodeIdx_(a.0)).hit_count = 1; + assert_eq!( + slru.get_priority(arena.node(NodeIdx_(a.0))), + PriorityKey(0, 5) + ); +} + +#[test] +fn get_eviction_strategy_resolves_each_policy_name() { + let (arena, a) = arena_with_node(); + let node = arena.node(NodeIdx_(a.0)); + // Distinct node fields make each policy's key identify its strategy. + let cases = [ + ("lru", PriorityKey(5, 0)), + ("LFU", PriorityKey(3, 5)), + ("fifo", PriorityKey(7, 0)), + ("mru", PriorityKey(-5, 0)), + ("filo", PriorityKey(-7, 0)), + ("priority", PriorityKey(9, 5)), + ("slru", PriorityKey(1, 5)), + ]; + for (policy, expected) in cases { + assert_eq!( + get_eviction_strategy::>(policy).get_priority(node), + expected, + "policy {policy}" + ); + } +} + +#[test] +fn eviction_policy_names_are_case_insensitive() { + let (arena, a) = arena_with_node(); + let node = arena.node(NodeIdx_(a.0)); + // Mixed-case names resolve to the same strategies as their lowercase forms. + assert_eq!( + get_eviction_strategy::>("LRU").get_priority(node), + PriorityKey(5, 0) + ); + assert_eq!( + get_eviction_strategy::>("Priority").get_priority(node), + PriorityKey(9, 5) + ); +} + +#[test] +fn get_eviction_strategy_slru_default_threshold_is_two() { + let (mut arena, a) = arena_with_node(); + let slru = get_eviction_strategy::>("slru"); + // Exactly 2 hits is protected under the factory default; 1 is not. + arena.node_mut(NodeIdx_(a.0)).hit_count = 2; + assert_eq!( + slru.get_priority(arena.node(NodeIdx_(a.0))), + PriorityKey(1, 5) + ); + arena.node_mut(NodeIdx_(a.0)).hit_count = 1; + assert_eq!( + slru.get_priority(arena.node(NodeIdx_(a.0))), + PriorityKey(0, 5) + ); +} + +#[test] +#[should_panic(expected = "Unknown eviction policy: random. Supported policies:")] +fn get_eviction_strategy_panics_on_an_unknown_policy() { + get_eviction_strategy::>("Random"); +} + +#[test] +fn priority_keys_order_lexicographically() { + assert!(PriorityKey(0, 9) < PriorityKey(1, 0)); + assert!(PriorityKey(1, 2) < PriorityKey(1, 3)); +} diff --git a/rust/mem-cache/src/tests/unified_tree_core.rs b/rust/mem-cache/src/tests/unified_tree_core.rs new file mode 100644 index 000000000..61fec2de9 --- /dev/null +++ b/rust/mem-cache/src/tests/unified_tree_core.rs @@ -0,0 +1,7739 @@ +use std::sync::Mutex; + +use tch::Tensor; + +use super::*; +use crate::components::{FULL, MAMBA, SWA}; +use crate::node::ValueSlotIdx; +use crate::test_utils::{accumulate_step, action_kinds}; + +fn core() -> UnifiedTreeCore> { + UnifiedTreeCore::new(CacheInitParams::default(), vec![FULL]) +} + +// Records every refresh_lru dispatch; unrelated hooks stay unimplemented. +#[derive(Default)] +struct RecordingComponentForTest { + refreshes: Mutex>, + host_eviction_calls: Mutex>, +} + +impl TreeComponent> for RecordingComponentForTest { + fn component_type(&self) -> ComponentType { + SWA + } + + fn refresh_lru( + &self, + _tree_core: &mut UnifiedTreeCore>, + phase: LRURefreshPhase, + node_id: NodeIdx_, + ) { + self.refreshes.lock().unwrap().push((phase, node_id)); + } + + fn create_match_validator( + &self, + _tree_core: &UnifiedTreeCore>, + _match_device_only: bool, + ) -> Box>, NodeIdx_) -> bool> { + Box::new(|_, _| true) + } + + fn redistribute_on_node_split( + &self, + _tree_core: &mut UnifiedTreeCore>, + _new_parent_id: NodeIdx_, + _child_id: NodeIdx_, + ) { + } + + fn evict_component( + &self, + _tree_core: &mut UnifiedTreeCore>, + _node_id: NodeIdx_, + _device_frees: &mut HashMap>, + _host_frees: &mut HashMap>, + _target: EvictLayer, + ) -> (usize, usize) { + unimplemented!() + } + + fn evict_device_start(&self, _tree_core: &mut UnifiedTreeCore>, _request_cnt: usize) { + unimplemented!() + } + + fn evict_device_next_node( + &self, + _tree_core: &mut UnifiedTreeCore>, + _tracker: &mut HashMap, + _device_frees: &mut HashMap>, + _host_frees: &mut HashMap>, + ) -> Option { + unimplemented!() + } + + fn evict_device_end(&self, _tree_core: &mut UnifiedTreeCore>) { + unimplemented!() + } + + fn acquire_component_lock( + &self, + _tree_core: &mut UnifiedTreeCore>, + _node_id: NodeIdx_, + _result: IncLockRefResult, + _lock_host: bool, + ) -> IncLockRefResult { + unimplemented!() + } + + fn release_component_lock( + &self, + _tree_core: &mut UnifiedTreeCore>, + _node_id: NodeIdx_, + _params: Option<&DecLockRefParams>, + _lock_host: bool, + ) { + unimplemented!() + } + + fn reclaim_coexisting_host_values( + &self, + _tree_core: &mut UnifiedTreeCore>, + _num_tokens: usize, + tracker: &mut HashMap, + _device_frees: &mut HashMap>, + _host_frees: &mut HashMap>, + ) { + let tracked = tracker[&SWA]; + self.host_eviction_calls + .lock() + .unwrap() + .push(("reclaim", tracked)); + tracker.insert(SWA, tracked + 2); + } + + fn drive_host_eviction( + &self, + _tree_core: &mut UnifiedTreeCore>, + _num_tokens: usize, + tracker: &mut HashMap, + _device_frees: &mut HashMap>, + _host_frees: &mut HashMap>, + ) { + let tracked = tracker[&SWA]; + self.host_eviction_calls + .lock() + .unwrap() + .push(("drive", tracked)); + tracker.insert(SWA, tracked + 3); + } +} + +// Counts every match-validator invocation; unrelated hooks stay unimplemented. +#[derive(Default)] +struct CountingComponentForTest { + validator_calls: Arc>, +} + +impl TreeComponent> for CountingComponentForTest { + fn component_type(&self) -> ComponentType { + SWA + } + + fn refresh_lru( + &self, + _tree_core: &mut UnifiedTreeCore>, + _phase: LRURefreshPhase, + _node_id: NodeIdx_, + ) { + } + + fn create_match_validator( + &self, + _tree_core: &UnifiedTreeCore>, + _match_device_only: bool, + ) -> Box>, NodeIdx_) -> bool> { + let calls = Arc::clone(&self.validator_calls); + Box::new(move |_, _| { + *calls.lock().unwrap() += 1; + true + }) + } + + fn redistribute_on_node_split( + &self, + _tree_core: &mut UnifiedTreeCore>, + _new_parent_id: NodeIdx_, + _child_id: NodeIdx_, + ) { + unimplemented!() + } + + fn evict_component( + &self, + _tree_core: &mut UnifiedTreeCore>, + _node_id: NodeIdx_, + _device_frees: &mut HashMap>, + _host_frees: &mut HashMap>, + _target: EvictLayer, + ) -> (usize, usize) { + unimplemented!() + } + + fn evict_device_start(&self, _tree_core: &mut UnifiedTreeCore>, _request_cnt: usize) { + unimplemented!() + } + + fn evict_device_next_node( + &self, + _tree_core: &mut UnifiedTreeCore>, + _tracker: &mut HashMap, + _device_frees: &mut HashMap>, + _host_frees: &mut HashMap>, + ) -> Option { + unimplemented!() + } + + fn evict_device_end(&self, _tree_core: &mut UnifiedTreeCore>) { + unimplemented!() + } + + fn acquire_component_lock( + &self, + _tree_core: &mut UnifiedTreeCore>, + _node_id: NodeIdx_, + _result: IncLockRefResult, + _lock_host: bool, + ) -> IncLockRefResult { + unimplemented!() + } + + fn release_component_lock( + &self, + _tree_core: &mut UnifiedTreeCore>, + _node_id: NodeIdx_, + _params: Option<&DecLockRefParams>, + _lock_host: bool, + ) { + unimplemented!() + } +} + +// Swa-flavored stub driver: any dispatched call panics as unimplemented. +// A Mamba-slot double with internal priority 0; its release panics so a +// test can pin that dec_swa_lock_only dispatches lower-priority releases. +struct LowPriorityComponentForTest; + +impl TreeComponent> for LowPriorityComponentForTest { + fn component_type(&self) -> ComponentType { + MAMBA + } + + fn eviction_priority(&self, _is_leaf: bool) -> i64 { + 0 + } + + fn create_match_validator( + &self, + _tree_core: &UnifiedTreeCore>, + _match_device_only: bool, + ) -> Box>, NodeIdx_) -> bool> { + unimplemented!() + } + + fn redistribute_on_node_split( + &self, + _tree_core: &mut UnifiedTreeCore>, + _new_parent_id: NodeIdx_, + _child_id: NodeIdx_, + ) { + } + + fn evict_component( + &self, + _tree_core: &mut UnifiedTreeCore>, + _node_id: NodeIdx_, + _device_frees: &mut HashMap>, + _host_frees: &mut HashMap>, + _target: EvictLayer, + ) -> (usize, usize) { + unimplemented!() + } + + fn evict_device_start(&self, _tree_core: &mut UnifiedTreeCore>, _request_cnt: usize) { + unimplemented!() + } + + fn evict_device_next_node( + &self, + _tree_core: &mut UnifiedTreeCore>, + _tracker: &mut HashMap, + _device_frees: &mut HashMap>, + _host_frees: &mut HashMap>, + ) -> Option { + unimplemented!() + } + + fn evict_device_end(&self, _tree_core: &mut UnifiedTreeCore>) { + unimplemented!() + } + + fn acquire_component_lock( + &self, + _tree_core: &mut UnifiedTreeCore>, + _node_id: NodeIdx_, + _result: IncLockRefResult, + _lock_host: bool, + ) -> IncLockRefResult { + unimplemented!() + } + + fn release_component_lock( + &self, + _tree_core: &mut UnifiedTreeCore>, + _node_id: NodeIdx_, + params: Option<&DecLockRefParams>, + lock_host: bool, + ) { + assert!(!lock_host); + assert!(params.is_some_and(|p| p.swa_uuid_for_lock.is_some())); + panic!("low-priority release dispatched"); + } +} + +struct SwaComponentForTest; + +impl TreeComponent> for SwaComponentForTest { + fn component_type(&self) -> ComponentType { + SWA + } + + fn create_match_validator( + &self, + _tree_core: &UnifiedTreeCore>, + _match_device_only: bool, + ) -> Box>, NodeIdx_) -> bool> { + unimplemented!() + } + + fn redistribute_on_node_split( + &self, + _tree_core: &mut UnifiedTreeCore>, + _new_parent_id: NodeIdx_, + _child_id: NodeIdx_, + ) { + } + + fn evict_component( + &self, + _tree_core: &mut UnifiedTreeCore>, + _node_id: NodeIdx_, + _device_frees: &mut HashMap>, + _host_frees: &mut HashMap>, + _target: EvictLayer, + ) -> (usize, usize) { + unimplemented!() + } + + fn evict_device_start(&self, _tree_core: &mut UnifiedTreeCore>, _request_cnt: usize) { + unimplemented!() + } + + fn evict_device_next_node( + &self, + _tree_core: &mut UnifiedTreeCore>, + _tracker: &mut HashMap, + _device_frees: &mut HashMap>, + _host_frees: &mut HashMap>, + ) -> Option { + unimplemented!() + } + + fn evict_device_end(&self, _tree_core: &mut UnifiedTreeCore>) { + unimplemented!() + } + + fn acquire_component_lock( + &self, + _tree_core: &mut UnifiedTreeCore>, + _node_id: NodeIdx_, + _result: IncLockRefResult, + _lock_host: bool, + ) -> IncLockRefResult { + unimplemented!() + } + + fn release_component_lock( + &self, + _tree_core: &mut UnifiedTreeCore>, + _node_id: NodeIdx_, + _params: Option<&DecLockRefParams>, + _lock_host: bool, + ) { + unimplemented!() + } +} + +// Swa-flavored driver with working eviction hooks: evict_component frees the +// Swa slot and records the call; eviction priorities are configurable. +struct SwaEvictionComponentForTest { + leaf_priority: i64, + internal_priority: i64, + evictions: Mutex>, +} + +impl SwaEvictionComponentForTest { + fn new(leaf_priority: i64, internal_priority: i64) -> Self { + SwaEvictionComponentForTest { + leaf_priority, + internal_priority, + evictions: Mutex::new(Vec::new()), + } + } +} + +impl TreeComponent> for SwaEvictionComponentForTest { + fn component_type(&self) -> ComponentType { + SWA + } + + fn eviction_priority(&self, is_leaf: bool) -> i64 { + if is_leaf { + self.leaf_priority + } else { + self.internal_priority + } + } + + fn evict_component( + &self, + tree_core: &mut UnifiedTreeCore>, + node_id: NodeIdx_, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + target: EvictLayer, + ) -> (usize, usize) { + self.evictions.lock().unwrap().push((node_id, target)); + let node = tree_core.arena.node_mut(node_id); + let mut device_freed = 0; + let mut host_freed = 0; + if target.contains(EvictLayer::Device) + && let Some(value) = node.values[SWA.idx()].value.take() + { + device_freed = value.size()[0] as usize; + device_frees.entry(SWA).or_default().push(value); + } + if target.contains(EvictLayer::Host) + && let Some(value) = node.state_mut_(ValueSlotIdx::host(SWA)).value.take() + { + host_freed = value.size()[0] as usize; + host_frees.entry(SWA).or_default().push(value); + } + (device_freed, host_freed) + } + + fn create_match_validator( + &self, + _tree_core: &UnifiedTreeCore>, + _match_device_only: bool, + ) -> Box>, NodeIdx_) -> bool> { + unimplemented!() + } + + fn redistribute_on_node_split( + &self, + _tree_core: &mut UnifiedTreeCore>, + _new_parent_id: NodeIdx_, + _child_id: NodeIdx_, + ) { + } + + fn evict_device_start(&self, _tree_core: &mut UnifiedTreeCore>, _request_cnt: usize) { + unimplemented!() + } + + fn evict_device_next_node( + &self, + _tree_core: &mut UnifiedTreeCore>, + _tracker: &mut HashMap, + _device_frees: &mut HashMap>, + _host_frees: &mut HashMap>, + ) -> Option { + unimplemented!() + } + + fn evict_device_end(&self, _tree_core: &mut UnifiedTreeCore>) { + unimplemented!() + } + + fn acquire_component_lock( + &self, + _tree_core: &mut UnifiedTreeCore>, + _node_id: NodeIdx_, + _result: IncLockRefResult, + _lock_host: bool, + ) -> IncLockRefResult { + unimplemented!() + } + + fn release_component_lock( + &self, + _tree_core: &mut UnifiedTreeCore>, + _node_id: NodeIdx_, + _params: Option<&DecLockRefParams>, + _lock_host: bool, + ) { + unimplemented!() + } +} + +// A locked non-root anchor; component dispatch bypasses roots entirely. +fn locked_anchor_for_dispatch(tc: &mut UnifiedTreeCore>) -> NodeIdx_ { + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 11], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1])); + tc.component_state_mut(FULL).evictable_size = 2; + tc.inc_lock_ref(tc.arena.node(n1).id); + n1 +} + +#[test] +fn dec_lock_ref_skip_swa_skips_the_swa_component() { + let mut tc = core(); + let n1 = locked_anchor_for_dispatch(&mut tc); + tc.register_component_(Arc::new(SwaComponentForTest)); + // The skipped Swa driver is never dispatched, so its stub cannot panic. + tc.dec_lock_ref( + tc.arena.node(n1).id, + /* params = */ None, + /* skip_swa = */ true, + ); + assert_eq!(tc.arena.device_lock_ref(n1, FULL), 0); +} + +#[test] +#[should_panic(expected = "not implemented")] +fn inc_lock_ref_reaches_every_component() { + let mut tc = core(); + tc.register_component_(Arc::new(SwaComponentForTest)); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 11], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1])); + tc.component_state_mut(FULL).evictable_size = 2; + tc.inc_lock_ref(tc.arena.node(n1).id); +} + +#[test] +#[should_panic(expected = "not implemented")] +fn dec_lock_ref_without_skip_swa_reaches_every_component() { + let mut tc = core(); + let n1 = locked_anchor_for_dispatch(&mut tc); + tc.register_component_(Arc::new(SwaComponentForTest)); + tc.dec_lock_ref( + tc.arena.node(n1).id, + /* params = */ None, + /* skip_swa = */ false, + ); +} + +#[test] +fn set_component_device_value_sizes_by_the_value_length() { + let mut tc = core(); + tc.register_component_(Arc::new(SwaComponentForTest)); + let root = tc.arena.root(); + let node = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2, 3], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.set_component_device_value( + tc.arena.node(node).id, + SWA, + Tensor::from_slice(&[7i64, 8, 9]), + ); + assert!( + tc.arena + .device_value(node, SWA) + .equal(&Tensor::from_slice(&[7i64, 8, 9])) + ); + assert_eq!(tc.evictable_size_(SWA), 3); +} + +#[test] +#[should_panic(expected = "slot already set")] +fn set_component_device_value_rejects_an_occupied_slot() { + let mut tc = core(); + tc.register_component_(Arc::new(SwaComponentForTest)); + let root = tc.arena.root(); + let node = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.set_component_device_value(tc.arena.node(node).id, SWA, Tensor::from_slice(&[7i64])); + tc.device_lru_list_mut(SWA).remove_node(node); + tc.set_component_device_value(tc.arena.node(node).id, SWA, Tensor::from_slice(&[8i64])); +} + +#[test] +#[should_panic(expected = "Swa component is not enabled")] +fn set_component_device_value_rejects_a_disabled_component() { + let mut tc = core(); + let root = tc.arena.root(); + tc.set_component_device_value(tc.arena.node(root).id, SWA, Tensor::from_slice(&[1i64])); +} + +#[test] +#[should_panic(expected = "low-priority release dispatched")] +fn dec_swa_lock_only_dispatches_lower_priority_releases() { + let mut tc: UnifiedTreeCore> = UnifiedTreeCore::new( + CacheInitParams { + swa_sliding_window_size: Some(4), + ..Default::default() + }, + vec![FULL, SWA], + ); + tc.register_component_(Arc::new(LowPriorityComponentForTest)); + let root = tc.arena.root(); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + tc.dec_swa_lock_only( + tc.arena.node(root).id, + Some(7), + &mut device_frees, + &mut host_frees, + ); +} + +#[test] +fn dec_swa_lock_only_returns_device_frees_in_the_device_dict() { + let params = CacheInitParams { + swa_sliding_window_size: Some(2), + ..Default::default() + }; + let mut tc: UnifiedTreeCore> = UnifiedTreeCore::new(params, vec![FULL, SWA]); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(a, FULL, Tensor::from_slice(&[9i64])); + tc.set_component_device_value(tc.arena.node(a).id, SWA, Tensor::from_slice(&[7i64])); + let swa = SwaComponent::new(&CacheInitParams { + swa_sliding_window_size: Some(2), + ..Default::default() + }); + let result = swa.acquire_component_lock( + &mut tc, + a, + IncLockRefResult::default(), + /* lock_host = */ false, + ); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + tc.dec_swa_lock_only( + tc.arena.node(a).id, + result.swa_uuid_for_lock, + &mut device_frees, + &mut host_frees, + ); + // The fully unlocked D-leaf's SWA value is device-evicted on release; + // the freed span is reported as the node's Full indices. + assert!(!tc.arena.has_device_value(a, SWA)); + assert_eq!(device_frees[&SWA].len(), 1); + assert!(device_frees[&SWA][0].equal(&Tensor::from_slice(&[9i64]))); + assert!(host_frees.is_empty()); +} + +#[test] +fn next_swa_uuid_counts_up_from_two() { + let mut tc = core(); + assert_eq!(tc.next_swa_uuid_(), 2); + assert_eq!(tc.next_swa_uuid_(), 3); +} + +#[test] +fn inc_lock_ref_result_defaults_carry_no_uuids() { + let result = IncLockRefResult::default(); + assert_eq!(result.swa_uuid_for_lock, None); + assert_eq!(result.swa_uuid_for_host_lock, None); +} + +// A tree with the Swa stub registered and a node carrying an SWA device value. +fn swa_valued_node(tc: &mut UnifiedTreeCore>) -> NodeIdx_ { + tc.register_component_(Arc::new(SwaComponentForTest)); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena.node_mut(a).values[SWA.idx()].value = Some(Tensor::from_slice(&[0i64])); + a +} + +#[test] +fn for_each_component_lru_visits_valued_aux_components_only() { + let mut tc = core(); + let a = swa_valued_node(&mut tc); + // A Full device value must not draw an LRU visit (Full uses leaf sets). + tc.arena + .set_device_value(a, FULL, Tensor::from_slice(&[0i64])); + let mut visited = Vec::new(); + tc.for_each_component_lru_( + a, + &mut |lru, node_id| { + lru.insert_mru(node_id); + visited.push(node_id); + }, + EvictLayer::Device, + /* skip_existing = */ false, + ); + assert_eq!(visited, vec![a]); + assert!(tc.device_lru_list(SWA).in_list(Some(a))); + assert!(!tc.device_lru_list(FULL).in_list(Some(a))); +} + +#[test] +fn for_each_component_lru_skips_valueless_components() { + let mut tc = core(); + tc.register_component_(Arc::new(SwaComponentForTest)); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let mut visits = 0; + tc.for_each_component_lru_( + a, + &mut |_, _| visits += 1, + EvictLayer::Device, + /* skip_existing = */ false, + ); + assert_eq!(visits, 0); +} + +#[test] +fn for_each_component_lru_targets_the_host_tier() { + let mut tc = core(); + tc.register_component_(Arc::new(SwaComponentForTest)); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .node_mut(a) + .state_mut_(ValueSlotIdx::host(SWA)) + .value = Some(Tensor::from_slice(&[0i64])); + let mut visited = Vec::new(); + tc.for_each_component_lru_( + a, + &mut |lru, node_id| { + lru.insert_mru(node_id); + visited.push(node_id); + }, + EvictLayer::Host, + /* skip_existing = */ false, + ); + assert_eq!(visited, vec![a]); + assert!(tc.host_lru_list(SWA).in_list(Some(a))); + // The device-tier walk sees no device value on the node. + let mut device_visits = 0; + tc.for_each_component_lru_( + a, + &mut |_, _| device_visits += 1, + EvictLayer::Device, + /* skip_existing = */ false, + ); + assert_eq!(device_visits, 0); +} + +#[test] +fn for_each_component_lru_skip_existing_spares_listed_nodes() { + let mut tc = core(); + let a = swa_valued_node(&mut tc); + tc.device_lru_list_mut(SWA).insert_mru(a); + let mut visits = 0; + tc.for_each_component_lru_( + a, + &mut |_, _| visits += 1, + EvictLayer::Device, + /* skip_existing = */ true, + ); + assert_eq!(visits, 0); + // Without the flag, the listed node is visited again. + tc.for_each_component_lru_( + a, + &mut |_, _| visits += 1, + EvictLayer::Device, + /* skip_existing = */ false, + ); + assert_eq!(visits, 1); +} + +#[test] +fn new_node_allocates_a_half_linked_stamped_node() { + let mut tc = core(); + let root = tc.arena.root(); + let before = tc.arena.node(root).last_access_counter; + let a = tc.new_node_( + /* key = */ vec![5], + root, + /* priority = */ 7, + /* hit_count = */ 3, + /* creation_counter = */ None, + /* extra_key = */ None, + ); + let node = tc.arena.node(a); + assert_eq!(node.priority, 7); + assert_eq!(node.hit_count, 3); + assert_eq!(node.key, vec![5]); + assert_eq!(node.parent(), root); + assert!(node.children.is_empty()); + // Half-linked: the parent's child map does not know the node yet. + assert!(tc.arena.root_child(None, &[5]).is_none()); + // Creation and access stamps share one fresh tick. + assert_eq!(node.creation_counter, node.last_access_counter); + assert!(node.last_access_counter > before); +} + +#[test] +fn new_node_ids_are_distinct_live_slots() { + let mut tc = core(); + let root = tc.arena.root(); + let a = tc.new_node_( + /* key = */ vec![1], + root, + /* priority = */ 0, + /* hit_count = */ 0, + /* creation_counter = */ None, + /* extra_key = */ None, + ); + let b = tc.new_node_( + /* key = */ vec![2], + root, + /* priority = */ 0, + /* hit_count = */ 0, + /* creation_counter = */ None, + /* extra_key = */ None, + ); + assert_ne!(a, b); + assert_eq!(tc.arena.resolve(tc.arena.node(a).id), a); + assert_eq!(tc.arena.resolve(tc.arena.node(b).id), b); +} + +// Chain root -> c with a 3-atom key and FULL device value, seeded as a D-leaf. +fn split_setup(tc: &mut UnifiedTreeCore>) -> NodeIdx_ { + let root = tc.arena.root(); + let c = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2, 3], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(c, FULL, Tensor::from_slice(&[10i64, 11, 12])); + tc.evictable_device_leaves.add(c); + c +} + +#[test] +fn split_wires_the_new_node_between_parent_and_child() { + let mut tc = core(); + let root = tc.arena.root(); + let c = split_setup(&mut tc); + let (new_node, action) = tc.split_node_(c, /* split_len = */ 2); + assert!(action.is_none()); + assert_eq!(tc.arena.node(new_node).key, vec![1, 2]); + assert_eq!(tc.arena.node(c).key, vec![3]); + assert_eq!(tc.arena.root_child(None, &[1]).as_ref(), Some(&new_node)); + assert_eq!( + tc.arena + .node(new_node) + .children + .get(&(KeyNamespace::default(), vec![3])), + Some(&c) + ); + assert_eq!(tc.arena.node(new_node).parent(), root); + assert_eq!(tc.arena.node(c).parent(), new_node); +} + +#[test] +fn split_redistributes_the_device_value_and_locks() { + let mut tc = core(); + let c = split_setup(&mut tc); + tc.arena + .node_mut(c) + .set_lock_ref_(ValueSlotIdx::device(FULL), 2); + let (new_node, _) = tc.split_node_(c, /* split_len = */ 2); + assert_eq!(tc.arena.device_value_len(new_node, FULL), 2); + assert_eq!(tc.arena.device_value_len(c, FULL), 1); + assert_eq!(tc.arena.device_lock_ref(new_node, FULL), 2); + assert_eq!(tc.arena.device_lock_ref(c, FULL), 2); +} + +#[test] +fn split_copies_stats_and_restamps_the_child() { + let mut tc = core(); + let c = split_setup(&mut tc); + tc.arena.node_mut(c).hit_count = 5; + let creation = tc.arena.node(c).creation_counter; + let access_before = tc.arena.node(c).last_access_counter; + let (new_node, _) = tc.split_node_(c, /* split_len = */ 2); + // The prefix node inherits hits and creation; the child gets a fresh access tick. + assert_eq!(tc.arena.node(new_node).hit_count, 5); + assert_eq!(tc.arena.node(new_node).creation_counter, creation); + assert!(tc.arena.node(c).last_access_counter > access_before); +} + +#[test] +fn split_propagates_the_child_priority() { + let mut tc = core(); + let c = split_setup(&mut tc); + tc.arena.node_mut(c).priority = 7; + let (new_node, _) = tc.split_node_(c, /* split_len = */ 2); + assert_eq!(tc.arena.node(new_node).priority, 7); + assert_eq!(tc.arena.node(c).priority, 7); +} + +#[test] +fn split_updates_the_leaf_sets() { + let mut tc = core(); + let c = split_setup(&mut tc); + let (new_node, _) = tc.split_node_(c, /* split_len = */ 2); + // The child stays the D-leaf; the prefix node has a valued child. + assert!(tc.evictable_device_leaves.contains(c)); + assert!(!tc.evictable_device_leaves.contains(new_node)); +} + +#[test] +fn split_readmits_aux_lru_cells() { + let mut tc = core(); + tc.register_component_(Arc::new(SwaComponentForTest)); + let c = split_setup(&mut tc); + tc.arena.node_mut(c).values[SWA.idx()].value = Some(Tensor::from_slice(&[0i64])); + tc.device_lru_list_mut(SWA).insert_mru(c); + // A second listed node makes the child's detach-and-readmit observable. + let root = tc.arena.root(); + let s = tc + .arena + .alloc_child( + root, + /* key = */ vec![9], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.device_lru_list_mut(SWA).insert_mru(s); + let (new_node, _) = tc.split_node_(c, /* split_len = */ 2); + // The child re-enters the SWA LRU at MRU; the value-less prefix node does not. + assert!(tc.device_lru_list(SWA).in_list(Some(c))); + assert!(!tc.device_lru_list(SWA).in_list(Some(new_node))); + assert_eq!(tc.device_lru_list(SWA).get_lru_where(|_| true), Some(s)); +} + +#[test] +#[should_panic(expected = "split_node_: the parent's page entry must map to the split child")] +fn split_panics_when_the_parent_entry_is_not_the_child() { + let mut tc = core(); + let root = tc.arena.root(); + let c = split_setup(&mut tc); + // A corrupted parent map (the page entry no longer points at c) must fail loudly. + let imposter = tc + .arena + .alloc_child( + root, + /* key = */ vec![9], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena.insert_child_edge(root, vec![1], imposter); + tc.split_node_(c, /* split_len = */ 2); +} + +#[test] +#[should_panic(expected = "split_at: split_idx 3 out of range (0, 3)")] +fn split_panics_on_a_boundary_split() { + let mut tc = core(); + let c = split_setup(&mut tc); + tc.split_node_(c, /* split_len = */ 3); +} + +#[test] +#[should_panic(expected = "split_node_: split_len 0 must be a nonzero page multiple")] +fn split_panics_on_a_zero_split_len() { + let mut tc = core(); + let c = split_setup(&mut tc); + tc.split_node_(c, /* split_len = */ 0); +} + +#[test] +fn add_new_node_creates_a_valued_child() { + let mut tc = core(); + let root = tc.arena.root(); + let mut source = Tensor::from_slice(&[10i64, 11]); + let a = tc.add_new_node_( + root, + /* key = */ vec![1, 2], + &source, + /* priority = */ 3, + /* extra_key = */ None, + ); + let node = tc.arena.node(a); + assert_eq!(node.key, vec![1, 2]); + assert_eq!(node.parent(), root); + assert_eq!(node.priority, 3); + assert_eq!(tc.arena.root_child(None, &[1]).as_ref(), Some(&a)); + assert_eq!(tc.evictable_size_(FULL), 2); + assert!(tc.evictable_device_leaves.contains(a)); + assert!(!tc.evictable_device_leaves.contains(root)); + // The stored value is a deep copy of the insert slice. + let _ = source.fill_(99); + assert!( + tc.arena + .device_value(a, FULL) + .equal(&Tensor::from_slice(&[10i64, 11])) + ); +} + +#[test] +fn add_new_node_retires_the_parent_from_the_leaf_set() { + let mut tc = core(); + let root = tc.arena.root(); + let p = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(p, FULL, Tensor::from_slice(&[0i64])); + tc.evictable_device_leaves.add(p); + let value = Tensor::from_slice(&[10i64]); + let a = tc.add_new_node_( + p, + /* key = */ vec![2], + &value, + /* priority = */ 0, + /* extra_key = */ None, + ); + // The new leaf takes over; the parent now has a valued child. + assert!(tc.evictable_device_leaves.contains(a)); + assert!(!tc.evictable_device_leaves.contains(p)); +} + +#[test] +#[should_panic(expected = "already has a child on the new node's page")] +fn add_new_node_panics_when_the_page_is_taken() { + let mut tc = core(); + let root = tc.arena.root(); + let value = Tensor::from_slice(&[10i64]); + tc.add_new_node_( + root, + /* key = */ vec![1], + &value, + /* priority = */ 0, + /* extra_key = */ None, + ); + tc.add_new_node_( + root, + /* key = */ vec![1], + &value, + /* priority = */ 0, + /* extra_key = */ None, + ); +} + +#[test] +fn unevict_restores_the_value_and_the_leaf_sets() { + // Chain root -> p (valued) -> c (evicted): p is the D-leaf until c revives. + let mut tc = core(); + let root = tc.arena.root(); + let p = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let c = tc + .arena + .alloc_child( + p, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(p, FULL, Tensor::from_slice(&[0i64])); + tc.evictable_device_leaves.add(p); + let mut fresh = Tensor::from_slice(&[20i64]); + tc.unevict_node_on_insert_(c, &fresh); + assert_eq!(tc.evictable_size_(FULL), 1); + assert!(tc.evictable_device_leaves.contains(c)); + assert!(!tc.evictable_device_leaves.contains(p)); + // The restored value is a deep copy of the fresh indices. + let _ = fresh.fill_(99); + assert!( + tc.arena + .device_value(c, FULL) + .equal(&Tensor::from_slice(&[20i64])) + ); +} + +#[test] +#[should_panic(expected = "slot already set")] +fn unevict_panics_on_a_node_that_still_has_its_value() { + let mut tc = core(); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(a, FULL, Tensor::from_slice(&[0i64])); + tc.unevict_node_on_insert_(a, &Tensor::from_slice(&[1i64])); +} + +fn match_params(key: &Vec) -> MatchPrefixParams<'_, Vec> { + MatchPrefixParams { + key, + namespace: Default::default(), + } +} + +fn match_params_in_namespace<'a>( + key: &'a Vec, + extra_key: Option<&'a str>, + cache_salt: Option<&'a str>, +) -> MatchPrefixParams<'a, Vec> { + MatchPrefixParams { + key, + namespace: KeyNamespaceRef::new(extra_key, cache_salt), + } +} + +// root -> a (key [1,2], kv [10,11]) -> b (key [3], kv [12]). +fn matched_chain(tc: &mut UnifiedTreeCore>) -> (NodeIdx_, NodeIdx_) { + let root = tc.arena.root(); + let a = tc.add_new_node_( + root, + /* key = */ vec![1, 2], + &Tensor::from_slice(&[10i64, 11]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let b = tc.add_new_node_( + a, + /* key = */ vec![3], + &Tensor::from_slice(&[12i64]), + /* priority = */ 0, + /* extra_key = */ None, + ); + (a, b) +} + +#[test] +fn match_prefix_returns_the_full_hit() { + let mut tc = core(); + let (_a, b) = matched_chain(&mut tc); + let result = tc.match_prefix(&match_params(&vec![1, 2, 3])); + assert!( + result + .device_indices + .equal(&Tensor::from_slice(&[10i64, 11, 12])) + ); + assert_eq!(result.last_device_node_id, tc.arena.node(b).id); + assert_eq!(result.last_host_node_id, tc.arena.node(b).id); + assert_eq!(result.best_match_node_id, tc.arena.node(b).id); + assert_eq!(result.host_hit_length, 0); + assert!(result.cache_actions.is_empty()); +} + +#[test] +fn match_prefix_stops_at_the_matched_depth() { + let mut tc = core(); + let (a, _b) = matched_chain(&mut tc); + let result = tc.match_prefix(&match_params(&vec![1, 2, 9])); + assert!( + result + .device_indices + .equal(&Tensor::from_slice(&[10i64, 11])) + ); + assert_eq!(result.best_match_node_id, tc.arena.node(a).id); +} + +#[test] +fn match_prefix_miss_anchors_at_the_root() { + let mut tc = core(); + matched_chain(&mut tc); + let root = tc.arena.root(); + let result = tc.match_prefix(&match_params(&vec![9])); + assert_eq!(result.device_indices.numel(), 0); + assert_eq!(result.best_match_node_id, tc.arena.node(root).id); + assert_eq!(result.last_device_node_id, tc.arena.node(root).id); +} + +#[test] +fn match_prefix_splits_on_a_partial_match() { + let mut tc = core(); + let (a, _b) = matched_chain(&mut tc); + let result = tc.match_prefix(&match_params(&vec![1, 9])); + // The walk split a at 1: the new prefix node holds [10] and anchors the result. + let prefix_node = result.best_match_node_id; + assert_ne!(prefix_node, tc.arena.node(a).id); + assert!(result.device_indices.equal(&Tensor::from_slice(&[10i64]))); + assert_eq!(tc.arena.node(tc.arena.resolve(prefix_node)).key, vec![1]); + assert_eq!(tc.arena.node(a).key, vec![2]); + assert_eq!(tc.arena.node(a).parent(), tc.arena.resolve(prefix_node)); +} + +#[test] +fn match_prefix_stops_at_a_dead_node() { + // An evicted, unbackuped child ends the traversal before it. + let mut tc = core(); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let _ = a; + let result = tc.match_prefix(&match_params(&vec![1, 2])); + assert_eq!(result.device_indices.numel(), 0); + assert_eq!(result.best_match_node_id, tc.arena.node(root).id); +} + +#[test] +fn match_prefix_page_aligns_the_query() { + let params = CacheInitParams { + page_size: 2, + ..Default::default() + }; + let mut tc: UnifiedTreeCore> = UnifiedTreeCore::new(params, vec![FULL]); + let root = tc.arena.root(); + let a = tc.add_new_node_( + root, + /* key = */ vec![1, 2], + &Tensor::from_slice(&[10i64, 11]), + /* priority = */ 0, + /* extra_key = */ None, + ); + // The trailing partial page is dropped before the walk. + let result = tc.match_prefix(&match_params(&vec![1, 2, 3])); + assert!( + result + .device_indices + .equal(&Tensor::from_slice(&[10i64, 11])) + ); + assert_eq!(result.best_match_node_id, tc.arena.node(a).id); +} + +#[test] +fn match_prefix_restamps_the_matched_path_newest_first() { + let mut tc = core(); + let (a, b) = matched_chain(&mut tc); + let before = tc.arena.node(b).last_access_counter; + tc.match_prefix(&match_params(&vec![1, 2, 3])); + let root_id = tc.arena.root(); + let root_tick = tc.arena.node(root_id).last_access_counter; + let a_tick = tc.arena.node(a).last_access_counter; + let b_tick = tc.arena.node(b).last_access_counter; + assert!(b_tick > before); + assert!(b_tick > a_tick); + assert!(a_tick > root_tick); +} + +#[test] +fn insert_first_write_creates_the_namespace() { + let mut tc = core(); + matched_chain(&mut tc); + tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some("lora-1"), None), + ..insert_params(&vec![7, 8], &[40, 41]) + }); + // The namespace is isolated under its own root edges, created by the write. + assert!(tc.arena.namespace_exists(Some("lora-1"))); + let result = tc.match_prefix(&MatchPrefixParams { + key: &vec![1, 2], + namespace: KeyNamespaceRef::new(Some("lora-1"), None), + }); + assert_eq!(result.device_indices.numel(), 0); + assert_eq!(result.best_match_node_id, tc.root_node_handle(None)); +} + +#[test] +fn match_prefix_empty_query_anchors_at_the_root() { + let mut tc = core(); + let root_handle = tc.root_node_handle(None); + for extra_key in [Some("lora-1"), None] { + let result = tc.match_prefix(&MatchPrefixParams { + key: &vec![], + namespace: KeyNamespaceRef::new(extra_key, None), + }); + assert_eq!(result.best_match_node_id, root_handle); + assert_eq!(result.last_device_node_id, root_handle); + assert_eq!(result.last_host_node_id, root_handle); + } + // A read never creates a namespace. + assert!(!tc.arena.namespace_exists(Some("lora-1"))); +} + +#[test] +fn match_prefix_skips_host_only_nodes_without_hicache() { + // Chain a(valued) -> b(evicted but backuped): the walk passes b but the + // device-only validator keeps the boundary at a. + let mut tc = core(); + let (a, b) = matched_chain(&mut tc); + let taken = tc.arena.take_device_value(b, FULL); + tc.arena.set_host_value(b, FULL, taken); + let result = tc.match_prefix(&match_params(&vec![1, 2, 3])); + assert!( + result + .device_indices + .equal(&Tensor::from_slice(&[10i64, 11])) + ); + assert_eq!(result.best_match_node_id, tc.arena.node(a).id); + assert_eq!(result.last_device_node_id, tc.arena.node(a).id); +} + +#[test] +fn match_prefix_with_hicache_advances_best_match_onto_host_nodes() { + // Chain a(valued) -> b(evicted but backuped): the device anchor stays at + // a while the hicache validators carry the best match onto b. + let mut tc = core(); + tc.set_hicache_enabled(); + let (a, b) = matched_chain(&mut tc); + let taken = tc.arena.take_device_value(b, FULL); + tc.arena.set_host_value(b, FULL, taken); + let result = tc.match_prefix(&match_params(&vec![1, 2, 3])); + assert!( + result + .device_indices + .equal(&Tensor::from_slice(&[10i64, 11])) + ); + assert_eq!(result.last_device_node_id, tc.arena.node(a).id); + assert_eq!(result.best_match_node_id, tc.arena.node(b).id); + assert_eq!(result.last_host_node_id, tc.arena.node(b).id); + assert_eq!(result.host_hit_length, 1); + assert!(result.cache_actions.is_empty()); +} + +#[test] +fn match_prefix_with_hicache_restamps_the_host_best_match() { + // Chain a(valued) -> b(host-only): the restamp walk anchors at the + // consensus best match b, not the shallower device anchor a. + let mut tc = core(); + tc.set_hicache_enabled(); + let (a, b) = matched_chain(&mut tc); + let taken = tc.arena.take_device_value(b, FULL); + tc.arena.set_host_value(b, FULL, taken); + let before = tc.arena.node(b).last_access_counter; + tc.match_prefix(&match_params(&vec![1, 2, 3])); + let a_tick = tc.arena.node(a).last_access_counter; + let b_tick = tc.arena.node(b).last_access_counter; + assert!(b_tick > before); + assert!(b_tick > a_tick); +} + +#[test] +fn match_prefix_with_hicache_sums_the_host_span_length() { + let mut tc = core(); + tc.set_hicache_enabled(); + let (a, b) = matched_chain(&mut tc); + let c = tc.add_new_node_( + b, + /* key = */ vec![4, 5], + &Tensor::from_slice(&[13i64, 14]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let taken = tc.arena.take_device_value(b, FULL); + tc.arena.set_host_value(b, FULL, taken); + let taken = tc.arena.take_device_value(c, FULL); + tc.arena.set_host_value(c, FULL, taken); + let result = tc.match_prefix(&match_params(&vec![1, 2, 3, 4, 5])); + assert!( + result + .device_indices + .equal(&Tensor::from_slice(&[10i64, 11])) + ); + assert_eq!(result.last_device_node_id, tc.arena.node(a).id); + assert_eq!(result.best_match_node_id, tc.arena.node(c).id); + assert_eq!(result.last_host_node_id, tc.arena.node(c).id); + assert_eq!(result.host_hit_length, 3); +} + +#[test] +fn match_walk_runs_every_validator_at_a_host_only_node() { + // Chain a(valued) -> b(host-only): Full's device-only validator is + // false at b, yet the aux validator must still observe b. + let mut tc = core(); + let counter = Arc::new(CountingComponentForTest::default()); + tc.register_component_(counter.clone()); + let (_a, b) = matched_chain(&mut tc); + let taken = tc.arena.take_device_value(b, FULL); + tc.arena.set_host_value(b, FULL, taken); + tc.match_prefix(&match_params(&vec![1, 2, 3])); + assert_eq!(*counter.validator_calls.lock().unwrap(), 2); +} + +#[test] +fn match_walk_runs_every_device_validator_under_hicache() { + // Both hicache folds observe both nodes: Full's device validator is + // false at the host-only b, yet the aux device validator still runs there. + let mut tc = core(); + tc.set_hicache_enabled(); + let counter = Arc::new(CountingComponentForTest::default()); + tc.register_component_(counter.clone()); + let (_a, b) = matched_chain(&mut tc); + let taken = tc.arena.take_device_value(b, FULL); + tc.arena.set_host_value(b, FULL, taken); + tc.match_prefix(&match_params(&vec![1, 2, 3])); + assert_eq!(*counter.validator_calls.lock().unwrap(), 4); +} + +#[test] +fn match_end_refresh_anchors_at_the_consensus_best_match() { + // Chain a(valued) -> b(host-only): the device anchor stays at a while + // the MatchEnd refresh dispatches on the consensus best match b. + let mut tc = core(); + tc.set_hicache_enabled(); + let recorder = Arc::new(RecordingComponentForTest::default()); + tc.register_component_(recorder.clone()); + let (a, b) = matched_chain(&mut tc); + let taken = tc.arena.take_device_value(b, FULL); + tc.arena.set_host_value(b, FULL, taken); + let result = tc.match_prefix(&match_params(&vec![1, 2, 3])); + assert_eq!(result.last_device_node_id, tc.arena.node(a).id); + assert_eq!(result.best_match_node_id, tc.arena.node(b).id); + let refreshes = recorder.refreshes.lock().unwrap(); + assert!( + refreshes + .iter() + .any(|&(phase, node)| phase == LRURefreshPhase::MatchEnd && node == b) + ); +} + +// A [Full, Swa] core with the given sliding window (page size 1). +fn swa_match_core(window: usize) -> UnifiedTreeCore> { + UnifiedTreeCore::new( + CacheInitParams { + swa_sliding_window_size: Some(window), + ..Default::default() + }, + vec![FULL, SWA], + ) +} + +// Stamp a key-covering SWA device value on the node and list it in the SWA LRU. +fn set_swa_device_and_list(tc: &mut UnifiedTreeCore>, node: NodeIdx_) { + let len = tc.arena.node(node).key.atom_len(); + tc.arena.node_mut(node).values[SWA.idx()].value = Some(Tensor::from_slice(&vec![0i64; len])); + tc.device_lru_list_mut(SWA).insert_mru(node); +} + +#[test] +fn match_walk_swa_state_resets_at_a_full_rejected_node() { + // Chain a(SWA on) -> b(Full host-only, SWA tombstone) -> c(SWA on, + // below the window): Full's validator rejects b, but the SWA validator + // must still observe b so its window run restarts before c. + let mut tc = swa_match_core(/* window = */ 2); + let root = tc.arena.root(); + let a = tc.add_new_node_( + root, + /* key = */ vec![1, 2], + &Tensor::from_slice(&[10i64, 11]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let b = tc.add_new_node_( + a, + /* key = */ vec![3], + &Tensor::from_slice(&[12i64]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let c = tc.add_new_node_( + b, + /* key = */ vec![4], + &Tensor::from_slice(&[13i64]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let taken = tc.arena.take_device_value(b, FULL); + tc.arena.set_host_value(b, FULL, taken); + set_swa_device_and_list(&mut tc, a); + set_swa_device_and_list(&mut tc, c); + let result = tc.match_prefix(&match_params(&vec![1, 2, 3, 4])); + assert_eq!(result.best_match_node_id, tc.arena.node(a).id); + assert_eq!(result.last_device_node_id, tc.arena.node(a).id); + assert!( + result + .device_indices + .equal(&Tensor::from_slice(&[10i64, 11])) + ); +} + +#[test] +fn match_prefix_swa_tombstone_holds_the_best_match_below_the_window() { + // a(SWA on) -> t(Full on, SWA tombstone) -> c(SWA on, span 1 < window): + // the best match stays at a even though Full accepts the deeper nodes. + let mut tc = swa_match_core(/* window = */ 2); + let root = tc.arena.root(); + let a = tc.add_new_node_( + root, + /* key = */ vec![1, 2], + &Tensor::from_slice(&[10i64, 11]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let t = tc.add_new_node_( + a, + /* key = */ vec![3], + &Tensor::from_slice(&[12i64]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let c = tc.add_new_node_( + t, + /* key = */ vec![4], + &Tensor::from_slice(&[13i64]), + /* priority = */ 0, + /* extra_key = */ None, + ); + set_swa_device_and_list(&mut tc, a); + set_swa_device_and_list(&mut tc, c); + let result = tc.match_prefix(&match_params(&vec![1, 2, 3, 4])); + assert_eq!(result.best_match_node_id, tc.arena.node(a).id); + assert!( + result + .device_indices + .equal(&Tensor::from_slice(&[10i64, 11])) + ); +} + +#[test] +fn match_prefix_swa_best_match_advances_at_the_window() { + // Same shape, but c spans the whole window: c revalidates and takes + // the best match past the SWA tombstone. + let mut tc = swa_match_core(/* window = */ 2); + let root = tc.arena.root(); + let a = tc.add_new_node_( + root, + /* key = */ vec![1, 2], + &Tensor::from_slice(&[10i64, 11]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let t = tc.add_new_node_( + a, + /* key = */ vec![3], + &Tensor::from_slice(&[12i64]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let c = tc.add_new_node_( + t, + /* key = */ vec![4, 5], + &Tensor::from_slice(&[13i64, 14]), + /* priority = */ 0, + /* extra_key = */ None, + ); + set_swa_device_and_list(&mut tc, a); + set_swa_device_and_list(&mut tc, c); + let result = tc.match_prefix(&match_params(&vec![1, 2, 3, 4, 5])); + assert_eq!(result.best_match_node_id, tc.arena.node(c).id); + assert!( + result + .device_indices + .equal(&Tensor::from_slice(&[10i64, 11, 12, 13, 14])) + ); +} + +#[test] +fn match_end_refresh_moves_the_swa_window_run() { + let mut tc = swa_match_core(/* window = */ 2); + let root = tc.arena.root(); + let a = tc.add_new_node_( + root, + /* key = */ vec![1, 2], + &Tensor::from_slice(&[10i64, 11]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let b = tc.add_new_node_( + a, + /* key = */ vec![3], + &Tensor::from_slice(&[12i64]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let c = tc.add_new_node_( + b, + /* key = */ vec![4], + &Tensor::from_slice(&[13i64]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let s = tc.add_new_node_( + root, + /* key = */ vec![9], + &Tensor::from_slice(&[19i64]), + /* priority = */ 0, + /* extra_key = */ None, + ); + for node in [c, b, a, s] { + set_swa_device_and_list(&mut tc, node); + } + let result = tc.match_prefix(&match_params(&vec![1, 2, 3, 4])); + assert_eq!(result.best_match_node_id, tc.arena.node(c).id); + // The walk window is sliding_window_size + page_size = 3: c, b, and + // the straddling a become the MRU run; the sentinel s stays behind. + let order: Vec = tc.device_lru_list(SWA).iter().collect(); + assert_eq!(order, vec![c, b, a, s]); +} + +#[test] +fn match_prefix_with_hicache_reports_swa_host_hits() { + // a(SWA device) -> b(Full host-only, SWA host-only): the best match + // advances onto b and finalize reports b's SWA host span. + let mut tc = swa_match_core(/* window = */ 4); + tc.set_hicache_enabled(); + let root = tc.arena.root(); + let a = tc.add_new_node_( + root, + /* key = */ vec![1, 2], + &Tensor::from_slice(&[10i64, 11]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let b = tc.add_new_node_( + a, + /* key = */ vec![3], + &Tensor::from_slice(&[12i64]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let taken = tc.arena.take_device_value(b, FULL); + tc.arena.set_host_value(b, FULL, taken); + set_swa_device_and_list(&mut tc, a); + tc.arena + .node_mut(b) + .state_mut_(ValueSlotIdx::host(SWA)) + .value = Some(Tensor::from_slice(&[0i64])); + let result = tc.match_prefix(&match_params(&vec![1, 2, 3])); + assert_eq!(result.best_match_node_id, tc.arena.node(b).id); + assert_eq!(result.last_device_node_id, tc.arena.node(a).id); + assert_eq!(result.host_hit_length, 1); + assert_eq!(result.swa_host_hit_length, 1); +} + +#[test] +fn touch_node_walkdown_keeps_the_swa_lru_order() { + let mut tc = swa_match_core(/* window = */ 2); + let root = tc.arena.root(); + let a = tc.add_new_node_( + root, + /* key = */ vec![1], + &Tensor::from_slice(&[10i64]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let b = tc.add_new_node_( + root, + /* key = */ vec![2], + &Tensor::from_slice(&[11i64]), + /* priority = */ 0, + /* extra_key = */ None, + ); + set_swa_device_and_list(&mut tc, a); + set_swa_device_and_list(&mut tc, b); + // The SWA walk-down refresh is a no-op: touching the LRU-tail node + // must not move it (window-bounded refresh runs at match/insert end). + tc.touch_node_(a); + let order: Vec = tc.device_lru_list(SWA).iter().collect(); + assert_eq!(order, vec![b, a]); +} + +#[test] +fn repeated_deep_swa_matches_keep_the_tree_sane() { + let mut tc = swa_match_core(/* window = */ 2); + let root = tc.arena.root(); + let a = tc.add_new_node_( + root, + /* key = */ vec![1, 2], + &Tensor::from_slice(&[10i64, 11]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let b = tc.add_new_node_( + a, + /* key = */ vec![3], + &Tensor::from_slice(&[12i64]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let c = tc.add_new_node_( + b, + /* key = */ vec![4], + &Tensor::from_slice(&[13i64]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let s = tc.add_new_node_( + root, + /* key = */ vec![9], + &Tensor::from_slice(&[19i64]), + /* priority = */ 0, + /* extra_key = */ None, + ); + // The size-accounted store keeps the SWA bookkeeping sanity-checkable. + for node in [c, b, a, s] { + let len = tc.arena.node(node).key.atom_len(); + tc.set_component_device_value( + tc.arena.node(node).id, + SWA, + Tensor::from_slice(&vec![0i64; len]), + ); + } + for _ in 0..3 { + let result = tc.match_prefix(&match_params(&vec![1, 2, 3, 4])); + assert!( + result + .device_indices + .equal(&Tensor::from_slice(&[10i64, 11, 12, 13])) + ); + tc.sanity_check(&[], &[]); + } +} + +#[test] +fn swa_host_backed_node_advances_best_match_but_keeps_the_device_anchor() { + let mut tc = swa_match_core(/* window = */ 2); + tc.set_hicache_enabled(); + // A wired host SWA pool makes host-only SWA gate the device match again. + tc.set_has_swa_host_pool(); + let root = tc.arena.root(); + let a = tc.add_new_node_( + root, + /* key = */ vec![1, 2], + &Tensor::from_slice(&[10i64, 11]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let b = tc.add_new_node_( + a, + /* key = */ vec![3], + &Tensor::from_slice(&[12i64]), + /* priority = */ 0, + /* extra_key = */ None, + ); + set_swa_device_and_list(&mut tc, a); + tc.arena + .node_mut(b) + .state_mut_(ValueSlotIdx::host(SWA)) + .value = Some(Tensor::from_slice(&[0i64])); + tc.host_lru_list_mut(SWA).insert_mru(b); + let result = tc.match_prefix(&match_params(&vec![1, 2, 3])); + assert!( + result + .device_indices + .equal(&Tensor::from_slice(&[10i64, 11])) + ); + assert_eq!(result.last_device_node_id, tc.arena.node(a).id); + assert_eq!(result.best_match_node_id, tc.arena.node(b).id); + assert_eq!(result.host_hit_length, 0); + assert_eq!(result.swa_host_hit_length, 1); +} + +fn insert_params<'k>(key: &'k Vec, value: &[i64]) -> InsertParams<'k, Vec> { + InsertParams { + key, + namespace: Default::default(), + value: Tensor::from_slice(value), + mamba_value: None, + prev_prefix_len: 0, + swa_evicted_seqlen: 0, + chunked: false, + priority: 0, + track_adopted_ranges: false, + } +} + +fn tracked_insert_params<'k>(key: &'k Vec, value: &[i64]) -> InsertParams<'k, Vec> { + InsertParams { + track_adopted_ranges: true, + ..insert_params(key, value) + } +} + +#[test] +fn adopted_ranges_are_opt_in_and_coalesce() { + let mut untracked = InsertResult::default(); + untracked.record_adopted_range(FULL, 0, 2); + assert_eq!(untracked.adopted_ranges, None); + + let mut tracked = InsertResult { + adopted_ranges: Some(HashMap::new()), + ..InsertResult::default() + }; + tracked.record_adopted_range(FULL, 2, 4); + tracked.record_adopted_range(FULL, 4, 6); + tracked.record_adopted_range(FULL, 8, 10); + tracked.record_adopted_range(FULL, 7, 9); + tracked.record_adopted_range(SWA, 3, 3); + assert_eq!( + tracked.adopted_ranges.as_ref().unwrap()[&FULL], + [(2, 6), (7, 10)] + ); + assert!(!tracked.adopted_ranges.as_ref().unwrap().contains_key(&SWA)); +} + +#[test] +fn insert_reports_new_and_unevicted_full_ranges() { + let mut tc = core(); + let first = tc.insert(&tracked_insert_params(&vec![1, 2, 3], &[10, 11, 12])); + assert_eq!(first.adopted_ranges.as_ref().unwrap()[&FULL], [(0, 3)]); + + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 3])) + .best_match_node_id; + let leaf = tc.arena.resolve(leaf); + let _ = tc.arena.take_device_value(leaf, FULL); + tc.component_state_mut(FULL).evictable_size = 0; + tc.evictable_device_leaves.discard(leaf); + + let restored = tc.insert(&tracked_insert_params(&vec![1, 2, 3, 4], &[20, 21, 22, 23])); + assert_eq!(restored.adopted_ranges.as_ref().unwrap()[&FULL], [(0, 4)]); +} + +fn insert_params_in_namespace<'a>( + key: &'a Vec, + value: &[i64], + extra_key: Option<&'a str>, + cache_salt: Option<&'a str>, +) -> InsertParams<'a, Vec> { + InsertParams { + namespace: KeyNamespaceRef::new(extra_key, cache_salt), + ..insert_params(key, value) + } +} + +#[test] +fn insert_creates_a_leaf_and_matches_back() { + let mut tc = core(); + let result = tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + assert_eq!(result.prefix_len, 0); + assert!(!result.mamba_exist); + assert!(result.cache_actions.is_empty()); + assert_eq!(tc.evictable_size_(FULL), 3); + let matched = tc.match_prefix(&match_params(&vec![1, 2, 3])); + assert_eq!(result.last_device_node_id, Some(matched.best_match_node_id)); + assert!( + matched + .device_indices + .equal(&Tensor::from_slice(&[10i64, 11, 12])) + ); +} + +#[test] +fn insert_full_overlap_frees_the_duplicates() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + let result = tc.insert(&insert_params(&vec![1, 2, 3], &[20, 21, 22])); + assert_eq!(result.prefix_len, 3); + // Nothing was consumed by any component: the whole overlap is duplicate. + let [CacheAction::FreeDeviceKV(freed)] = result.cache_actions.as_slice() else { + panic!( + "expected one FreeDeviceKV action, got {:?}", + action_kinds(&result.cache_actions) + ); + }; + assert_eq!(freed.len(), 1); + assert!(freed[0].equal(&Tensor::from_slice(&[20i64, 21, 22]))); +} + +#[test] +fn insert_prev_prefix_len_narrows_the_dup_window() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + let result = tc.insert(&InsertParams { + prev_prefix_len: 2, + ..insert_params(&vec![1, 2, 3], &[20, 21, 22]) + }); + let [CacheAction::FreeDeviceKV(freed)] = result.cache_actions.as_slice() else { + panic!( + "expected one FreeDeviceKV action, got {:?}", + action_kinds(&result.cache_actions) + ); + }; + assert!(freed[0].equal(&Tensor::from_slice(&[22i64]))); +} + +#[test] +fn insert_extends_an_existing_prefix() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + let result = tc.insert(&insert_params(&vec![1, 2, 3, 4, 5], &[20, 21, 22, 13, 14])); + assert_eq!(result.prefix_len, 3); + // The existing node's overlap is duplicate; the new suffix is kept. + let [CacheAction::FreeDeviceKV(freed)] = result.cache_actions.as_slice() else { + panic!( + "expected one FreeDeviceKV action, got {:?}", + action_kinds(&result.cache_actions) + ); + }; + assert_eq!(freed.len(), 1); + assert!(freed[0].equal(&Tensor::from_slice(&[20i64, 21, 22]))); + let matched = tc.match_prefix(&match_params(&vec![1, 2, 3, 4, 5])); + assert!( + matched + .device_indices + .equal(&Tensor::from_slice(&[10i64, 11, 12, 13, 14])) + ); +} + +#[test] +fn insert_prev_prefix_len_spans_a_multi_node_walk() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + tc.insert(&insert_params(&vec![1, 2, 3, 4, 5], &[20, 21, 22, 13, 14])); + // The request already matched [1,2,3]: only the second node's overlap + // is duplicate. + let result = tc.insert(&InsertParams { + prev_prefix_len: 3, + ..insert_params(&vec![1, 2, 3, 4, 5], &[30, 31, 32, 33, 34]) + }); + assert_eq!(result.prefix_len, 5); + let [CacheAction::FreeDeviceKV(freed)] = result.cache_actions.as_slice() else { + panic!( + "expected one FreeDeviceKV action, got {:?}", + action_kinds(&result.cache_actions) + ); + }; + assert_eq!(freed.len(), 1); + assert!(freed[0].equal(&Tensor::from_slice(&[33i64, 34]))); +} + +#[test] +fn insert_prev_prefix_len_narrows_mid_node_on_a_multi_node_walk() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + tc.insert(&insert_params(&vec![1, 2, 3, 4, 5], &[20, 21, 22, 13, 14])); + // prev_prefix_len 4 lands mid second node: only its last token is duplicate. + let result = tc.insert(&InsertParams { + prev_prefix_len: 4, + ..insert_params(&vec![1, 2, 3, 4, 5], &[30, 31, 32, 33, 34]) + }); + assert_eq!(result.prefix_len, 5); + let [CacheAction::FreeDeviceKV(freed)] = result.cache_actions.as_slice() else { + panic!( + "expected one FreeDeviceKV action, got {:?}", + action_kinds(&result.cache_actions) + ); + }; + assert_eq!(freed.len(), 1); + assert!(freed[0].equal(&Tensor::from_slice(&[34i64]))); +} + +#[test] +fn insert_splits_on_a_partial_overlap() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + let result = tc.insert(&insert_params(&vec![1, 2, 9], &[20, 21, 29])); + assert_eq!(result.prefix_len, 2); + // Both suffixes live under the split prefix node. + let matched = tc.match_prefix(&match_params(&vec![1, 2, 9])); + assert!( + matched + .device_indices + .equal(&Tensor::from_slice(&[10i64, 11, 29])) + ); + let matched = tc.match_prefix(&match_params(&vec![1, 2, 3])); + assert!( + matched + .device_indices + .equal(&Tensor::from_slice(&[10i64, 11, 12])) + ); +} + +#[test] +fn insert_unevicts_a_tombstoned_node() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + let a = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + let _ = tc.arena.take_device_value(tc.arena.resolve(a), FULL); + tc.component_state_mut(FULL).evictable_size = 0; + tc.evictable_device_leaves.discard(tc.arena.resolve(a)); + let result = tc.insert(&insert_params(&vec![1, 2], &[20, 21])); + assert_eq!(result.prefix_len, 2); + // The fresh KV revives the node; nothing is duplicate. + assert!(result.cache_actions.is_empty()); + assert!( + tc.arena + .device_value(tc.arena.resolve(a), FULL) + .equal(&Tensor::from_slice(&[20i64, 21])) + ); + assert_eq!(tc.evictable_size_(FULL), 2); + assert!(tc.evictable_device_leaves.contains(tc.arena.resolve(a))); +} + +#[test] +fn insert_priority_floor_applies_along_the_path() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + let a = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + tc.insert(&InsertParams { + priority: 5, + ..insert_params(&vec![1, 2], &[20, 21]) + }); + assert_eq!(tc.arena.node(tc.arena.resolve(a)).priority, 5); +} + +#[test] +fn insert_chunked_skips_the_hit_count() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + let a = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + let hits_before = tc.arena.node(tc.arena.resolve(a)).hit_count; + tc.insert(&InsertParams { + chunked: true, + ..insert_params(&vec![1, 2], &[20, 21]) + }); + assert_eq!(tc.arena.node(tc.arena.resolve(a)).hit_count, hits_before); +} + +#[test] +fn insert_extension_bumps_the_traversed_node_hit_count_once() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + let a = tc + .match_prefix(&match_params(&vec![1, 2, 3])) + .best_match_node_id; + let hits_before = tc.arena.node(tc.arena.resolve(a)).hit_count; + tc.insert(&insert_params(&vec![1, 2, 3, 4, 5], &[20, 21, 22, 13, 14])); + assert_eq!( + tc.arena.node(tc.arena.resolve(a)).hit_count, + hits_before + 1 + ); +} + +#[test] +fn insert_full_overlap_bumps_the_hit_count_once() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + let a = tc + .match_prefix(&match_params(&vec![1, 2, 3])) + .best_match_node_id; + let hits_before = tc.arena.node(tc.arena.resolve(a)).hit_count; + // The walk already counted the full overlap; the target is no new leaf. + tc.insert(&insert_params(&vec![1, 2, 3], &[20, 21, 22])); + assert_eq!( + tc.arena.node(tc.arena.resolve(a)).hit_count, + hits_before + 1 + ); +} + +#[test] +fn insert_threshold_crossing_emits_the_backup_kv_action() { + let params = CacheInitParams { + write_through_threshold: 1, + ..Default::default() + }; + let mut tc: UnifiedTreeCore> = UnifiedTreeCore::new(params, vec![FULL]); + tc.set_hicache_enabled(); + let result = tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 3])) + .best_match_node_id; + let backups: Vec<_> = result + .cache_actions + .iter() + .filter_map(|action| match action { + CacheAction::BackupKV(backup) => Some(backup.node_ids.clone()), + _ => None, + }) + .collect(); + assert_eq!(backups, vec![vec![leaf]]); +} + +#[test] +fn mark_write_through_pending_stamps_the_node_id_as_the_ack() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1], &[10])); + let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; + tc.mark_write_through_pending(leaf); + assert_eq!( + tc.arena + .node(tc.arena.resolve(leaf)) + .write_through_pending_id, + Some(leaf) + ); +} + +#[test] +fn finish_write_through_clears_only_the_matching_ack() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1], &[10])); + let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; + tc.mark_write_through_pending(leaf); + tc.finish_write_through(vec![leaf], /* ack_id = */ 999_999); + assert_eq!( + tc.arena + .node(tc.arena.resolve(leaf)) + .write_through_pending_id, + Some(leaf) + ); + tc.finish_write_through(vec![leaf], /* ack_id = */ leaf); + assert_eq!( + tc.arena + .node(tc.arena.resolve(leaf)) + .write_through_pending_id, + None + ); +} + +#[test] +fn backup_kv_action_chains_unbacked_ancestors_first() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1], &[10])); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + let a = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; + let b = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + let c = tc + .match_prefix(&match_params(&vec![1, 2, 3])) + .best_match_node_id; + // a is backuped: the chain stops there and orders ancestors first. + tc.arena + .set_host_value(tc.arena.resolve(a), FULL, Tensor::from_slice(&[20i64])); + let action = tc.build_backup_kv_action_( + tc.arena.node(tc.arena.resolve(c)), + /* write_back = */ false, + ); + assert_eq!(action.node_ids, vec![b, c]); + let action = tc.build_backup_kv_action_( + tc.arena.node(tc.arena.resolve(c)), + /* write_back = */ true, + ); + assert_eq!(action.node_ids, vec![c]); +} + +#[test] +fn split_of_a_pending_node_transfers_the_ack_and_emits_the_replace_action() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + let node = tc + .match_prefix(&match_params(&vec![1, 2, 3])) + .best_match_node_id; + tc.mark_write_through_pending(node); + let (new_node, action) = tc.split_node_(tc.arena.resolve(node), /* split_len = */ 1); + assert_eq!(tc.arena.node(new_node).write_through_pending_id, Some(node)); + assert_eq!( + tc.arena + .node(tc.arena.resolve(node)) + .write_through_pending_id, + Some(node) + ); + match action { + Some(CacheAction::ReplaceWriteThroughOnNodeSplit { + ack_id, + old_node_id, + new_node_id, + new_child_node_id, + }) => { + assert_eq!(ack_id, node); + assert_eq!(old_node_id, node); + assert_eq!(new_node_id, tc.arena.node(new_node).id); + assert_eq!(new_child_node_id, node); + } + other => panic!("expected the replace action, got {:?}", other.is_some()), + } +} + +#[test] +fn insert_does_not_hash_without_storage() { + let mut tc = UnifiedTreeCore::new( + CacheInitParams { + page_size: 2, + ..CacheInitParams::default() + }, + vec![FULL], + ); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + assert_eq!(tc.arena.node(tc.arena.resolve(leaf)).hash_value, None); +} + +#[test] +fn insert_hashes_pages_chained_from_the_parent_when_storage_is_on() { + // Expected values are literals produced by the python native hash. + let mut tc = UnifiedTreeCore::new( + CacheInitParams { + page_size: 2, + ..CacheInitParams::default() + }, + vec![FULL], + ); + tc.set_enable_storage(true); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + tc.insert(&insert_params(&vec![1, 2, 7, 8], &[10, 11, 12, 13])); + let parent = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + assert_eq!( + tc.arena.node(tc.arena.resolve(parent)).hash_value, + Some(vec![ + "34fb5c825de7ca4aea6e712f19d439c1da0c92c37b423936c5f618545ca4fa1f".to_string() + ]) + ); + let child = tc + .match_prefix(&match_params(&vec![1, 2, 7, 8])) + .best_match_node_id; + assert_eq!( + tc.arena.node(tc.arena.resolve(child)).hash_value, + Some(vec![ + "0bfa9b9c6fd727c7410b6d42b753439911022d34cc6ef99ac43ed7724aa48a75".to_string() + ]) + ); + // The prefix walk concatenates the chain in root-to-node order. + assert_eq!( + tc.arena.prefix_hash_values(Some(tc.arena.resolve(child))), + vec![ + "34fb5c825de7ca4aea6e712f19d439c1da0c92c37b423936c5f618545ca4fa1f".to_string(), + "0bfa9b9c6fd727c7410b6d42b753439911022d34cc6ef99ac43ed7724aa48a75".to_string(), + ] + ); + assert_eq!( + tc.arena.prefix_hash_values(Some(tc.arena.resolve(parent))), + vec!["34fb5c825de7ca4aea6e712f19d439c1da0c92c37b423936c5f618545ca4fa1f".to_string()] + ); +} + +fn events_core(page_size: usize) -> UnifiedTreeCore> { + UnifiedTreeCore::new( + CacheInitParams { + page_size, + enable_kv_cache_events: true, + ..CacheInitParams::default() + }, + vec![FULL], + ) +} + +#[test] +fn take_events_is_empty_when_events_are_disabled() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + assert_eq!(tc.take_events(), Vec::new()); +} + +#[test] +fn insert_coalesces_parent_linked_block_stores() { + let mut tc = events_core(2); + tc.insert(&insert_params(&vec![1, 2, 7, 8], &[10, 11, 12, 13])); + let hashes = crate::node::get_hash_str::>(&[1, 2, 7, 8], None, 2); + assert_eq!( + tc.take_events(), + vec![KvCacheEvent::BlockStored { + block_hashes: hashes + .iter() + .map(|hash| crate::node::hash_str_to_int64(hash)) + .collect(), + parent_block_hash: None, + token_ids: vec![1, 2, 7, 8], + block_size: 2, + medium: StorageMedium::Gpu, + cache_salt: None, + }] + ); + // Events hash lazily even though the storage tier is off. + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 7, 8])) + .best_match_node_id; + assert_eq!( + tc.arena.node(tc.arena.resolve(leaf)).hash_value, + Some(hashes) + ); + assert!(tc.salted_event_hashes.is_empty()); +} + +#[test] +fn salted_event_hashes_are_sparse_and_removed_with_the_node() { + let mut tc = events_core(2); + let key = vec![1, 2, 7, 8]; + tc.insert(&insert_params_in_namespace( + &key, + &[10, 11, 12, 13], + None, + Some("tenant-a"), + )); + tc.take_events(); + + let leaf = tc + .match_prefix(&match_params_in_namespace(&key, None, Some("tenant-a"))) + .best_match_node_id; + let leaf_idx = tc.arena.resolve(leaf); + assert_eq!(tc.salted_event_hashes[&leaf].len(), 2); + assert_eq!( + tc.arena.node(leaf_idx).hash_value, + Some(crate::node::get_hash_str::>(&key, None, 2)) + ); + + let mut tracker = HashMap::from([(FULL, 0)]); + let (mut device_frees, mut host_frees) = (HashMap::new(), HashMap::new()); + tc.evict_device_start(FULL, key.len()); + let (candidate, step) = tc.evict_device_next_node(FULL, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + let (_, step) = tc.evict_device_leaf(candidate.unwrap(), false); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + tc.evict_device_end(FULL); + tc.take_events(); + assert!(tc.salted_event_hashes.is_empty()); + + tc.insert(&insert_params_in_namespace( + &key, + &[20, 21, 22, 23], + None, + Some("tenant-a"), + )); + assert!(!tc.salted_event_hashes.is_empty()); + tc.reset(); + assert!(tc.salted_event_hashes.is_empty()); +} + +#[test] +fn salted_event_hashes_survive_node_split() { + let mut tc = events_core(2); + let original = vec![1, 2, 3, 4]; + tc.insert(&insert_params_in_namespace( + &original, + &[10, 11, 12, 13], + None, + Some("tenant-a"), + )); + let original_leaf = tc + .match_prefix(&match_params_in_namespace( + &original, + None, + Some("tenant-a"), + )) + .best_match_node_id; + let original_hashes = tc.salted_event_hashes[&original_leaf].clone(); + tc.take_events(); + + let branch = vec![1, 2, 5, 6]; + tc.insert(&insert_params_in_namespace( + &branch, + &[20, 21, 22, 23], + None, + Some("tenant-a"), + )); + tc.take_events(); + + let split_child = tc + .match_prefix(&match_params_in_namespace( + &original, + None, + Some("tenant-a"), + )) + .best_match_node_id; + let split_parent_idx = tc.arena.node(tc.arena.resolve(split_child)).parent(); + let split_parent = tc.arena.node(split_parent_idx).id; + assert_eq!(tc.salted_event_hashes[&split_parent], original_hashes[..1]); + assert_eq!(tc.salted_event_hashes[&split_child], original_hashes[1..]); +} + +#[test] +fn salted_event_hash_walk_is_iterative_and_on_demand() { + let mut tc = events_core(1); + let mut parent = tc.arena.root(); + for token in 1..=1100 { + parent = tc + .arena + .alloc_child_in_namespace( + parent, + vec![token], + 0, + KeyNamespaceRef::new(None, Some("tenant-a")), + ) + .unwrap(); + } + assert!(tc.salted_event_hashes.is_empty()); + tc.ensure_salted_event_hashes_(parent); + assert_eq!(tc.salted_event_hashes.len(), 1100); +} + +#[test] +fn event_coalescing_respects_store_remove_and_clear_boundaries() { + let mut tc = events_core(2); + tc.enqueue_kv_event_(KvCacheEvent::BlockStored { + block_hashes: vec![1], + parent_block_hash: None, + token_ids: vec![10, 11], + block_size: 2, + medium: StorageMedium::Gpu, + cache_salt: None, + }); + assert_eq!(tc.kv_event_queue.len(), 1); + // A different block size must not join the parent-linked store tail. + tc.enqueue_kv_event_(KvCacheEvent::BlockStored { + block_hashes: vec![2], + parent_block_hash: Some(1), + token_ids: vec![12], + block_size: 1, + medium: StorageMedium::Gpu, + cache_salt: None, + }); + assert_eq!(tc.kv_event_queue.len(), 2); + // Matching size and parent are still separated across media. + tc.enqueue_kv_event_(KvCacheEvent::BlockStored { + block_hashes: vec![3], + parent_block_hash: Some(2), + token_ids: vec![13], + block_size: 1, + medium: StorageMedium::Cpu, + cache_salt: None, + }); + assert_eq!(tc.kv_event_queue.len(), 3); + // Matching size and medium are still separated without the parent link. + tc.enqueue_kv_event_(KvCacheEvent::BlockStored { + block_hashes: vec![4], + parent_block_hash: None, + token_ids: vec![14], + block_size: 1, + medium: StorageMedium::Cpu, + cache_salt: None, + }); + assert_eq!(tc.kv_event_queue.len(), 4); + tc.enqueue_kv_event_(KvCacheEvent::BlockRemoved { + block_hashes: vec![1], + medium: StorageMedium::Gpu, + }); + assert_eq!(tc.kv_event_queue.len(), 5); + tc.enqueue_kv_event_(KvCacheEvent::BlockRemoved { + block_hashes: vec![2, 3], + medium: StorageMedium::Gpu, + }); + assert_eq!(tc.kv_event_queue.len(), 5); + assert!(matches!( + tc.kv_event_queue.last(), + Some(KvCacheEvent::BlockRemoved { block_hashes, .. }) + if block_hashes.as_slice() == [1, 2, 3] + )); + tc.enqueue_kv_event_(KvCacheEvent::BlockRemoved { + block_hashes: vec![4], + medium: StorageMedium::Cpu, + }); + assert_eq!(tc.kv_event_queue.len(), 6); + tc.record_all_cleared_event(); + assert_eq!(tc.kv_event_queue.len(), 7); + tc.enqueue_kv_event_(KvCacheEvent::BlockRemoved { + block_hashes: vec![5], + medium: StorageMedium::Cpu, + }); + assert_eq!(tc.kv_event_queue.len(), 8); + assert!(matches!( + &tc.kv_event_queue[6], + KvCacheEvent::AllBlocksCleared + )); + + tc.kv_event_queue.clear(); + tc.enqueue_kv_event_(KvCacheEvent::BlockStored { + block_hashes: vec![1], + parent_block_hash: None, + token_ids: vec![10, 11], + block_size: 2, + medium: StorageMedium::Gpu, + cache_salt: Some(Arc::from("tenant-a")), + }); + tc.enqueue_kv_event_(KvCacheEvent::BlockStored { + block_hashes: vec![2], + parent_block_hash: Some(1), + token_ids: vec![12, 13], + block_size: 2, + medium: StorageMedium::Gpu, + cache_salt: Some(Arc::from("tenant-b")), + }); + assert_eq!(tc.kv_event_queue.len(), 2); +} + +#[test] +fn eviction_emits_block_removed_with_all_page_hashes() { + let mut tc = events_core(2); + tc.insert(&insert_params(&vec![1, 2, 7, 8], &[10, 11, 12, 13])); + let _ = tc.take_events(); + let mut tracker = HashMap::from([(FULL, 0)]); + let (mut device_frees, mut host_frees) = (HashMap::new(), HashMap::new()); + tc.evict_device_start(FULL, 100); + loop { + let (node, step) = tc.evict_device_next_node(FULL, &tracker); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + let Some(node) = node else { break }; + let (_, step) = tc.evict_device_leaf(node, /* is_write_back = */ false); + accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees); + } + tc.evict_device_end(FULL); + let hashes = crate::node::get_hash_str::>(&[1, 2, 7, 8], None, 2); + assert_eq!( + tc.take_events(), + vec![KvCacheEvent::BlockRemoved { + block_hashes: hashes + .iter() + .map(|h| crate::node::hash_str_to_int64(h)) + .collect(), + medium: StorageMedium::Gpu, + }] + ); +} + +#[test] +fn bigram_insert_events_carry_pair_token_payloads() { + let mut tc = UnifiedTreeCore::>::new( + CacheInitParams { + enable_kv_cache_events: true, + ..CacheInitParams::default() + }, + vec![FULL], + ); + let key: Vec<(i64, i64)> = vec![(1, 2), (2, 3)]; + tc.insert(&InsertParams { + key: &key, + namespace: Default::default(), + value: Tensor::from_slice(&[10i64, 11]), + mamba_value: None, + prev_prefix_len: 0, + swa_evicted_seqlen: 0, + chunked: false, + priority: 0, + track_adopted_ranges: false, + }); + let hashes = crate::node::get_hash_str::>(&key, None, 1); + assert_eq!( + tc.take_events(), + vec![KvCacheEvent::BlockStored { + block_hashes: hashes + .iter() + .map(|hash| crate::node::hash_str_to_int64(hash)) + .collect(), + parent_block_hash: None, + token_ids: vec![(1, 2), (2, 3)], + block_size: 1, + medium: StorageMedium::Gpu, + cache_salt: None, + }] + ); +} + +#[test] +fn finish_write_through_emits_cpu_stored_events() { + let mut tc = events_core(1); + tc.set_hicache_enabled(); + tc.insert(&insert_params(&vec![1], &[10])); + let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; + let _ = tc.take_events(); + tc.finish_write_through(vec![leaf], leaf); + let hashes = crate::node::get_hash_str::>(&[1], None, 1); + assert_eq!( + tc.take_events(), + vec![KvCacheEvent::BlockStored { + block_hashes: vec![crate::node::hash_str_to_int64(&hashes[0])], + parent_block_hash: None, + token_ids: vec![1], + block_size: 1, + medium: StorageMedium::Cpu, + cache_salt: None, + }] + ); +} + +// A demoted (host-only) single-token leaf with the event queue drained. +fn demoted_events_leaf(tc: &mut UnifiedTreeCore>) -> NodeIdx_ { + tc.set_hicache_enabled(); + tc.insert(&insert_params(&vec![1], &[10])); + let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; + tc.commit_backup(leaf, Tensor::from_slice(&[100i64]), HashMap::new()); + tc.demote(leaf); + let _ = tc.take_events(); + tc.arena.resolve(leaf) +} + +#[test] +fn host_eviction_emits_a_cpu_block_removed() { + let mut tc = events_core(1); + demoted_events_leaf(&mut tc); + tc.drive_host_eviction(FULL, 1); + let hashes = crate::node::get_hash_str::>(&[1], None, 1); + assert_eq!( + tc.take_events(), + vec![KvCacheEvent::BlockRemoved { + block_hashes: vec![crate::node::hash_str_to_int64(&hashes[0])], + medium: StorageMedium::Cpu, + }] + ); +} + +#[test] +fn load_back_commit_emits_gpu_stored_events() { + let mut tc = events_core(1); + let leaf = demoted_events_leaf(&mut tc); + let (kv_xfer, comp_xfers) = tc.build_load_back_spec(tc.arena.node(leaf).id, None); + tc.commit_load_back( + tc.arena.node(leaf).id, + Tensor::from_slice(&[50i64]), + kv_xfer, + comp_xfers, + ); + let hashes = crate::node::get_hash_str::>(&[1], None, 1); + assert_eq!( + tc.take_events(), + vec![KvCacheEvent::BlockStored { + block_hashes: vec![crate::node::hash_str_to_int64(&hashes[0])], + parent_block_hash: None, + token_ids: vec![1], + block_size: 1, + medium: StorageMedium::Gpu, + cache_salt: None, + }] + ); +} + +#[test] +fn unevict_on_insert_emits_a_gpu_stored_event() { + let mut tc = events_core(1); + demoted_events_leaf(&mut tc); + tc.insert(&insert_params(&vec![1], &[60])); + let hashes = crate::node::get_hash_str::>(&[1], None, 1); + assert_eq!( + tc.take_events(), + vec![KvCacheEvent::BlockStored { + block_hashes: vec![crate::node::hash_str_to_int64(&hashes[0])], + parent_block_hash: None, + token_ids: vec![1], + block_size: 1, + medium: StorageMedium::Gpu, + cache_salt: None, + }] + ); +} + +#[test] +fn drop_subtree_emits_removals_for_host_descendants_then_the_leaf() { + let mut tc = UnifiedTreeCore::new( + CacheInitParams { + is_write_back: true, + enable_hicache: true, + enable_kv_cache_events: true, + ..CacheInitParams::default() + }, + vec![FULL], + ); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + let child = tc + .arena + .alloc_child( + tc.arena.resolve(leaf), + /* key = */ vec![3, 4], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(child, FULL, Tensor::from_slice(&[20i64, 21])); + tc.update_evictable_leaf_sets_(child); + tc.update_evictable_leaf_sets_(tc.arena.resolve(leaf)); + let _ = tc.take_events(); + let (dropped, _step) = tc.drop_subtree_no_host(leaf); + assert!(dropped); + // The leaf hashed lazily at its insert store event; the host-only + // child hashes lazily at removal, chaining from the leaf. + let leaf_hashes = crate::node::get_hash_str::>(&[1, 2], None, 1); + let child_hashes = + crate::node::get_hash_str::>(&[3, 4], leaf_hashes.last().map(String::as_str), 1); + assert_eq!( + tc.take_events(), + vec![ + KvCacheEvent::BlockRemoved { + block_hashes: child_hashes + .iter() + .map(|h| crate::node::hash_str_to_int64(h)) + .collect(), + medium: StorageMedium::Cpu, + }, + KvCacheEvent::BlockRemoved { + block_hashes: leaf_hashes + .iter() + .map(|h| crate::node::hash_str_to_int64(h)) + .collect(), + medium: StorageMedium::Gpu, + }, + ] + ); +} + +#[test] +fn all_cleared_event_queues_and_take_drains() { + let mut tc = events_core(1); + tc.record_all_cleared_event(); + assert_eq!(tc.take_events(), vec![KvCacheEvent::AllBlocksCleared]); + assert_eq!(tc.take_events(), Vec::new()); +} + +#[test] +fn split_insert_stores_only_the_new_block_chained_to_the_split_parent() { + let mut tc = events_core(2); + tc.insert(&insert_params(&vec![1, 2, 3, 4], &[10, 11, 12, 13])); + let _ = tc.take_events(); + tc.insert(&insert_params(&vec![1, 2, 5, 6], &[20, 21, 22, 23])); + let base_hashes = crate::node::get_hash_str::>(&[1, 2, 3, 4], None, 2); + let leaf_hashes = + crate::node::get_hash_str::>(&[5, 6], Some(base_hashes[0].as_str()), 2); + // Only the diverging suffix is stored; the matched prefix is not re-published. + assert_eq!( + tc.take_events(), + vec![KvCacheEvent::BlockStored { + block_hashes: vec![crate::node::hash_str_to_int64(&leaf_hashes[0])], + parent_block_hash: Some(crate::node::hash_str_to_int64(&base_hashes[0])), + token_ids: vec![5, 6], + block_size: 2, + medium: StorageMedium::Gpu, + cache_salt: None, + }] + ); + // The split divided the page hashes between the two fragments. + let parent = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + assert_eq!( + tc.arena.node(tc.arena.resolve(parent)).hash_value, + Some(vec![base_hashes[0].clone()]) + ); + let child = tc + .match_prefix(&match_params(&vec![1, 2, 3, 4])) + .best_match_node_id; + assert_eq!( + tc.arena.node(tc.arena.resolve(child)).hash_value, + Some(vec![base_hashes[1].clone()]) + ); +} + +#[test] +fn finish_write_through_after_a_split_publishes_both_fragments() { + let mut tc = events_core(2); + tc.set_hicache_enabled(); + tc.insert(&insert_params(&vec![1, 2, 3, 4], &[10, 11, 12, 13])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 3, 4])) + .best_match_node_id; + tc.mark_write_through_pending(leaf); + let _ = tc.take_events(); + let result = tc.insert(&insert_params(&vec![1, 2, 5, 6], &[20, 21, 22, 23])); + let new_node_id = result + .cache_actions + .iter() + .find_map(|action| match action { + CacheAction::ReplaceWriteThroughOnNodeSplit { + ack_id, + new_node_id, + .. + } => { + assert_eq!(*ack_id, leaf); + Some(*new_node_id) + } + _ => None, + }) + .expect("the split relocates the pending write-through"); + // Nothing reaches the host tier before the ack. + assert!(tc.take_events().iter().all(|event| !matches!( + event, + KvCacheEvent::BlockStored { + medium: StorageMedium::Cpu, + .. + } + ))); + tc.commit_backup( + new_node_id, + Tensor::from_slice(&[100i64, 101]), + HashMap::new(), + ); + tc.commit_backup(leaf, Tensor::from_slice(&[102i64, 103]), HashMap::new()); + tc.finish_write_through(vec![new_node_id, leaf], /* ack_id = */ leaf); + let hashes = crate::node::get_hash_str::>(&[1, 2, 3, 4], None, 2); + assert_eq!( + tc.take_events(), + vec![KvCacheEvent::BlockStored { + block_hashes: hashes + .iter() + .map(|hash| crate::node::hash_str_to_int64(hash)) + .collect(), + parent_block_hash: None, + token_ids: vec![1, 2, 3, 4], + block_size: 2, + medium: StorageMedium::Cpu, + cache_salt: None, + }] + ); + // The matching ack cleared the pending mark on both fragments. + assert_eq!( + tc.arena + .node(tc.arena.resolve(new_node_id)) + .write_through_pending_id, + None + ); + assert_eq!( + tc.arena + .node(tc.arena.resolve(leaf)) + .write_through_pending_id, + None + ); +} + +#[test] +fn prefetch_anchor_info_maps_the_namespace() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some("chat"), Some("tenant-a")), + ..insert_params(&vec![7, 8], &[20, 21]) + }); + let plain = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + assert_eq!(tc.prefetch_anchor_info(plain), (None, None)); + let salted = tc + .match_prefix(&MatchPrefixParams { + key: &vec![7, 8], + namespace: KeyNamespaceRef::new(Some("chat"), Some("tenant-a")), + }) + .best_match_node_id; + assert_eq!( + tc.prefetch_anchor_info(salted), + (Some("chat".to_string()), Some("tenant-a".to_string())) + ); + // A root anchor carries no namespace: the single root serves them all. + let root = tc.arena.root(); + assert_eq!( + tc.prefetch_anchor_info(tc.arena.node(root).id), + (None, None) + ); + // A node minted by a split inherits the namespace. + tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some("chat"), Some("tenant-a")), + ..insert_params(&vec![7], &[30]) + }); + let split_mid = tc + .match_prefix(&MatchPrefixParams { + key: &vec![7], + namespace: KeyNamespaceRef::new(Some("chat"), Some("tenant-a")), + }) + .best_match_node_id; + assert_ne!(split_mid, salted); + assert_eq!( + tc.prefetch_anchor_info(split_mid), + (Some("chat".to_string()), Some("tenant-a".to_string())) + ); +} + +#[test] +fn mamba_core_constructs_through_the_factory() { + let tc = UnifiedTreeCore::>::new( + CacheInitParams { + mamba_cache_chunk_size: Some(256), + ..CacheInitParams::default() + }, + vec![FULL, MAMBA], + ); + assert_eq!(tc.components.len(), 2); +} + +#[test] +fn mamba_sizes_and_flatten_read_the_component_state() { + let mut tc = UnifiedTreeCore::>::new( + CacheInitParams { + mamba_cache_chunk_size: Some(256), + ..CacheInitParams::default() + }, + vec![FULL, MAMBA], + ); + assert_eq!(tc.mamba_evictable_size(), 0); + assert_eq!(tc.mamba_protected_size(), 0); + assert_eq!(tc.all_mamba_values_flatten().numel(), 0); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let b = tc + .arena + .alloc_child( + a, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(a, MAMBA, Tensor::from_slice(&[7i64])); + tc.arena + .set_device_value(b, MAMBA, Tensor::from_slice(&[9i64])); + let mut slots = Vec::::try_from(tc.all_mamba_values_flatten()).unwrap(); + slots.sort_unstable(); + assert_eq!(slots, vec![7, 9]); + // A Full-only tree reports empty mamba state. + let full_only = core(); + assert_eq!(full_only.mamba_evictable_size(), 0); + assert_eq!(full_only.all_mamba_values_flatten().numel(), 0); +} + +#[test] +fn prefetch_node_accessors_cover_gate_and_hash_chain() { + let mut tc = UnifiedTreeCore::new( + CacheInitParams { + page_size: 2, + enable_hicache: true, + ..CacheInitParams::default() + }, + vec![FULL], + ); + tc.set_enable_storage(true); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + tc.insert(&insert_params(&vec![1, 2, 7, 8], &[10, 11, 12, 13])); + let child = tc + .match_prefix(&match_params(&vec![1, 2, 7, 8])) + .best_match_node_id; + + assert!(!tc.node_backuped(child)); + assert!(!tc.is_root(child)); + assert_eq!( + tc.get_last_hash_value(child).as_deref(), + Some("0bfa9b9c6fd727c7410b6d42b753439911022d34cc6ef99ac43ed7724aa48a75") + ); + assert_eq!( + tc.get_prefix_hash_values(child), + vec!["34fb5c825de7ca4aea6e712f19d439c1da0c92c37b423936c5f618545ca4fa1f".to_string()] + ); + + tc.commit_backup(child, Tensor::from_slice(&[102i64, 103]), HashMap::new()); + assert!(tc.node_backuped(child)); + + // Roots have no hashes of their own. + let root = tc.arena.root(); + assert!(tc.is_root(tc.arena.node(root).id)); + assert_eq!(tc.get_last_hash_value(tc.arena.node(root).id), None); + assert_eq!( + tc.get_prefix_hash_values(tc.arena.node(root).id), + Vec::::new() + ); +} + +#[test] +fn storage_backup_spec_is_none_for_an_unbackuped_node() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + assert!( + tc.build_storage_backup_spec(tc.arena.node(tc.arena.resolve(leaf)).id, true) + .is_none() + ); +} + +#[test] +fn storage_backup_spec_gathers_the_chained_node() { + let mut tc = UnifiedTreeCore::new( + CacheInitParams { + page_size: 2, + enable_hicache: true, + ..CacheInitParams::default() + }, + vec![FULL], + ); + tc.set_enable_storage(true); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + tc.insert(&insert_params(&vec![1, 2, 7, 8], &[10, 11, 12, 13])); + let parent = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + let child = tc + .match_prefix(&match_params(&vec![1, 2, 7, 8])) + .best_match_node_id; + tc.commit_backup(parent, Tensor::from_slice(&[100i64, 101]), HashMap::new()); + tc.commit_backup(child, Tensor::from_slice(&[102i64, 103]), HashMap::new()); + + let spec = tc + .build_storage_backup_spec(tc.arena.node(tc.arena.resolve(child)).id, true) + .unwrap(); + assert!(spec.host_value.equal(&Tensor::from_slice(&[102i64, 103]))); + assert_eq!(spec.token_ids, vec![7, 8]); + assert_eq!( + spec.hash_value, + Some(vec![ + "0bfa9b9c6fd727c7410b6d42b753439911022d34cc6ef99ac43ed7724aa48a75".to_string() + ]) + ); + assert_eq!( + spec.prefix_keys, + Some(vec![ + "34fb5c825de7ca4aea6e712f19d439c1da0c92c37b423936c5f618545ca4fa1f".to_string() + ]) + ); + assert!(spec.comp_xfers.is_empty()); + + let spec = tc + .build_storage_backup_spec(tc.arena.node(tc.arena.resolve(child)).id, false) + .unwrap(); + assert_eq!(spec.prefix_keys, None); +} + +#[test] +fn prefix_hash_walk_stops_below_an_unhashed_ancestor() { + let mut tc = core(); + let (a, b) = matched_chain(&mut tc); + tc.arena.node_mut(b).hash_value = Some(vec!["b0".to_string()]); + assert_eq!(tc.arena.node(a).hash_value, None); + assert_eq!(tc.arena.prefix_hash_values(Some(b)), vec!["b0".to_string()]); + assert_eq!(tc.arena.prefix_hash_values(None), Vec::::new()); +} + +#[test] +fn insert_host_attaches_a_host_only_leaf_under_the_root() { + let mut tc = core(); + let root = tc.arena.root(); + let result = tc.insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1, 2], + Tensor::from_slice(&[100i64, 101]), + vec!["h0".to_string(), "h1".to_string()], + ); + assert_eq!(result.prefix_len, 0); + assert_eq!(result.total_len, 2); + assert!(!result.host_insert_dropped); + let new_node = result.inserted_host_node.unwrap(); + let node = tc.arena.node(tc.arena.resolve(new_node)); + assert!(node.evicted() && node.backuped()); + assert!( + node.host_value(FULL) + .equal(&Tensor::from_slice(&[100i64, 101])) + ); + assert_eq!( + node.hash_value, + Some(vec!["h0".to_string(), "h1".to_string()]) + ); + assert!( + tc.evictable_host_leaves + .contains(tc.arena.resolve(new_node)) + ); + tc.sanity_check(&[], &[]); +} + +#[test] +fn insert_host_allows_a_suffix_under_an_unbacked_write_back_parent() { + let mut tc = core(); + tc.is_write_back = true; + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + let root = tc.arena.root(); + let result = tc.insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1, 2, 3, 4], + Tensor::from_slice(&[100i64, 101, 102, 103]), + vec!["h0", "h1", "h2", "h3"] + .into_iter() + .map(String::from) + .collect(), + ); + assert_eq!(result.prefix_len, 2); + assert_eq!(result.total_len, 4); + assert!(!result.host_insert_dropped); + let new_node = tc + .arena + .node(tc.arena.resolve(result.inserted_host_node.unwrap())); + assert!( + new_node + .host_value(FULL) + .equal(&Tensor::from_slice(&[102i64, 103])) + ); + assert_eq!( + new_node.hash_value, + Some(vec!["h2".to_string(), "h3".to_string()]) + ); +} + +#[test] +fn insert_host_drops_a_suffix_under_an_unbacked_write_through_parent() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + let root = tc.arena.root(); + let nodes_before = tc.arena.len(); + let result = tc.insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1, 2, 3, 4], + Tensor::from_slice(&[100i64, 101, 102, 103]), + vec!["h0", "h1", "h2", "h3"] + .into_iter() + .map(String::from) + .collect(), + ); + + assert_eq!(result.prefix_len, 2); + assert_eq!(result.total_len, 4); + assert_eq!(result.inserted_host_node, None); + assert!(result.host_insert_dropped); + assert!(result.cache_actions.is_empty()); + assert_eq!(tc.arena.len(), nodes_before); +} + +#[test] +fn insert_host_drop_preserves_split_actions_and_lengths() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 3])) + .best_match_node_id; + tc.mark_write_through_pending(leaf); + let root = tc.arena.root(); + let result = tc.insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1, 9], + Tensor::from_slice(&[100i64, 101]), + vec!["h0".to_string(), "h1".to_string()], + ); + + assert_eq!(result.prefix_len, 1); + assert_eq!(result.total_len, 2); + assert_eq!(result.inserted_host_node, None); + assert!(result.host_insert_dropped); + assert!(matches!( + result.cache_actions.as_slice(), + [CacheAction::ReplaceWriteThroughOnNodeSplit { ack_id, .. }] if *ack_id == leaf + )); +} + +#[test] +fn insert_host_splits_a_host_chain_and_divides_the_hash() { + let mut tc = core(); + let root = tc.arena.root(); + tc.insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1, 2, 3], + Tensor::from_slice(&[100i64, 101, 102]), + vec!["h0", "h1", "h2"] + .into_iter() + .map(String::from) + .collect(), + ); + let result = tc.insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1, 9], + Tensor::from_slice(&[200i64, 201]), + vec!["g0".to_string(), "g1".to_string()], + ); + assert_eq!(result.prefix_len, 1); + let new_node = tc + .arena + .node(tc.arena.resolve(result.inserted_host_node.unwrap())); + assert!( + new_node + .host_value(FULL) + .equal(&Tensor::from_slice(&[201i64])) + ); + assert_eq!(new_node.hash_value, Some(vec!["g1".to_string()])); + // The split divided the chain's host value and hash at the boundary. + let split_parent = new_node.parent(); + let parent = tc.arena.node(split_parent); + assert!( + parent + .host_value(FULL) + .equal(&Tensor::from_slice(&[100i64])) + ); + assert_eq!(parent.hash_value, Some(vec!["h0".to_string()])); + tc.sanity_check(&[], &[]); +} + +#[test] +fn insert_host_hash_slices_by_pages_not_atoms() { + let params = CacheInitParams { + page_size: 2, + ..Default::default() + }; + let mut tc: UnifiedTreeCore> = UnifiedTreeCore::new(params, vec![FULL]); + let root = tc.arena.root(); + tc.insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1, 2], + Tensor::from_slice(&[100i64, 101]), + vec!["h0".to_string()], + ); + let result = tc.insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1, 2, 3, 4], + Tensor::from_slice(&[200i64, 201, 202, 203]), + vec!["g0".to_string(), "g1".to_string()], + ); + assert_eq!(result.prefix_len, 2); + let new_node = tc + .arena + .node(tc.arena.resolve(result.inserted_host_node.unwrap())); + // Two matched atoms are ONE page: only g0 is consumed. + assert_eq!(new_node.hash_value, Some(vec!["g1".to_string()])); + assert!( + new_node + .host_value(FULL) + .equal(&Tensor::from_slice(&[202i64, 203])) + ); +} + +#[test] +fn insert_host_full_match_reports_only_a_backuped_node() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + let root = tc.arena.root(); + // The device-only match reports no host node. + let result = tc.insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1, 2], + Tensor::from_slice(&[100i64, 101]), + vec!["h0".to_string(), "h1".to_string()], + ); + assert_eq!(result.prefix_len, 2); + assert_eq!(result.inserted_host_node, None); + assert!(!result.host_insert_dropped); + // Once backuped, the same insert reports the node. + tc.arena.set_host_value( + tc.arena.resolve(leaf), + FULL, + Tensor::from_slice(&[20i64, 21]), + ); + let result = tc.insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1, 2], + Tensor::from_slice(&[100i64, 101]), + vec!["h0".to_string(), "h1".to_string()], + ); + assert_eq!(result.inserted_host_node, Some(leaf)); + assert!(!result.host_insert_dropped); +} + +#[test] +#[ignore] +#[should_panic(expected = "insert_host: parent")] +fn insert_host_panics_on_a_colliding_page() { + // TODO: unconstructible today — the insert_host walk (like the python one) + // follows any child on the suffix page instead of breaking at a dead node, + // so the add path never sees an occupied page; the assert is defensive-only. + let mut tc = core(); + tc.set_hicache_enabled(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + let root = tc.arena.root(); + // A host suffix colliding with the device-valued child's first page. + tc.insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![1, 9], + Tensor::from_slice(&[100i64, 101]), + vec!["h0".to_string(), "h1".to_string()], + ); +} + +#[test] +fn insert_host_empty_key_is_a_noop() { + let mut tc = core(); + let root = tc.arena.root(); + let result = tc.insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + vec![], + Tensor::from_slice(&[0i64; 0]), + vec![], + ); + assert_eq!(result.prefix_len, 0); + assert!(result.mamba_exist); + assert_eq!(result.inserted_host_node, None); + assert!(!result.host_insert_dropped); + assert_eq!(tc.arena.len(), 1); +} + +#[test] +fn commit_backup_attaches_the_host_value() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + tc.commit_backup(leaf, Tensor::from_slice(&[100i64, 101]), HashMap::new()); + let node = tc.arena.node(tc.arena.resolve(leaf)); + assert!(node.backuped()); + assert!( + node.host_value(FULL) + .equal(&Tensor::from_slice(&[100i64, 101])) + ); +} + +#[test] +fn build_backup_spec_reads_the_device_value() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + let (device_value, comp_xfers) = tc.build_backup_spec(leaf); + assert!(device_value.equal(&Tensor::from_slice(&[10i64, 11]))); + assert!(comp_xfers.is_empty()); +} + +#[test] +fn build_backup_spec_skips_full_kv_for_an_already_backuped_node() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + + tc.commit_backup(leaf, Tensor::from_slice(&[100i64, 101]), HashMap::new()); + assert!(tc.arena.node(tc.arena.resolve(leaf)).backuped()); + + let (device_value, comp_xfers) = tc.build_backup_spec(leaf); + assert_eq!(device_value.numel(), 0); + assert!(comp_xfers.is_empty()); +} + +#[test] +fn commit_backup_preserves_full_kv_when_host_indices_are_empty() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + + tc.commit_backup(leaf, Tensor::from_slice(&[100i64, 101]), HashMap::new()); + tc.commit_backup(leaf, Tensor::from_slice(&[] as &[i64]), HashMap::new()); + + let node = tc.arena.node(tc.arena.resolve(leaf)); + assert!( + node.host_value(FULL) + .equal(&Tensor::from_slice(&[100i64, 101])) + ); +} + +// Two backuped device nodes [1,2] -> [3,4]; returns (parent, child) ids. +fn backuped_chain(tc: &mut UnifiedTreeCore>) -> (NodeIdx_, NodeIdx_) { + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + tc.insert(&insert_params(&vec![1, 2, 3, 4], &[10, 11, 12, 13])); + let parent = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + let child = tc + .match_prefix(&match_params(&vec![1, 2, 3, 4])) + .best_match_node_id; + tc.commit_backup(parent, Tensor::from_slice(&[20i64, 21]), HashMap::new()); + tc.commit_backup(child, Tensor::from_slice(&[22i64, 23]), HashMap::new()); + (tc.arena.resolve(parent), tc.arena.resolve(child)) +} + +// Demote `node_id` (device release of a backuped node), discarding the frees. +fn demote_node(tc: &mut UnifiedTreeCore>, node_id: NodeIdx_) { + tc.demote(tc.arena.node(node_id).id); +} + +#[test] +fn demote_releases_the_device_value_and_keeps_the_host_copy() { + let mut tc = core(); + let (_parent, child) = backuped_chain(&mut tc); + demote_node(&mut tc, child); + let node = tc.arena.node(child); + assert!(node.evicted() && node.backuped()); + assert_eq!(tc.full_evictable_size(), 2); + tc.sanity_check(&[], &[]); +} + +#[test] +fn build_load_back_spec_collects_the_evicted_chain_ancestors_first() { + let mut tc = core(); + let (parent, child) = backuped_chain(&mut tc); + demote_node(&mut tc, child); + demote_node(&mut tc, parent); + let (kv_xfer, comp_xfers) = + tc.build_load_back_spec(tc.arena.node(child).id, /* req = */ None); + assert_eq!(kv_xfer.name, PoolName::Kv); + assert!( + kv_xfer + .host_indices + .unwrap() + .equal(&Tensor::from_slice(&[20i64, 21, 22, 23])) + ); + assert!(kv_xfer.device_indices.is_none()); + assert_eq!( + kv_xfer.nodes_to_load, + Some(vec![tc.arena.node(parent).id, tc.arena.node(child).id]) + ); + assert!(comp_xfers.is_empty()); +} + +#[test] +fn build_load_back_spec_returns_an_empty_transfer_for_a_device_backed_node() { + let mut tc = core(); + let (_parent, child) = backuped_chain(&mut tc); + let (kv_xfer, comp_xfers) = + tc.build_load_back_spec(tc.arena.node(child).id, /* req = */ None); + let host_indices = kv_xfer.host_indices.unwrap(); + assert_eq!(host_indices.numel(), 0); + assert_eq!(host_indices.kind(), Kind::Int64); + assert_eq!(kv_xfer.nodes_to_load, Some(vec![])); + assert!(comp_xfers.is_empty()); +} + +#[test] +fn commit_load_back_reattaches_device_slices_and_restores_the_match() { + let mut tc = core(); + let (parent, child) = backuped_chain(&mut tc); + demote_node(&mut tc, child); + demote_node(&mut tc, parent); + // Coexistence tracking is intentionally lazy and may retain stale entries + // after demotion. Remove them so this test observes the ack-time refresh. + tc.full_coexisting_host_nodes.discard(parent); + tc.full_coexisting_host_nodes.discard(child); + let (kv_xfer, comp_xfers) = + tc.build_load_back_spec(tc.arena.node(child).id, /* req = */ None); + let actions = tc.commit_load_back( + tc.arena.node(child).id, + Tensor::from_slice(&[50i64, 51, 52, 53]), + kv_xfer, + comp_xfers, + ); + assert!(actions.is_empty()); + assert!( + tc.arena + .device_value(parent, FULL) + .equal(&Tensor::from_slice(&[50i64, 51])) + ); + assert!( + tc.arena + .device_value(child, FULL) + .equal(&Tensor::from_slice(&[52i64, 53])) + ); + // Write-through does not need an in-flight Full host pin. Duplicate + // tracking is refreshed only when the orchestrator acknowledges the load. + assert!(!tc.arena.node(parent).is_load_back_pending()); + assert!(!tc.arena.node(child).is_load_back_pending()); + assert!(!tc.full_coexisting_host_nodes.contains(parent)); + assert!(!tc.full_coexisting_host_nodes.contains(child)); + assert_eq!(tc.full_evictable_size(), 4); + // The orchestrator re-locks the loaded path right after commit; that lock walk + // also re-evaluates the parent's transient D-leaf membership. + tc.inc_lock_ref(tc.arena.node(child).id); + tc.dec_lock_ref( + tc.arena.node(child).id, + /* params = */ None, + /* skip_swa = */ false, + ); + tc.finish_load_back(tc.arena.node(child).id); + assert!(tc.full_coexisting_host_nodes.contains(parent)); + assert!(tc.full_coexisting_host_nodes.contains(child)); + tc.sanity_check(&[], &[]); + let result = tc.match_prefix(&match_params(&vec![1, 2, 3, 4])); + assert!( + result + .device_indices + .equal(&Tensor::from_slice(&[50i64, 51, 52, 53])) + ); +} + +#[test] +fn device_eviction_and_demote_skip_a_load_back_pinned_chain() { + let mut tc = UnifiedTreeCore::new( + CacheInitParams { + is_write_back: true, + ..CacheInitParams::default() + }, + vec![FULL], + ); + let (parent, child) = backuped_chain(&mut tc); + demote_node(&mut tc, child); + demote_node(&mut tc, parent); + tc.full_coexisting_host_nodes.discard(parent); + tc.full_coexisting_host_nodes.discard(child); + let (kv_xfer, comp_xfers) = + tc.build_load_back_spec(tc.arena.node(child).id, /* req = */ None); + tc.commit_load_back( + tc.arena.node(child).id, + Tensor::from_slice(&[50i64, 51, 52, 53]), + kv_xfer, + comp_xfers, + ); + let anchor_id = tc.arena.node(child).id; + assert_eq!(tc.arena.node(parent).load_back_pending_id, Some(anchor_id)); + assert_eq!(tc.arena.node(child).load_back_pending_id, Some(anchor_id)); + assert!(!tc.full_coexisting_host_nodes.contains(parent)); + assert!(!tc.full_coexisting_host_nodes.contains(child)); + // The pin alone keeps the in-flight chain out of device eviction. + tc.evict_device_start(FULL, 4); + let (next, _) = tc.evict_device_next_node(FULL, &HashMap::new()); + assert_eq!(next, None); + tc.evict_device_end(FULL); + tc.demote(tc.arena.node(child).id); + assert!(tc.arena.has_device_value(child, FULL)); + tc.finish_load_back(tc.arena.node(child).id); + assert!(!tc.arena.node(parent).is_load_back_pending()); + assert!(!tc.arena.node(child).is_load_back_pending()); + assert!(tc.full_coexisting_host_nodes.contains(parent)); + assert!(tc.full_coexisting_host_nodes.contains(child)); + // The ack re-arms leaf-set membership and demotion. + tc.evict_device_start(FULL, 4); + let (next, _) = tc.evict_device_next_node(FULL, &HashMap::new()); + assert_eq!(next, Some(tc.arena.node(child).id)); + tc.evict_device_end(FULL); + tc.demote(tc.arena.node(child).id); + assert!(!tc.arena.has_device_value(child, FULL)); + tc.sanity_check(&[], &[]); +} + +#[test] +fn component_has_host_value_only_tracks_the_demote_and_load_back_cycle() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1], &[10])); + let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; + assert!(!tc.component_has_host_value_only(leaf, FULL)); + tc.commit_backup(leaf, Tensor::from_slice(&[20i64]), HashMap::new()); + // Device value still present: backuped but not host-only. + assert!(!tc.component_has_host_value_only(leaf, FULL)); + let leaf_idx = tc.arena.resolve(leaf); + demote_node(&mut tc, leaf_idx); + assert!(tc.component_has_host_value_only(leaf, FULL)); + let (kv_xfer, comp_xfers) = tc.build_load_back_spec(leaf, /* req = */ None); + tc.commit_load_back(leaf, Tensor::from_slice(&[30i64]), kv_xfer, comp_xfers); + assert!(!tc.component_has_host_value_only(leaf, FULL)); + tc.finish_load_back(leaf); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "!node.evicted() && node.backuped()")] +fn demote_panics_on_an_unbackuped_node() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1], &[10])); + let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; + tc.demote(leaf); +} + +#[test] +fn try_demote_rejects_unbackuped_and_evicted_nodes() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1], &[10])); + let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; + + assert!(matches!( + tc.try_demote(leaf), + Err(TreeCoreRuntimeError::InvalidDemoteState { + node_id, + evicted: false, + backuped: false, + }) if node_id == leaf + )); + + tc.commit_backup(leaf, Tensor::from_slice(&[20i64]), HashMap::new()); + tc.demote(leaf); + assert!(matches!( + tc.try_demote(leaf), + Err(TreeCoreRuntimeError::InvalidDemoteState { + node_id, + evicted: true, + backuped: true, + }) if node_id == leaf + )); +} + +#[test] +fn fallible_node_boundaries_reject_stale_handles() { + let mut tc = core(); + let stale_root = tc.root_node_handle(/* extra_key = */ None); + tc.reset(); + + assert!(matches!( + tc.try_demote(stale_root), + Err(TreeCoreRuntimeError::NodeNotAllocated { node_id }) if node_id == stale_root + )); + assert!(matches!( + tc.try_build_hicache_transfers( + FULL, + stale_root, + CacheTransferPhase::BackupStorage, + /* host_indices = */ None, + /* token_ids = */ None, + /* prefetch_tokens = */ 0, + /* last_hash = */ None, + ), + Err(TreeCoreRuntimeError::NodeNotAllocated { node_id }) if node_id == stale_root + )); + assert!(matches!( + tc.try_build_load_back_spec(stale_root, /* req = */ None), + Err(TreeCoreRuntimeError::NodeNotAllocated { node_id }) if node_id == stale_root + )); + assert!(matches!( + tc.try_get_hash_values(stale_root), + Err(TreeCoreRuntimeError::NodeNotAllocated { node_id }) if node_id == stale_root + )); + assert!(matches!( + tc.try_dfs_weight_order(&[stale_root]), + Err(TreeCoreRuntimeError::NodeNotAllocated { node_id }) if node_id == stale_root + )); + + let live_root = tc.root_node_handle(/* extra_key = */ None); + assert!(tc.is_root(live_root)); +} + +#[test] +fn match_prefix_with_hicache_splits_a_host_only_backuped_node() { + let mut tc = core(); + tc.set_hicache_enabled(); + tc.insert(&insert_params(&vec![1, 2, 3, 4], &[10, 11, 12, 13])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 3, 4])) + .best_match_node_id; + tc.commit_backup( + leaf, + Tensor::from_slice(&[100i64, 101, 102, 103]), + HashMap::new(), + ); + tc.demote(leaf); + // The partial match splits the host-only node; the host prefix stays usable. + let result = tc.match_prefix(&match_params(&vec![1, 2, 9])); + assert_eq!(result.device_indices.numel(), 0); + let root = tc.arena.root(); + assert_eq!(result.last_device_node_id, tc.arena.node(root).id); + assert_eq!(result.host_hit_length, 2); + assert_eq!(result.best_match_node_id, result.last_host_node_id); + let parent = tc.arena.resolve(result.best_match_node_id); + let child = tc.arena.node(parent).children[&(KeyNamespace::default(), vec![3])]; + { + let parent_node = tc.arena.node(parent); + assert_eq!(parent_node.key, vec![1, 2]); + assert!(parent_node.evicted() && parent_node.backuped()); + assert!( + parent_node + .host_value(FULL) + .equal(&Tensor::from_slice(&[100i64, 101])) + ); + } + let child_node = tc.arena.node(child); + assert_eq!(child_node.key, vec![3, 4]); + assert!(child_node.evicted() && child_node.backuped()); + assert!( + child_node + .host_value(FULL) + .equal(&Tensor::from_slice(&[102i64, 103])) + ); + tc.sanity_check(&[], &[]); +} + +#[test] +fn mixed_backup_evict_insert_keeps_the_leaf_sets_disjoint() { + let mut tc = core(); + tc.set_hicache_enabled(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + tc.insert(&insert_params(&vec![101, 102], &[20, 21])); + tc.insert(&insert_params(&vec![201, 202], &[30, 31])); + tc.insert(&insert_params(&vec![301, 302], &[40, 41])); + tc.insert(&insert_params(&vec![401, 402], &[50, 51])); + // Backing up (and thereby re-stamping) the first three chains leaves the + // two unbacked chains as the oldest eviction victims. + let first = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + tc.commit_backup(first, Tensor::from_slice(&[100i64, 101]), HashMap::new()); + let second = tc + .match_prefix(&match_params(&vec![101, 102])) + .best_match_node_id; + tc.commit_backup(second, Tensor::from_slice(&[102i64, 103]), HashMap::new()); + let third = tc + .match_prefix(&match_params(&vec![201, 202])) + .best_match_node_id; + tc.commit_backup(third, Tensor::from_slice(&[104i64, 105]), HashMap::new()); + let mut tracker = HashMap::from([(FULL, 0)]); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.evict_device_start(FULL, /* request_cnt = */ 4); + loop { + let (leaf, step) = tc.evict_device_next_node(FULL, &tracker); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + let Some(leaf) = leaf else { break }; + let (_, step) = tc.evict_device_leaf(leaf, /* is_write_back = */ false); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + } + tc.evict_device_end(FULL); + assert_eq!(tracker[&FULL], 4); + // The unbacked chains died outright. + assert_eq!( + tc.match_prefix(&match_params(&vec![301, 302])) + .device_indices + .numel(), + 0 + ); + assert_eq!( + tc.match_prefix(&match_params(&vec![401, 402])) + .device_indices + .numel(), + 0 + ); + tc.insert(&insert_params(&vec![501, 502], &[60, 61])); + tc.insert(&insert_params(&vec![601, 602], &[70, 71])); + tc.insert(&insert_params(&vec![701, 702], &[80, 81])); + // D-leaf / H-leaf membership stays mutually exclusive after the mixed traffic. + for node in tc.collect_all_nodes_() { + assert!( + !(tc.evictable_device_leaves.contains(node) && tc.evictable_host_leaves.contains(node)) + ); + } + tc.sanity_check(&[], &[]); +} + +fn write_back_core() -> UnifiedTreeCore> { + UnifiedTreeCore::new( + CacheInitParams { + is_write_back: true, + ..Default::default() + }, + vec![FULL], + ) +} + +// A device-on unbacked leaf [1,2] with a host-only child [3,4] under it. +fn unbacked_leaf_with_host_child(tc: &mut UnifiedTreeCore>) -> (NodeIdx_, NodeIdx_) { + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + let child = tc + .arena + .alloc_child( + tc.arena.resolve(leaf), + /* key = */ vec![3, 4], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(child, FULL, Tensor::from_slice(&[20i64, 21])); + tc.update_evictable_leaf_sets_(child); + tc.update_evictable_leaf_sets_(tc.arena.resolve(leaf)); + (tc.arena.resolve(leaf), child) +} + +#[test] +fn drop_subtree_no_host_frees_the_leaf_and_its_host_descendants() { + let mut tc = write_back_core(); + let (leaf, _child) = unbacked_leaf_with_host_child(&mut tc); + let mut tracker = HashMap::from([(FULL, 0)]); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + let (dropped, step) = tc.drop_subtree_no_host(tc.arena.node(leaf).id); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + assert!(dropped); + // Under EvictLayer::All only device tokens enter the tracker; host + // frees ride the host_frees tensors. + assert_eq!(tracker[&FULL], 2); + assert_eq!(df[&FULL].len(), 1); + assert!(df[&FULL][0].equal(&Tensor::from_slice(&[10i64, 11]))); + assert_eq!(hf[&FULL].len(), 1); + assert!(hf[&FULL][0].equal(&Tensor::from_slice(&[20i64, 21]))); + assert_eq!(tc.arena.len(), 1); + let result = tc.match_prefix(&match_params(&vec![1, 2])); + assert_eq!(result.device_indices.numel(), 0); + tc.sanity_check(&[], &[]); +} + +#[test] +fn drop_subtree_no_host_removes_a_deeper_host_chain_child_first() { + let mut tc = write_back_core(); + let (leaf, child) = unbacked_leaf_with_host_child(&mut tc); + let grandchild = tc + .arena + .alloc_child( + child, + /* key = */ vec![5], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(grandchild, FULL, Tensor::from_slice(&[22i64])); + tc.update_evictable_leaf_sets_(grandchild); + tc.update_evictable_leaf_sets_(child); + let mut tracker = HashMap::from([(FULL, 0)]); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + let (dropped, step) = tc.drop_subtree_no_host(tc.arena.node(leaf).id); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + assert!(dropped); + assert_eq!(tracker[&FULL], 2); + assert_eq!(hf[&FULL].len(), 2); + assert_eq!(tc.arena.len(), 1); + tc.sanity_check(&[], &[]); +} + +#[test] +fn drop_subtree_no_host_bails_on_a_locked_descendant() { + let mut tc = write_back_core(); + let (leaf, child) = unbacked_leaf_with_host_child(&mut tc); + tc.arena + .node_mut(child) + .set_lock_ref_(ValueSlotIdx::host(FULL), 1); + let mut tracker = HashMap::from([(FULL, 0)]); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + let (dropped, step) = tc.drop_subtree_no_host(tc.arena.node(leaf).id); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + assert!(!dropped); + assert_eq!(tracker[&FULL], 0); + assert!(df.is_empty() && hf.is_empty()); + assert_eq!(tc.arena.len(), 3); +} + +#[test] +fn drop_subtree_no_host_bails_on_a_host_locked_root() { + let mut tc = write_back_core(); + let (leaf, _child) = unbacked_leaf_with_host_child(&mut tc); + tc.arena + .node_mut(leaf) + .set_lock_ref_(ValueSlotIdx::host(FULL), 1); + let (dropped, _step) = tc.drop_subtree_no_host(tc.arena.node(leaf).id); + assert!(!dropped); + assert_eq!(tc.arena.len(), 3); +} + +#[test] +#[should_panic(expected = "is not a D-leaf")] +fn drop_subtree_no_host_panics_on_a_non_device_leaf() { + let mut tc = write_back_core(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + tc.insert(&insert_params(&vec![1, 2, 3, 4], &[10, 11, 12, 13])); + let parent = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + tc.drop_subtree_no_host(parent); +} + +#[test] +#[should_panic(expected = "assertion failed")] +fn drop_subtree_no_host_panics_on_a_backuped_leaf() { + let mut tc = write_back_core(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + tc.commit_backup(leaf, Tensor::from_slice(&[20i64, 21]), HashMap::new()); + tc.drop_subtree_no_host(leaf); +} + +#[test] +fn write_back_eviction_frees_the_device_value_exactly_once() { + // The backup pass must not free the device value the DMA still reads; + // the post-ack pass demotes and frees it exactly once. + let mut tc = write_back_core(); + tc.insert(&insert_params(&vec![1], &[10])); + let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; + let mut tracker = HashMap::from([(FULL, 0)]); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + let (action, step) = tc.evict_device_leaf(leaf, true); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + assert_eq!(action.unwrap().node_ids, vec![leaf]); + assert!(df.is_empty() && hf.is_empty()); + tc.commit_backup(leaf, Tensor::from_slice(&[20i64]), HashMap::new()); + let (action, step) = tc.evict_device_leaf(leaf, true); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + assert!(action.is_none()); + assert_eq!(tracker[&FULL], 1); + assert_eq!(df[&FULL].len(), 1); + assert!(df[&FULL][0].equal(&Tensor::from_slice(&[10i64]))); + assert!( + tc.arena.node(tc.arena.resolve(leaf)).evicted() + && tc.arena.node(tc.arena.resolve(leaf)).backuped() + ); + tc.sanity_check(&[], &[]); +} + +#[test] +fn insert_empty_key_is_a_noop_and_mints_no_namespace_root() { + let mut tc = core(); + let result = tc.insert(&insert_params(&vec![], &[])); + assert_eq!(result.prefix_len, 0); + // Vacuously true: there is no sequence whose mamba state could be missing. + assert!(result.mamba_exist); + assert_eq!(tc.arena.len(), 1); + // A namespaced empty insert never creates the namespace root. + let result = tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some("ghost"), None), + ..insert_params(&vec![], &[]) + }); + assert!(result.mamba_exist); + assert_eq!(tc.arena.len(), 1); + assert!(!tc.arena.namespace_exists(Some("ghost"))); +} + +#[test] +fn root_node_handle_is_namespace_independent() { + let mut tc = core(); + let root_handle = tc.arena.node(tc.arena.root()).id; + assert_eq!(tc.root_node_handle(None), root_handle); + // The single root serves every namespace, seen or not. + assert_eq!(tc.root_node_handle(Some("ghost")), root_handle); + assert!(!tc.arena.namespace_exists(Some("ghost"))); + tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some("chat"), None), + ..insert_params(&vec![1], &[10]) + }); + assert_eq!(tc.root_node_handle(Some("chat")), root_handle); +} + +#[test] +fn dfs_weight_order_groups_the_heaviest_subtree_first() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 10], &[10, 11])); + tc.insert(&insert_params(&vec![1, 11], &[10, 12])); + tc.insert(&insert_params(&vec![2, 20], &[20, 21])); + + let branch_a = tc + .match_prefix(&MatchPrefixParams { + key: &vec![1, 99], + namespace: KeyNamespaceRef::default(), + }) + .last_device_node_id; + let leaf_a1 = tc + .match_prefix(&MatchPrefixParams { + key: &vec![1, 10], + namespace: KeyNamespaceRef::default(), + }) + .last_device_node_id; + let leaf_a2 = tc + .match_prefix(&MatchPrefixParams { + key: &vec![1, 11], + namespace: KeyNamespaceRef::default(), + }) + .last_device_node_id; + let leaf_b = tc + .match_prefix(&MatchPrefixParams { + key: &vec![2, 20], + namespace: KeyNamespaceRef::default(), + }) + .last_device_node_id; + + assert_eq!( + tc.dfs_weight_order(&[leaf_b, leaf_a2, leaf_a1, leaf_a1, branch_a]), + vec![2, 3, 1, 4, 0] + ); + assert_eq!(tc.dfs_weight_order(&[leaf_b, leaf_a2]), vec![1, 0]); +} + +#[test] +fn get_hash_values_reads_the_nodes_own_hashes() { + let mut tc = core(); + let (a, _b) = matched_chain(&mut tc); + assert_eq!( + tc.get_hash_values(tc.arena.node(a).id), + Vec::::new() + ); + tc.arena.node_mut(a).hash_value = Some(vec!["h0".to_string(), "h1".to_string()]); + assert_eq!( + tc.get_hash_values(tc.arena.node(a).id), + vec!["h0".to_string(), "h1".to_string()] + ); +} + +#[test] +fn insert_empty_key_still_touches_the_existing_root() { + let mut tc = core(); + tc.insert(&InsertParams { + priority: 7, + ..insert_params(&vec![], &[]) + }); + assert_eq!(tc.arena.node(tc.arena.root()).priority, 7); + // A namespaced empty insert touches the same single root. + tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some("chat"), None), + priority: 9, + ..insert_params(&vec![], &[]) + }); + assert_eq!(tc.arena.node(tc.arena.root()).priority, 9); +} + +#[test] +fn insert_into_a_named_namespace_is_isolated() { + let mut tc = core(); + tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some("lora-1"), None), + ..insert_params(&vec![1, 2], &[10, 11]) + }); + // The default namespace stays empty; the named namespace hits. + let miss = tc.match_prefix(&match_params(&vec![1, 2])); + assert_eq!(miss.device_indices.numel(), 0); + let hit = tc.match_prefix(&MatchPrefixParams { + key: &vec![1, 2], + namespace: KeyNamespaceRef::new(Some("lora-1"), None), + }); + assert!(hit.device_indices.equal(&Tensor::from_slice(&[10i64, 11]))); +} + +fn page2_core() -> UnifiedTreeCore> { + UnifiedTreeCore::new( + CacheInitParams { + page_size: 2, + ..Default::default() + }, + vec![FULL], + ) +} + +#[test] +fn insert_sub_page_key_is_a_noop_and_mints_no_namespace_root() { + let mut tc = page2_core(); + let result = tc.insert(&insert_params(&vec![1], &[10])); + assert_eq!(result.prefix_len, 0); + assert!(result.mamba_exist); + assert_eq!(tc.arena.len(), 1); + let result = tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some("ghost"), None), + ..insert_params(&vec![1], &[10]) + }); + assert!(result.mamba_exist); + assert_eq!(tc.arena.len(), 1); + assert!(!tc.arena.namespace_exists(Some("ghost"))); +} + +#[test] +fn insert_page_size_two_drops_the_unaligned_tail() { + let mut tc = page2_core(); + let result = tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + assert_eq!(result.prefix_len, 0); + assert_eq!(tc.evictable_size_(FULL), 2); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + assert_eq!(tc.arena.node(tc.arena.resolve(leaf)).key, vec![1, 2]); + assert!( + tc.arena + .device_value(tc.arena.resolve(leaf), FULL) + .equal(&Tensor::from_slice(&[10i64, 11])) + ); + tc.sanity_check(&[], &[]); +} + +#[test] +fn insert_page_size_two_splits_mid_page_divergence_at_the_page_boundary() { + let mut tc = page2_core(); + tc.insert(&insert_params(&vec![1, 2, 3, 4], &[10, 11, 12, 13])); + // The keys share 3 atoms but pages quantize the overlap down to 2. + let result = tc.insert(&insert_params(&vec![1, 2, 3, 9], &[20, 21, 22, 29])); + assert_eq!(result.prefix_len, 2); + let [CacheAction::FreeDeviceKV(freed)] = result.cache_actions.as_slice() else { + panic!( + "expected one FreeDeviceKV action, got {:?}", + action_kinds(&result.cache_actions) + ); + }; + assert_eq!(freed.len(), 1); + assert!(freed[0].equal(&Tensor::from_slice(&[20i64, 21]))); + let prefix = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + assert_eq!(tc.arena.node(tc.arena.resolve(prefix)).key, vec![1, 2]); + let matched = tc.match_prefix(&match_params(&vec![1, 2, 3, 4])); + assert!( + matched + .device_indices + .equal(&Tensor::from_slice(&[10i64, 11, 12, 13])) + ); + let matched = tc.match_prefix(&match_params(&vec![1, 2, 3, 9])); + assert!( + matched + .device_indices + .equal(&Tensor::from_slice(&[10i64, 11, 22, 29])) + ); + tc.sanity_check(&[], &[]); +} + +#[test] +fn match_prefix_page_size_two_splits_at_a_page_boundary() { + let mut tc = page2_core(); + tc.insert(&insert_params(&vec![1, 2, 3, 4], &[10, 11, 12, 13])); + // The query shares 3 atoms; pages quantize the split down to 2. + let result = tc.match_prefix(&match_params(&vec![1, 2, 3, 9])); + assert!( + result + .device_indices + .equal(&Tensor::from_slice(&[10i64, 11])) + ); + assert_eq!( + tc.arena + .node(tc.arena.resolve(result.best_match_node_id)) + .key, + vec![1, 2] + ); + // The split child stays reachable through its own page key. + let matched = tc.match_prefix(&match_params(&vec![1, 2, 3, 4])); + assert!( + matched + .device_indices + .equal(&Tensor::from_slice(&[10i64, 11, 12, 13])) + ); + tc.sanity_check(&[], &[]); +} + +#[test] +fn match_prefix_page_size_two_sub_page_query_is_an_empty_match() { + let mut tc = page2_core(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + let root = tc.arena.root(); + let result = tc.match_prefix(&match_params(&vec![1])); + assert_eq!(result.device_indices.numel(), 0); + assert_eq!(result.best_match_node_id, tc.arena.node(root).id); +} + +#[test] +fn evict_walk_page_size_two_empties_the_tree() { + let mut tc = page2_core(); + tc.insert(&insert_params(&vec![1, 2, 3, 4], &[10, 11, 12, 13])); + tc.insert(&insert_params(&vec![5, 6], &[14, 15])); + let mut tracker = HashMap::from([(FULL, 0)]); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.evict_device_start(FULL, /* request_cnt = */ 100); + let mut evicted = 0; + loop { + let (leaf, step) = tc.evict_device_next_node(FULL, &tracker); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + let Some(leaf) = leaf else { break }; + let (_, step) = tc.evict_device_leaf(leaf, /* is_write_back = */ false); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + evicted += 1; + } + tc.evict_device_end(FULL); + assert_eq!(evicted, 2); + assert_eq!(tracker[&FULL], 6); + assert_eq!(tc.arena.len(), 1); + tc.sanity_check(&[], &[]); +} + +#[test] +fn evict_and_detach_frees_the_device_value_and_tracks() { + let mut tc = core(); + let root = tc.arena.root(); + let a = tc.add_new_node_( + root, + /* key = */ vec![1, 2], + &Tensor::from_slice(&[10, 11]), + /* priority = */ 0, + /* extra_key = */ None, + ); + // A mechanical LRU entry pins the detach branch (Full itself never lists). + tc.device_lru_list_mut(FULL).insert_mru(a); + let mut tracker = HashMap::from([(FULL, 0)]); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + let (device_freed, host_freed) = tc.evict_component_and_detach_lru_( + a, + FULL, + &mut df, + &mut hf, + EvictLayer::Device, + Some(&mut tracker), + ); + assert_eq!((device_freed, host_freed), (2, 0)); + assert_eq!(tracker[&FULL], 2); + assert_eq!(df[&FULL].len(), 1); + assert!(!tc.device_lru_list(FULL).in_list(Some(a))); +} + +#[test] +fn evict_and_detach_host_targets_the_host_tier() { + let mut tc = core(); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(a, FULL, Tensor::from_slice(&[20i64, 21])); + let mut tracker = HashMap::from([(FULL, 0)]); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + let (device_freed, host_freed) = tc.evict_component_and_detach_lru_( + a, + FULL, + &mut df, + &mut hf, + EvictLayer::Host, + Some(&mut tracker), + ); + assert_eq!((device_freed, host_freed), (0, 2)); + assert_eq!(tracker[&FULL], 2); + assert!(df.is_empty()); + assert_eq!(hf[&FULL].len(), 1); +} + +#[test] +fn remove_leaf_from_parent_unlinks_and_recycles() { + let mut tc = core(); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.remove_leaf_from_parent_(a); + assert!(tc.arena.node(root).children.is_empty()); + assert_eq!(tc.arena.len(), 1); +} + +#[test] +#[should_panic(expected = "remove_leaf_from_parent_: a deletable leaf")] +fn remove_leaf_from_parent_panics_on_an_internal_node() { + let mut tc = core(); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .alloc_child( + a, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.remove_leaf_from_parent_(a); +} + +// Chain root -> a -> b where a is a valueless tombstone; b's deletion +// hands the walk a's id. +fn tombstone_chain(tc: &mut UnifiedTreeCore>) -> NodeIdx_ { + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let b = tc + .arena + .alloc_child( + a, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.remove_leaf_from_parent_(b); + a +} + +fn delete_walk(tc: &mut UnifiedTreeCore>, from: NodeIdx_) { + let mut tracker = HashMap::from([(FULL, 0)]); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.iteratively_delete_tombstone_leaf_(from, &mut tracker, &mut df, &mut hf); +} + +#[test] +fn tombstone_walk_deletes_childless_valueless_ancestors() { + let mut tc = core(); + let a = tombstone_chain(&mut tc); + delete_walk(&mut tc, a); + assert_eq!(tc.arena.len(), 1); + let root = tc.arena.root(); + assert!(tc.arena.node(root).children.is_empty()); +} + +#[test] +fn tombstone_walk_keeps_a_device_valued_ancestor_as_a_leaf() { + let mut tc = core(); + let a = tombstone_chain(&mut tc); + tc.arena + .set_device_value(a, FULL, Tensor::from_slice(&[10i64])); + delete_walk(&mut tc, a); + assert_eq!(tc.arena.len(), 2); + assert!(tc.evictable_device_leaves.contains(a)); +} + +#[test] +fn tombstone_walk_keeps_a_host_backed_ancestor_as_an_h_leaf() { + let mut tc = core(); + let a = tombstone_chain(&mut tc); + tc.arena + .set_host_value(a, FULL, Tensor::from_slice(&[10i64])); + delete_walk(&mut tc, a); + assert_eq!(tc.arena.len(), 2); + assert!(tc.evictable_host_leaves.contains(a)); + assert!(!tc.evictable_device_leaves.contains(a)); +} + +#[test] +fn tombstone_walk_stops_at_a_locked_ancestor() { + let mut tc = core(); + let a = tombstone_chain(&mut tc); + tc.arena + .node_mut(a) + .set_lock_ref_(ValueSlotIdx::device(FULL), 1); + delete_walk(&mut tc, a); + assert_eq!(tc.arena.len(), 2); +} + +#[test] +fn tombstone_walk_cascades_multiple_levels() { + let mut tc = core(); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let b = tc + .arena + .alloc_child( + a, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let c = tc + .arena + .alloc_child( + b, + /* key = */ vec![3], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.remove_leaf_from_parent_(c); + delete_walk(&mut tc, b); + assert_eq!(tc.arena.len(), 1); +} + +#[test] +fn tombstone_walk_stops_at_a_host_locked_ancestor() { + // A host lock (e.g. an in-flight write-back) pins a valueless ancestor. + let mut tc = core(); + let a = tombstone_chain(&mut tc); + tc.arena + .node_mut(a) + .set_lock_ref_(ValueSlotIdx::host(FULL), 1); + delete_walk(&mut tc, a); + assert_eq!(tc.arena.len(), 2); +} + +#[test] +fn tombstone_walk_sweeps_orphaned_aux_device_data() { + let mut tc = core(); + let recorder = Arc::new(SwaEvictionComponentForTest::new( + /* leaf_priority = */ 0, /* internal_priority = */ 2, + )); + tc.register_component_(recorder.clone()); + let a = tombstone_chain(&mut tc); + tc.arena.node_mut(a).values[SWA.idx()].value = Some(Tensor::from_slice(&[0i64, 1])); + tc.device_lru_list_mut(SWA).insert_mru(a); + let mut tracker = HashMap::new(); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.iteratively_delete_tombstone_leaf_(a, &mut tracker, &mut df, &mut hf); + // The orphaned Swa device data is evicted before the node is deleted. + assert_eq!(tc.arena.len(), 1); + assert_eq!( + *recorder.evictions.lock().unwrap(), + vec![(a, EvictLayer::Device)] + ); + assert_eq!(tracker[&SWA], 2); + assert_eq!(df[&SWA].len(), 1); + assert!(df[&SWA][0].equal(&Tensor::from_slice(&[0i64, 1]))); + assert!(!tc.device_lru_list(SWA).in_list(Some(a))); +} + +#[test] +fn tombstone_walk_sweeps_orphaned_aux_host_data() { + let mut tc = core(); + let recorder = Arc::new(SwaEvictionComponentForTest::new( + /* leaf_priority = */ 0, /* internal_priority = */ 2, + )); + tc.register_component_(recorder.clone()); + let a = tombstone_chain(&mut tc); + tc.arena + .node_mut(a) + .state_mut_(ValueSlotIdx::host(SWA)) + .value = Some(Tensor::from_slice(&[5i64])); + tc.host_lru_list_mut(SWA).insert_mru(a); + let mut tracker = HashMap::new(); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.iteratively_delete_tombstone_leaf_(a, &mut tracker, &mut df, &mut hf); + // The orphaned Swa host data is evicted before the node is deleted. + assert_eq!(tc.arena.len(), 1); + assert_eq!( + *recorder.evictions.lock().unwrap(), + vec![(a, EvictLayer::Host)] + ); + assert_eq!(tracker[&SWA], 1); + assert_eq!(hf[&SWA].len(), 1); + assert!(hf[&SWA][0].equal(&Tensor::from_slice(&[5i64]))); + assert!(!tc.host_lru_list(SWA).in_list(Some(a))); +} + +#[test] +fn cascade_tombstones_full_after_the_component_sweep() { + // Full's device value is cleared by the cascade, not by evict_component + // (aux components read it while freeing); Full-only trees sweep nothing. + let mut tc = core(); + let root = tc.arena.root(); + let a = tc.add_new_node_( + root, + /* key = */ vec![1], + &Tensor::from_slice(&[10]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let mut tracker = HashMap::new(); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.cascade_evict_( + a, + BASE_COMPONENT_TYPE, + &mut tracker, + &mut df, + &mut hf, + EvictLayer::Device, + ); + assert!(!tc.arena.has_device_value(a, FULL)); + // The trigger itself is excluded from the sweep: nothing freed here. + assert!(df.is_empty()); + assert!(tracker.is_empty()); + assert!(!tc.evictable_device_leaves.contains(a)); +} + +#[test] +#[should_panic(expected = "cascade_evict_: EvictLayer::All is not a single layer")] +fn cascade_rejects_the_all_layer() { + let mut tc = core(); + let root = tc.arena.root(); + let a = tc.add_new_node_( + root, + /* key = */ vec![1], + &Tensor::from_slice(&[10]), + /* priority = */ 0, + /* extra_key = */ None, + ); + let mut tracker = HashMap::new(); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.cascade_evict_( + a, + BASE_COMPONENT_TYPE, + &mut tracker, + &mut df, + &mut hf, + EvictLayer::All, + ); +} + +#[test] +fn cascade_host_target_keeps_the_device_value() { + let mut tc = core(); + let root = tc.arena.root(); + let a = tc.add_new_node_( + root, + /* key = */ vec![1], + &Tensor::from_slice(&[10]), + /* priority = */ 0, + /* extra_key = */ None, + ); + tc.arena + .set_host_value(a, FULL, Tensor::from_slice(&[20i64])); + let mut tracker = HashMap::new(); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.cascade_evict_( + a, + BASE_COMPONENT_TYPE, + &mut tracker, + &mut df, + &mut hf, + EvictLayer::Host, + ); + // No tombstone on a host-target cascade; the D-leaf survives. + assert!(tc.arena.has_device_value(a, FULL)); + assert!(tc.evictable_device_leaves.contains(a)); +} + +#[test] +fn cascade_moves_a_host_backed_node_into_the_h_leaf_set() { + let mut tc = core(); + let root = tc.arena.root(); + let a = tc.add_new_node_( + root, + /* key = */ vec![1], + &Tensor::from_slice(&[10]), + /* priority = */ 0, + /* extra_key = */ None, + ); + tc.arena + .set_host_value(a, FULL, Tensor::from_slice(&[20i64])); + let mut tracker = HashMap::new(); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.cascade_evict_( + a, + BASE_COMPONENT_TYPE, + &mut tracker, + &mut df, + &mut hf, + EvictLayer::Device, + ); + // Evicted but backuped: the node leaves the D-set and joins the H-set. + assert!(!tc.evictable_device_leaves.contains(a)); + assert!(tc.evictable_host_leaves.contains(a)); +} + +// A D-leaf carrying both a Full device value and an Swa device value. +fn cascade_aux_setup(tc: &mut UnifiedTreeCore>) -> NodeIdx_ { + let root = tc.arena.root(); + let a = tc.add_new_node_( + root, + /* key = */ vec![1], + &Tensor::from_slice(&[10]), + /* priority = */ 0, + /* extra_key = */ None, + ); + tc.arena.node_mut(a).values[SWA.idx()].value = Some(Tensor::from_slice(&[0i64])); + a +} + +#[test] +fn cascade_sweeps_an_equal_priority_aux_component() { + // Swa leaf priority 0 equals Full's trigger priority: swept, not spared. + let mut tc = core(); + let recorder = Arc::new(SwaEvictionComponentForTest::new( + /* leaf_priority = */ 0, /* internal_priority = */ 2, + )); + tc.register_component_(recorder.clone()); + let a = cascade_aux_setup(&mut tc); + let mut tracker = HashMap::new(); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.cascade_evict_( + a, + BASE_COMPONENT_TYPE, + &mut tracker, + &mut df, + &mut hf, + EvictLayer::Device, + ); + assert_eq!( + *recorder.evictions.lock().unwrap(), + vec![(a, EvictLayer::Device)] + ); + assert!(tc.arena.node(a).values[SWA.idx()].value.is_none()); + assert!(!tc.arena.has_device_value(a, FULL)); + assert_eq!(tracker[&SWA], 1); + assert!(df[&SWA][0].equal(&Tensor::from_slice(&[0i64]))); +} + +#[test] +fn cascade_sweeps_a_lower_priority_aux_component() { + let mut tc = core(); + let recorder = Arc::new(SwaEvictionComponentForTest::new( + /* leaf_priority = */ -1, /* internal_priority = */ 2, + )); + tc.register_component_(recorder.clone()); + let a = cascade_aux_setup(&mut tc); + let mut tracker = HashMap::new(); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.cascade_evict_( + a, + BASE_COMPONENT_TYPE, + &mut tracker, + &mut df, + &mut hf, + EvictLayer::Device, + ); + assert_eq!( + *recorder.evictions.lock().unwrap(), + vec![(a, EvictLayer::Device)] + ); + assert!(tc.arena.node(a).values[SWA.idx()].value.is_none()); +} + +#[test] +fn cascade_spares_a_higher_priority_aux_component() { + // Swa leaf priority 1 outranks Full's trigger priority 0: kept intact. + let mut tc = core(); + let recorder = Arc::new(SwaEvictionComponentForTest::new( + /* leaf_priority = */ 1, /* internal_priority = */ 2, + )); + tc.register_component_(recorder.clone()); + let a = cascade_aux_setup(&mut tc); + let mut tracker = HashMap::new(); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.cascade_evict_( + a, + BASE_COMPONENT_TYPE, + &mut tracker, + &mut df, + &mut hf, + EvictLayer::Device, + ); + assert!(recorder.evictions.lock().unwrap().is_empty()); + assert!(tc.arena.node(a).values[SWA.idx()].value.is_some()); + // The trigger's deferred Full tombstone still lands. + assert!(!tc.arena.has_device_value(a, FULL)); +} + +#[test] +fn cascade_spares_a_locked_component_of_equal_internal_priority() { + // The Swa lock is a legit pin: leaf-collapse flattened priorities, but + // its true internal priority matches the trigger's. + let mut tc = core(); + let recorder = Arc::new(SwaEvictionComponentForTest::new( + /* leaf_priority = */ 0, /* internal_priority = */ 2, + )); + tc.register_component_(recorder.clone()); + let a = cascade_aux_setup(&mut tc); + tc.arena.node_mut(a).values[SWA.idx()].lock_ref = 1; + let mut tracker = HashMap::new(); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.cascade_evict_( + a, + BASE_COMPONENT_TYPE, + &mut tracker, + &mut df, + &mut hf, + EvictLayer::Device, + ); + assert!(recorder.evictions.lock().unwrap().is_empty()); + assert!(tc.arena.node(a).values[SWA.idx()].value.is_some()); + assert_eq!(tc.arena.node(a).values[SWA.idx()].lock_ref, 1); +} + +#[test] +#[should_panic(expected = "a Swa device lock strands node")] +fn cascade_panics_on_a_locked_lower_internal_priority_component() { + // A lock on a strictly-lower-priority tier is a real strand. + let mut tc = core(); + tc.register_component_(Arc::new(SwaEvictionComponentForTest::new( + /* leaf_priority = */ 0, /* internal_priority = */ 1, + ))); + let a = cascade_aux_setup(&mut tc); + tc.arena.node_mut(a).values[SWA.idx()].lock_ref = 1; + let mut tracker = HashMap::new(); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.cascade_evict_( + a, + BASE_COMPONENT_TYPE, + &mut tracker, + &mut df, + &mut hf, + EvictLayer::Device, + ); +} + +// An H-tier aux carrier: a raw child holding only an SWA host value. +fn cascade_host_aux_setup(tc: &mut UnifiedTreeCore>) -> NodeIdx_ { + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena.set_host_value(a, SWA, Tensor::from_slice(&[0i64])); + a +} + +#[test] +fn cascade_host_spares_a_locked_component_of_equal_internal_priority() { + // The Swa host lock is a legit pin: its true internal priority matches + // the trigger's, so the host cascade skips it and keeps the host value. + let mut tc = core(); + let recorder = Arc::new(SwaEvictionComponentForTest::new( + /* leaf_priority = */ 0, /* internal_priority = */ 2, + )); + tc.register_component_(recorder.clone()); + let a = cascade_host_aux_setup(&mut tc); + tc.arena + .node_mut(a) + .set_lock_ref_(ValueSlotIdx::host(SWA), 1); + let mut tracker = HashMap::new(); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.cascade_evict_( + a, + BASE_COMPONENT_TYPE, + &mut tracker, + &mut df, + &mut hf, + EvictLayer::Host, + ); + assert!(recorder.evictions.lock().unwrap().is_empty()); + assert!(tc.arena.has_host_value(a, SWA)); + assert_eq!(tc.arena.host_lock_ref(a, SWA), 1); +} + +#[test] +#[should_panic(expected = "a Swa host lock strands node")] +fn cascade_host_panics_on_a_locked_lower_internal_priority_component() { + // A host lock on a strictly-lower-priority tier is a real strand. + let mut tc = core(); + tc.register_component_(Arc::new(SwaEvictionComponentForTest::new( + /* leaf_priority = */ 0, /* internal_priority = */ 1, + ))); + let a = cascade_host_aux_setup(&mut tc); + tc.arena + .node_mut(a) + .set_lock_ref_(ValueSlotIdx::host(SWA), 1); + let mut tracker = HashMap::new(); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.cascade_evict_( + a, + BASE_COMPONENT_TYPE, + &mut tracker, + &mut df, + &mut hf, + EvictLayer::Host, + ); +} + +#[test] +fn evict_walk_and_driver_empty_the_tree_end_to_end() { + // The full eviction loop: insert three leaves, walk them in LRU + // order, evict each through the driver, and end with a bare root. + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + tc.insert(&insert_params(&vec![3], &[12])); + tc.insert(&insert_params(&vec![4], &[13])); + let mut tracker = HashMap::from([(FULL, 0)]); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.evict_device_start(FULL, /* request_cnt = */ 100); + let mut evicted = 0; + loop { + let (leaf, step) = tc.evict_device_next_node(FULL, &tracker); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + let Some(leaf) = leaf else { break }; + let (backup, step) = tc.evict_device_leaf(leaf, /* is_write_back = */ false); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + assert!(backup.is_none()); + evicted += 1; + } + tc.evict_device_end(FULL); + assert_eq!(evicted, 3); + assert_eq!(tracker[&FULL], 4); + assert_eq!(df[&FULL].len(), 3); + assert_eq!(tc.arena.len(), 1); + assert_eq!(tc.evictable_device_leaves.len(), 0); + assert_eq!(tc.evictable_size_(FULL), 0); +} + +#[test] +fn evict_driver_readmits_the_parent_into_the_walk() { + // A two-level chain evicts leaf-first: the child's eviction turns the + // prefix node into the next walkable D-leaf. + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + tc.insert(&insert_params(&vec![1, 2, 9], &[10, 11, 29])); + let mut tracker = HashMap::from([(FULL, 0)]); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.evict_device_start(FULL, /* request_cnt = */ 100); + let mut evicted = 0; + loop { + let (leaf, step) = tc.evict_device_next_node(FULL, &tracker); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + let Some(leaf) = leaf else { break }; + let (_, step) = tc.evict_device_leaf(leaf, false); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + evicted += 1; + } + tc.evict_device_end(FULL); + // Two suffix leaves plus the readmitted split-prefix node. + assert_eq!(evicted, 3); + assert_eq!(tracker[&FULL], 4); + assert_eq!(tc.arena.len(), 1); +} + +#[test] +fn release_all_component_layers_scrubs_a_host_leaf_from_both_leaf_sets() { + // Host leaves only reach this helper via subtree drops; pin the discard. + let mut tc = core(); + let (_a, b) = matched_chain(&mut tc); + let taken = tc.arena.take_device_value(b, FULL); + tc.arena.set_host_value(b, FULL, taken); + tc.update_evictable_leaf_sets_(b); + assert!(tc.evictable_host_leaves.contains(b)); + + let mut tracker = HashMap::new(); + let (mut device_frees, mut host_frees) = (HashMap::new(), HashMap::new()); + tc.release_all_component_layers_( + b, + StorageMedium::Cpu, + &mut tracker, + &mut device_frees, + &mut host_frees, + ); + assert!(!tc.evictable_host_leaves.contains(b)); + assert!(!tc.evictable_device_leaves.contains(b)); + assert!(host_frees[&FULL][0].equal(&Tensor::from_slice(&[12i64]))); +} + +#[test] +fn evict_driver_deletes_through_a_tombstone_parent() { + // A leaf under a valueless internal node cascades the delete upward. + let mut tc = core(); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let b = tc.add_new_node_( + a, + /* key = */ vec![2], + &Tensor::from_slice(&[10]), + /* priority = */ 0, + /* extra_key = */ None, + ); + tc.evict_device_leaf(tc.arena.node(b).id, false); + assert_eq!(tc.arena.len(), 1); +} + +#[test] +fn insert_after_eviction_reuses_the_slot_but_never_the_handle() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + tc.insert(&insert_params(&vec![4], &[13])); + let leaf = tc.match_prefix(&match_params(&vec![4])).best_match_node_id; + let leaf_idx = tc.arena.resolve(leaf); + tc.evict_device_leaf(leaf, /* is_write_back = */ false); + assert_eq!(tc.arena.len(), 2); + // The freed handle no longer resolves. + assert!(tc.arena.try_resolve(leaf).is_none()); + // The splitting insert allocates its prefix node into the freed slot. + tc.insert(&insert_params(&vec![1, 2, 9], &[20, 21, 29])); + assert_eq!(tc.arena.len(), 4); + let prefix = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + // Slot recycled, but the stale handle can never alias the new node. + assert_eq!(tc.arena.resolve(prefix), leaf_idx); + assert_ne!(prefix, leaf); +} + +#[test] +#[should_panic(expected = "is not allocated")] +fn stale_handle_panics_after_its_node_is_freed() { + let mut tc = core(); + tc.insert(&insert_params(&vec![4], &[13])); + let leaf = tc.match_prefix(&match_params(&vec![4])).best_match_node_id; + tc.evict_device_leaf(leaf, /* is_write_back = */ false); + tc.inc_lock_ref(leaf); +} + +#[test] +#[should_panic(expected = "is not allocated")] +fn pre_reset_handle_panics_after_reset() { + let mut tc = core(); + tc.insert(&insert_params(&vec![4], &[13])); + let leaf = tc.match_prefix(&match_params(&vec![4])).best_match_node_id; + tc.reset(); + tc.arena.resolve(leaf); +} + +#[test] +#[should_panic(expected = "is not a D-leaf")] +fn evict_driver_rejects_a_non_leaf() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + tc.insert(&insert_params(&vec![1, 2, 9], &[10, 11, 29])); + // The split-prefix node has valued children: not a D-leaf. + let prefix = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + tc.evict_device_leaf(prefix, false); +} + +#[test] +fn evict_driver_write_back_returns_the_backup_action_for_an_unbacked_leaf() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1], &[10])); + let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; + let mut tracker = HashMap::from([(FULL, 0)]); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + let (action, step) = tc.evict_device_leaf(leaf, /* is_write_back = */ true); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + // Write-back carries only the leaf itself; nothing is freed yet. + assert_eq!(action.unwrap().node_ids, vec![leaf]); + assert!(!tc.arena.node(tc.arena.resolve(leaf)).evicted()); + assert_eq!(tracker[&FULL], 0); + assert!(df.is_empty() && hf.is_empty()); +} + +#[test] +fn evict_driver_demotes_a_backuped_leaf_to_host_only() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1], &[10])); + let leaf = tc.match_prefix(&match_params(&vec![1])).best_match_node_id; + tc.arena + .set_host_value(tc.arena.resolve(leaf), FULL, Tensor::from_slice(&[20i64])); + let mut tracker = HashMap::from([(FULL, 0)]); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + let (action, step) = tc.evict_device_leaf(leaf, false); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + assert!(action.is_none()); + // The node stays in the tree, now host-only. + let node = tc.arena.node(tc.arena.resolve(leaf)); + assert!(node.evicted() && node.backuped()); + assert_eq!(tracker[&FULL], 1); + assert_eq!(df[&FULL].len(), 1); + assert!(hf.is_empty()); + assert!(!tc.evictable_device_leaves.contains(tc.arena.resolve(leaf))); + assert!(tc.evictable_host_leaves.contains(tc.arena.resolve(leaf))); + tc.sanity_check(&[], &[]); +} + +#[test] +fn evict_device_leaf_step_counts_are_independent_of_prior_evictions() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + tc.insert(&insert_params(&vec![4], &[13])); + let first = tc + .match_prefix(&match_params(&vec![1, 2, 3])) + .best_match_node_id; + let second = tc.match_prefix(&match_params(&vec![4])).best_match_node_id; + let (_, step) = tc.evict_device_leaf(first, /* is_write_back = */ false); + assert_eq!(step.tracker[&FULL], 3); + // The second step reports only its own leaf, not a running total. + let (_, step) = tc.evict_device_leaf(second, /* is_write_back = */ false); + assert_eq!(step.tracker[&FULL], 1); + assert_eq!(step.device_frees[&FULL].len(), 1); +} + +#[test] +fn empty_tree_device_evict_walk_returns_nothing_for_each_component() { + let mut tc = core(); + tc.evict_device_start(FULL, /* request_cnt = */ 10); + let (leaf, step) = tc.evict_device_next_node(FULL, &HashMap::from([(FULL, 0)])); + assert_eq!(leaf, None); + assert!(step.tracker.is_empty()); + assert!(step.device_frees.is_empty() && step.host_frees.is_empty()); + tc.evict_device_end(FULL); + assert_eq!(tc.evictable_size_(FULL), 0); + + let mut tc = swa_match_core(/* window = */ 4); + tc.evict_device_start(SWA, /* request_cnt = */ 10); + let (leaf, step) = tc.evict_device_next_node(SWA, &HashMap::from([(FULL, 0), (SWA, 0)])); + assert_eq!(leaf, None); + assert!(step.tracker.is_empty()); + assert!(step.device_frees.is_empty() && step.host_frees.is_empty()); + tc.evict_device_end(SWA); + assert_eq!(tc.evictable_size_(SWA), 0); + + let mut tc = UnifiedTreeCore::>::new( + CacheInitParams { + mamba_cache_chunk_size: Some(256), + ..CacheInitParams::default() + }, + vec![FULL, MAMBA], + ); + tc.evict_device_start(MAMBA, /* request_cnt = */ 10); + let (leaf, step) = tc.evict_device_next_node(MAMBA, &HashMap::from([(FULL, 0), (MAMBA, 0)])); + assert_eq!(leaf, None); + assert!(step.tracker.is_empty()); + assert!(step.device_frees.is_empty() && step.host_frees.is_empty()); + tc.evict_device_end(MAMBA); + assert_eq!(tc.evictable_size_(MAMBA), 0); +} + +#[test] +fn empty_match_result_anchors_all_boundaries_at_the_root() { + let tc = core(); + let root = tc.arena.root(); + let result = tc.empty_match_result(); + assert_eq!(result.last_device_node_id, tc.arena.node(root).id); + assert_eq!(result.last_host_node_id, tc.arena.node(root).id); + assert_eq!(result.best_match_node_id, tc.arena.node(root).id); + assert_eq!(result.host_hit_length, 0); + assert_eq!(result.device_indices.numel(), 0); + assert_eq!(result.device_indices.kind(), Kind::Int64); +} + +#[test] +fn touch_node_stamps_a_fresh_access_tick() { + let mut tc = core(); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let created_at = tc.arena.node(a).last_access_counter; + tc.touch_node_(a); + let first_touch = tc.arena.node(a).last_access_counter; + assert!(first_touch > created_at); + tc.touch_node_(a); + assert!(tc.arena.node(a).last_access_counter > first_touch); +} + +#[test] +#[should_panic(expected = "not implemented")] +fn touch_node_refreshes_aux_component_lrus() { + let mut tc = core(); + tc.register_component_(Arc::new(SwaComponentForTest)); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.touch_node_(a); +} + +#[test] +fn touch_node_on_a_root_only_stamps_the_tick() { + // Even with an aux component present, a root touch skips the refresh loop. + let mut tc = core(); + tc.register_component_(Arc::new(SwaComponentForTest)); + let root = tc.arena.root(); + let before = tc.arena.node(root).last_access_counter; + tc.touch_node_(root); + assert!(tc.arena.node(root).last_access_counter > before); +} + +#[test] +fn inc_hit_count_bumps_and_stays_quiet_without_hicache() { + let mut tc = core(); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(a, FULL, Tensor::from_slice(&[0i64])); + assert!(!tc.inc_hit_count_and_check_(a, /* chunked = */ false)); + assert_eq!(tc.arena.node(a).hit_count, 1); + // The tree defaults keep the host tier off and the threshold at 256. + assert!(!tc.enable_hicache); + assert_eq!(tc.write_through_threshold, 256); +} + +#[test] +fn inc_hit_count_skips_evicted_or_chunked_nodes() { + let mut tc = core(); + let root = tc.arena.root(); + let evicted = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let chunked = tc + .arena + .alloc_child( + root, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(chunked, FULL, Tensor::from_slice(&[0i64])); + assert!(!tc.inc_hit_count_and_check_(evicted, /* chunked = */ false)); + assert_eq!(tc.arena.node(evicted).hit_count, 0); + assert!(!tc.inc_hit_count_and_check_(chunked, /* chunked = */ true)); + assert_eq!(tc.arena.node(chunked).hit_count, 0); +} + +#[test] +fn inc_hit_count_is_a_noop_in_write_back_mode() { + let params = CacheInitParams { + is_write_back: true, + ..Default::default() + }; + let mut tc: UnifiedTreeCore> = UnifiedTreeCore::new(params, vec![FULL]); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(a, FULL, Tensor::from_slice(&[0i64])); + assert!(!tc.inc_hit_count_and_check_(a, /* chunked = */ false)); + assert_eq!(tc.arena.node(a).hit_count, 0); +} + +#[test] +fn inc_hit_count_fires_the_write_through_check() { + let mut tc = core(); + tc.set_hicache_enabled(); + tc.write_through_threshold = 2; + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(a, FULL, Tensor::from_slice(&[0i64])); + assert!(!tc.inc_hit_count_and_check_(a, /* chunked = */ false)); + assert!(tc.inc_hit_count_and_check_(a, /* chunked = */ false)); + // A backuped node never re-fires. + tc.arena + .set_host_value(a, FULL, Tensor::from_slice(&[0i64])); + assert!(!tc.inc_hit_count_and_check_(a, /* chunked = */ false)); + assert_eq!(tc.arena.node(a).hit_count, 3); +} + +#[test] +#[should_panic(expected = "Swa component is not enabled")] +fn dispatch_panics_for_a_disabled_component() { + let mut tc = core(); + tc.evict_device_start(SWA, /* request_cnt = */ 1); +} + +#[test] +fn registry_and_map_share_the_drivers() { + let tc = core(); + let component = &tc.components[0]; + assert_eq!(component.component_type(), FULL); + assert!(Arc::ptr_eq( + component, + tc.components_by_type[FULL.idx()].as_ref().unwrap() + )); +} + +#[test] +#[should_panic(expected = "at least one component type is required")] +fn new_rejects_an_empty_component_list() { + UnifiedTreeCore::>::new(CacheInitParams::default(), vec![]); +} + +#[test] +#[should_panic(expected = "duplicate component type Full")] +fn new_rejects_duplicate_component_types() { + UnifiedTreeCore::>::new(CacheInitParams::default(), vec![FULL, FULL]); +} + +#[test] +#[should_panic(expected = "the base (Full) component is required")] +fn new_requires_the_full_component() { + UnifiedTreeCore::>::new(CacheInitParams::default(), vec![SWA]); +} + +#[test] +fn new_resolves_the_eviction_policy() { + let mut tc = UnifiedTreeCore::new( + CacheInitParams { + eviction_policy: "FIFO".to_string(), + ..Default::default() + }, + vec![FULL], + ); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena.node_mut(a).creation_counter = 7; + assert_eq!( + tc.eviction_strategy.get_priority(tc.arena.node(a)), + crate::unified_lru_list::PriorityKey(7, 0) + ); +} + +#[test] +fn new_builds_independent_lru_lists() { + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.device_lru_list_mut(FULL).insert_mru(n1); + // Every component/tier list exists as an independent container. + assert!(tc.device_lru_list(FULL).in_list(Some(n1))); + assert!(!tc.host_lru_list(FULL).in_list(Some(n1))); + assert!(!tc.device_lru_list(SWA).in_list(Some(n1))); +} + +#[test] +fn update_adds_an_unlocked_device_valued_leaf() { + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(n1, FULL, Tensor::from_slice(&[10i64])); + tc.update_evictable_leaf_sets_(n1); + assert!(tc.evictable_device_leaves.contains(n1)); + assert!(!tc.evictable_host_leaves.contains(n1)); +} + +#[test] +fn update_discards_a_node_that_stops_being_a_device_leaf() { + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(n1, FULL, Tensor::from_slice(&[10i64])); + tc.update_evictable_leaf_sets_(n1); + assert!(tc.evictable_device_leaves.contains(n1)); + let _ = tc.arena.take_device_value(n1, FULL); + tc.update_evictable_leaf_sets_(n1); + assert!(!tc.evictable_device_leaves.contains(n1)); +} + +#[test] +fn device_leaf_excludes_the_root() { + let mut tc = core(); + let root = tc.arena.root(); + // Root keys are empty, so a key-aligned device value is the empty tensor. + let empty: [i64; 0] = []; + tc.arena + .set_device_value(root, FULL, Tensor::from_slice(&empty)); + tc.update_evictable_leaf_sets_(root); + assert!(!tc.evictable_device_leaves.contains(root)); + assert!(!tc.evictable_host_leaves.contains(root)); +} + +#[test] +fn device_leaf_excludes_locked_node_until_unlocked() { + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(n1, FULL, Tensor::from_slice(&[10i64])); + tc.arena + .node_mut(n1) + .set_lock_ref_(ValueSlotIdx::device(FULL), 1); + tc.update_evictable_leaf_sets_(n1); + assert!(!tc.evictable_device_leaves.contains(n1)); + tc.arena + .node_mut(n1) + .set_lock_ref_(ValueSlotIdx::device(FULL), 0); + tc.update_evictable_leaf_sets_(n1); + assert!(tc.evictable_device_leaves.contains(n1)); +} + +#[test] +fn device_leaf_excludes_node_locked_by_another_component() { + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(n1, FULL, Tensor::from_slice(&[10i64])); + // The lock check spans every component, not just Full. + tc.arena.node_mut(n1).values[SWA.idx()].lock_ref = 1; + tc.update_evictable_leaf_sets_(n1); + assert!(!tc.evictable_device_leaves.contains(n1)); +} + +#[test] +fn device_leaf_excludes_parent_with_a_device_valued_child() { + let mut tc = core(); + let root = tc.arena.root(); + let parent = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let child = tc + .arena + .alloc_child( + parent, + /* key = */ vec![3], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(parent, FULL, Tensor::from_slice(&[10i64, 11])); + tc.arena + .set_device_value(child, FULL, Tensor::from_slice(&[30i64])); + tc.update_evictable_leaf_sets_(parent); + tc.update_evictable_leaf_sets_(child); + assert!(!tc.evictable_device_leaves.contains(parent)); + assert!(tc.evictable_device_leaves.contains(child)); + // Once the child's device value is evicted, the parent becomes the D-leaf. + let _ = tc.arena.take_device_value(child, FULL); + tc.update_evictable_leaf_sets_(parent); + tc.update_evictable_leaf_sets_(child); + assert!(tc.evictable_device_leaves.contains(parent)); + assert!(!tc.evictable_device_leaves.contains(child)); +} + +#[test] +fn update_discards_a_node_that_stops_being_a_host_leaf() { + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(n1, FULL, Tensor::from_slice(&[10i64])); + tc.update_evictable_leaf_sets_(n1); + assert!(tc.evictable_host_leaves.contains(n1)); + // Loading the device value back means the node is no longer evicted. + tc.arena + .set_device_value(n1, FULL, Tensor::from_slice(&[20i64])); + tc.update_evictable_leaf_sets_(n1); + assert!(!tc.evictable_host_leaves.contains(n1)); +} + +#[test] +fn host_leaf_excludes_the_root() { + let mut tc = core(); + let root = tc.arena.root(); + // Root keys are empty, so a key-aligned host value is the empty tensor. + let empty: [i64; 0] = []; + tc.arena + .set_host_value(root, FULL, Tensor::from_slice(&empty)); + tc.update_evictable_leaf_sets_(root); + assert!(!tc.evictable_host_leaves.contains(root)); +} + +#[test] +fn host_leaf_true_for_evicted_backuped_childless_node() { + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(n1, FULL, Tensor::from_slice(&[10i64])); + tc.update_evictable_leaf_sets_(n1); + assert!(tc.evictable_host_leaves.contains(n1)); + assert!(!tc.evictable_device_leaves.contains(n1)); +} + +#[test] +fn device_valued_node_is_not_a_host_leaf() { + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(n1, FULL, Tensor::from_slice(&[10i64])); + tc.arena + .set_host_value(n1, FULL, Tensor::from_slice(&[20i64])); + tc.update_evictable_leaf_sets_(n1); + assert!(tc.evictable_device_leaves.contains(n1)); + assert!(!tc.evictable_host_leaves.contains(n1)); +} + +#[test] +fn host_leaf_excludes_node_with_children() { + let mut tc = core(); + let root = tc.arena.root(); + let parent = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let _child = tc + .arena + .alloc_child( + parent, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(parent, FULL, Tensor::from_slice(&[10i64])); + tc.update_evictable_leaf_sets_(parent); + assert!(!tc.evictable_host_leaves.contains(parent)); +} + +#[test] +fn host_leaf_excludes_host_locked_node_until_unlocked() { + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(n1, FULL, Tensor::from_slice(&[10i64])); + tc.arena + .node_mut(n1) + .set_lock_ref_(ValueSlotIdx::host(FULL), 1); + tc.update_evictable_leaf_sets_(n1); + assert!(!tc.evictable_host_leaves.contains(n1)); + tc.arena + .node_mut(n1) + .set_lock_ref_(ValueSlotIdx::host(FULL), 0); + tc.update_evictable_leaf_sets_(n1); + assert!(tc.evictable_host_leaves.contains(n1)); +} + +#[test] +fn host_leaf_excludes_node_host_locked_by_another_component() { + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(n1, FULL, Tensor::from_slice(&[10i64])); + // The host-lock check spans every component, not just Full. + tc.arena + .node_mut(n1) + .state_mut_(ValueSlotIdx::host(SWA)) + .lock_ref = 1; + tc.update_evictable_leaf_sets_(n1); + assert!(!tc.evictable_host_leaves.contains(n1)); +} + +#[test] +fn tombstone_without_backup_is_in_neither_set() { + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.update_evictable_leaf_sets_(n1); + assert!(!tc.evictable_device_leaves.contains(n1)); + assert!(!tc.evictable_host_leaves.contains(n1)); +} + +#[test] +#[should_panic(expected = "out of bounds")] +fn update_panics_on_missing_node() { + let mut tc = core(); + tc.update_evictable_leaf_sets_(NodeIdx_(999)); +} + +#[test] +fn evict_host_leaf_frees_host_values_and_credits_the_tracker() { + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(n1, FULL, Tensor::from_slice(&[10i64, 11])); + tc.evictable_host_leaves.add(n1); + let mut tr = HashMap::from([(FULL, 0)]); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.evict_host_leaf_(n1, &mut tr, &mut df, &mut hf); + assert_eq!(tr[&FULL], 2); + assert_eq!(hf[&FULL].len(), 1); + assert!(df.is_empty()); + assert!(tc.evictable_host_leaves.is_empty()); + assert_eq!(tc.arena.len(), 1); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "is not an H-leaf")] +fn evict_host_leaf_panics_on_a_device_valued_node() { + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_device_value(n1, FULL, Tensor::from_slice(&[10i64])); + let mut tr = HashMap::new(); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.evict_host_leaf_(n1, &mut tr, &mut df, &mut hf); +} + +#[test] +fn evict_host_leaf_cascades_tombstone_ancestors() { + let mut tc = core(); + let root = tc.arena.root(); + let t = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let h = tc + .arena + .alloc_child( + t, + /* key = */ vec![2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(h, FULL, Tensor::from_slice(&[20i64])); + tc.evictable_host_leaves.add(h); + let mut tr = HashMap::from([(FULL, 0)]); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.evict_host_leaf_(h, &mut tr, &mut df, &mut hf); + // The valueless ancestor t is deleted by the tombstone walk. + assert_eq!(tc.arena.len(), 1); + assert_eq!(tr[&FULL], 1); + tc.sanity_check(&[], &[]); +} + +#[test] +fn drive_host_eviction_is_a_noop_for_an_absent_component() { + let mut tc = core(); + let root = tc.arena.root(); + let n1 = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena + .set_host_value(n1, FULL, Tensor::from_slice(&[10i64])); + tc.evictable_host_leaves.add(n1); + let mut tr = HashMap::from([(SWA, 0)]); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + accumulate_step( + tc.drive_host_eviction(SWA, /* num_tokens = */ 100), + &mut tr, + &mut df, + &mut hf, + ); + assert_eq!(tr[&SWA], 0); + assert!(hf.is_empty()); + assert!(tc.evictable_host_leaves.contains(n1)); +} + +#[test] +fn drive_host_eviction_keeps_zero_delta_tracker_entries() { + let mut tc = core(); + let result = tc.drive_host_eviction(FULL, /* num_tokens = */ 10); + assert_eq!(result.tracker, HashMap::from([(FULL, 0)])); + assert!(result.device_frees.is_empty()); + assert!(result.host_frees.is_empty()); +} + +#[test] +fn drive_host_eviction_dispatches_reclaim_only_under_write_back() { + let mut tc = UnifiedTreeCore::new( + CacheInitParams { + is_write_back: true, + ..CacheInitParams::default() + }, + vec![FULL], + ); + let recorder = Arc::new(RecordingComponentForTest::default()); + tc.register_component_(recorder.clone()); + + let result = tc.drive_host_eviction(SWA, /* num_tokens = */ 10); + + assert_eq!( + *recorder.host_eviction_calls.lock().unwrap(), + vec![("reclaim", 0), ("drive", 2)] + ); + assert_eq!(result.tracker, HashMap::from([(SWA, 5)])); + + recorder.host_eviction_calls.lock().unwrap().clear(); + tc.is_write_back = false; + let result = tc.drive_host_eviction(SWA, /* num_tokens = */ 10); + + assert_eq!( + *recorder.host_eviction_calls.lock().unwrap(), + vec![("drive", 0)] + ); + assert_eq!(result.tracker, HashMap::from([(SWA, 3)])); +} + +#[test] +fn drive_host_eviction_default_reclaim_hook_is_a_noop() { + let mut tc = UnifiedTreeCore::new( + CacheInitParams { + is_write_back: true, + ..CacheInitParams::default() + }, + vec![FULL], + ); + tc.register_component_(Arc::new(SwaComponentForTest)); + + let result = tc.drive_host_eviction(SWA, /* num_tokens = */ 10); + + assert_eq!(result.tracker, HashMap::from([(SWA, 0)])); + assert!(result.device_frees.is_empty()); + assert!(result.host_frees.is_empty()); +} + +#[test] +fn evict_layer_contains_matches_intflag_membership() { + assert!(EvictLayer::Device.contains(EvictLayer::Device)); + assert!(!EvictLayer::Device.contains(EvictLayer::Host)); + assert!(EvictLayer::Host.contains(EvictLayer::Host)); + assert!(!EvictLayer::Host.contains(EvictLayer::Device)); + assert!(EvictLayer::All.contains(EvictLayer::Device)); + assert!(EvictLayer::All.contains(EvictLayer::Host)); + assert!(EvictLayer::All.contains(EvictLayer::All)); + assert!(!EvictLayer::Device.contains(EvictLayer::All)); + assert!(!EvictLayer::Host.contains(EvictLayer::All)); +} + +#[test] +fn reset_restores_a_fresh_tree() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some("chat"), None), + ..insert_params(&vec![7, 8], &[20, 21]) + }); + let matched = tc.match_prefix(&match_params(&vec![1, 2, 3])); + tc.inc_lock_ref(matched.best_match_node_id); + assert_eq!(tc.protected_size(), 3); + // Seed aux LRU, host LRU, and host-leaf state so the reset must clear each. + let root = tc.arena.root(); + let d = tc + .arena + .alloc_child( + root, + /* key = */ vec![9], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + let h = tc + .arena + .alloc_child( + root, + /* key = */ vec![12], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.device_lru_list_mut(SWA).insert_mru(d); + tc.host_lru_list_mut(SWA).insert_mru(h); + tc.evictable_host_leaves.add(h); + tc.reset(); + assert_eq!(tc.arena.len(), 1); + assert_eq!(tc.evictable_size(), 0); + assert_eq!(tc.protected_size(), 0); + assert_eq!(tc.total_size(), (0, 0)); + assert!(tc.evictable_device_leaves.is_empty()); + assert!(tc.evictable_host_leaves.is_empty()); + assert_eq!(tc.device_lru_list(FULL).len(), 0); + assert_eq!(tc.device_lru_list(SWA).len(), 0); + assert_eq!(tc.host_lru_list(SWA).len(), 0); + let matched = tc.match_prefix(&match_params(&vec![1, 2, 3])); + assert_eq!(matched.device_indices.numel(), 0); + // The tree accepts fresh inserts after the reset. + tc.insert(&insert_params(&vec![4, 5], &[30, 31])); + assert_eq!(tc.evictable_size(), 2); +} + +#[test] +fn size_accessors_mirror_the_full_component_state() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + assert_eq!(tc.evictable_size(), 3); + assert_eq!(tc.full_evictable_size(), 3); + assert_eq!(tc.protected_size(), 0); + assert_eq!(tc.component_evictable_size(FULL), 3); + let matched = tc.match_prefix(&match_params(&vec![1, 2, 3])); + tc.inc_lock_ref(matched.best_match_node_id); + assert_eq!(tc.protected_size(), 3); + assert_eq!(tc.full_protected_size(), 3); + assert_eq!(tc.evictable_size(), 0); +} + +#[test] +fn swa_size_accessors_mirror_the_swa_component_state() { + let mut tc = core(); + tc.component_state_mut(SWA).evictable_size = 2; + tc.component_state_mut(SWA).protected_size = 1; + assert_eq!(tc.swa_evictable_size(), 2); + assert_eq!(tc.swa_protected_size(), 1); + // The Full accessors read their own slot, untouched by the SWA seed. + assert_eq!(tc.full_evictable_size(), 0); + assert_eq!(tc.full_protected_size(), 0); +} + +#[test] +fn total_size_spans_namespaces_and_aux_values() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some("chat"), None), + ..insert_params(&vec![7, 8], &[20, 21]) + }); + assert_eq!(tc.total_size(), (5, 0)); + // An SWA-valued node adds to the aux total only. + tc.register_component_(Arc::new(SwaComponentForTest)); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![9, 10, 11], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.arena.node_mut(a).values[SWA.idx()].value = Some(Tensor::from_slice(&[0i64, 1, 2])); + assert_eq!(tc.total_size(), (5, 3)); +} + +// Zip a canary walk into sorted (slot, position, prev_slot) rows; emission order is not a contract. +fn sorted_canary_rows(walk: KvCanaryWalkResult) -> Vec<(i64, i64, i64)> { + let mut rows: Vec<(i64, i64, i64)> = walk + .slot_indices + .into_iter() + .zip(walk.positions) + .zip(walk.prev_slot_indices) + .map(|((slot, position), prev)| (slot, position, prev)) + .collect(); + rows.sort_unstable(); + rows +} + +#[test] +fn walk_for_kv_canary_on_an_empty_tree_emits_nothing() { + let tc = core(); + assert_eq!( + sorted_canary_rows(tc.walk_for_kv_canary(false, false)), + Vec::<(i64, i64, i64)>::new() + ); +} + +#[test] +fn walk_for_kv_canary_chains_slots_across_namespaces() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some("chat"), None), + ..insert_params(&vec![7, 8], &[20, 21]) + }); + assert_eq!( + sorted_canary_rows(tc.walk_for_kv_canary(false, false)), + vec![ + (10, 0, -1), + (11, 1, 10), + (12, 2, 11), + (20, 0, -1), + (21, 1, 20) + ] + ); +} + +#[test] +fn walk_for_kv_canary_unlocked_only_skips_locked_nodes_but_keeps_the_chain() { + let mut tc = core(); + let (a, _b) = matched_chain(&mut tc); + tc.inc_lock_ref(tc.arena.node(a).id); + assert_eq!( + sorted_canary_rows(tc.walk_for_kv_canary(true, false)), + vec![(12, 2, 11)] + ); + assert_eq!( + sorted_canary_rows(tc.walk_for_kv_canary(false, false)), + vec![(10, 0, -1), (11, 1, 10), (12, 2, 11)] + ); +} + +#[test] +fn walk_for_kv_canary_spans_device_evicted_nodes_without_emitting_them() { + let mut tc = core(); + let (a, _b) = matched_chain(&mut tc); + let _ = tc.arena.take_device_value(a, FULL); + assert_eq!( + sorted_canary_rows(tc.walk_for_kv_canary(false, false)), + vec![(12, 2, -1)] + ); +} + +#[test] +fn walk_for_kv_canary_swa_resident_only_skips_swa_tombstoned_nodes() { + let mut tc = core(); + let (a, _b) = matched_chain(&mut tc); + tc.register_component_(Arc::new(SwaComponentForTest)); + tc.arena + .set_device_value(a, SWA, Tensor::from_slice(&[0i64, 1])); + assert_eq!( + sorted_canary_rows(tc.walk_for_kv_canary(false, true)), + vec![(10, 0, -1), (11, 1, 10)] + ); +} + +#[test] +fn walk_for_kv_canary_swa_filter_is_inert_without_the_swa_component() { + let mut tc = core(); + let (_a, _b) = matched_chain(&mut tc); + assert_eq!( + sorted_canary_rows(tc.walk_for_kv_canary(false, true)), + vec![(10, 0, -1), (11, 1, 10), (12, 2, 11)] + ); +} + +#[test] +fn walk_for_kv_canary_unlocked_only_gates_on_the_swa_lock_under_the_swa_filter() { + let mut tc = core(); + let (a, _b) = matched_chain(&mut tc); + tc.register_component_(Arc::new(SwaComponentForTest)); + tc.arena + .set_device_value(a, SWA, Tensor::from_slice(&[0i64, 1])); + // A node can hold Full KV for a running request while its SWA slots are unused. + tc.arena.node_mut(a).values[FULL.idx()].lock_ref = 1; + assert_eq!( + sorted_canary_rows(tc.walk_for_kv_canary(true, true)), + vec![(10, 0, -1), (11, 1, 10)] + ); + // A held SWA lock excludes the node even with Full unlocked. + tc.arena.node_mut(a).values[FULL.idx()].lock_ref = 0; + tc.arena.node_mut(a).values[SWA.idx()].lock_ref = 1; + assert_eq!( + sorted_canary_rows(tc.walk_for_kv_canary(true, true)), + Vec::<(i64, i64, i64)>::new() + ); +} + +#[test] +fn get_component_device_value_reads_the_full_value() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 3])) + .best_match_node_id; + assert!( + tc.get_component_device_value(leaf, FULL) + .unwrap() + .equal(&Tensor::from_slice(&[10i64, 11, 12])) + ); + let _ = tc.arena.take_device_value(tc.arena.resolve(leaf), FULL); + assert!(tc.get_component_device_value(leaf, FULL).is_none()); +} + +#[test] +#[should_panic(expected = "Swa component is not enabled")] +fn get_component_device_value_panics_on_an_unregistered_component() { + let tc = core(); + let root = tc.arena.root(); + tc.get_component_device_value(tc.arena.node(root).id, SWA); +} + +#[test] +fn get_component_device_value_reads_the_registered_components_slot() { + let mut tc = core(); + tc.register_component_(Arc::new(SwaComponentForTest)); + let (a, _b) = matched_chain(&mut tc); + assert!( + tc.get_component_device_value(tc.arena.node(a).id, SWA) + .is_none() + ); + tc.arena + .set_device_value(a, SWA, Tensor::from_slice(&[5i64, 6])); + assert_eq!( + Vec::::try_from( + tc.get_component_device_value(tc.arena.node(a).id, SWA) + .unwrap() + ) + .unwrap(), + vec![5, 6] + ); +} + +#[test] +fn component_evictable_size_is_zero_for_an_absent_component() { + assert_eq!(core().component_evictable_size(SWA), 0); +} + +#[test] +fn component_evictable_size_reads_the_registered_components_state() { + let mut tc = core(); + tc.register_component_(Arc::new(SwaComponentForTest)); + tc.component_state_mut(SWA).evictable_size = 7; + assert_eq!(tc.component_evictable_size(SWA), 7); +} + +#[test] +fn component_protected_size_is_zero_for_an_absent_component() { + assert_eq!(core().component_protected_size(SWA), 0); +} + +#[test] +fn component_protected_size_reads_the_registered_components_state() { + let mut tc = core(); + tc.register_component_(Arc::new(SwaComponentForTest)); + tc.component_state_mut(SWA).protected_size = 7; + assert_eq!(tc.component_protected_size(SWA), 7); +} + +#[test] +fn is_full_device_evicted_flips_when_the_value_tombstones() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 3])) + .best_match_node_id; + assert!(!tc.is_full_device_evicted(leaf)); + let _ = tc.arena.take_device_value(tc.arena.resolve(leaf), FULL); + assert!(tc.is_full_device_evicted(leaf)); +} + +#[test] +fn set_component_device_value_stores_and_restamps_the_lru() { + let mut tc = core(); + tc.register_component_(Arc::new(SwaComponentForTest)); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1, 2], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + assert!(!tc.arena.has_device_value(a, SWA)); + tc.set_component_device_value(tc.arena.node(a).id, SWA, Tensor::from_slice(&[5i64, 6])); + assert!( + tc.arena + .device_value(a, SWA) + .equal(&Tensor::from_slice(&[5i64, 6])) + ); + assert_eq!(tc.evictable_size_(SWA), 2); + assert!(tc.device_lru_list(SWA).in_list(Some(a))); + assert_eq!(tc.device_lru_list(SWA).len(), 1); + assert!(!tc.host_lru_list(SWA).in_list(Some(a))); +} + +#[test] +fn set_component_device_value_migrates_the_node_off_the_host_lru() { + let mut tc = core(); + tc.register_component_(Arc::new(SwaComponentForTest)); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.host_lru_list_mut(SWA).insert_mru(a); + tc.set_component_device_value(tc.arena.node(a).id, SWA, Tensor::from_slice(&[5i64])); + assert!(!tc.host_lru_list(SWA).in_list(Some(a))); + assert_eq!(tc.host_lru_list(SWA).len(), 0); + assert_eq!(tc.device_lru_list(SWA).len(), 1); +} + +#[test] +#[should_panic(expected = "auxiliary components only")] +fn set_component_device_value_rejects_the_base_component() { + let mut tc = core(); + let root = tc.arena.root(); + tc.set_component_device_value( + tc.arena.node(root).id, + BASE_COMPONENT_TYPE, + Tensor::from_slice(&[1i64]), + ); +} + +#[test] +fn collect_full_device_indices_concatenates_in_root_order() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + tc.insert(&insert_params(&vec![1, 2, 3, 4, 5], &[20, 21, 22, 13, 14])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 3, 4, 5])) + .best_match_node_id; + let parent = tc.arena.node(tc.arena.resolve(leaf)).parent(); + let root = tc.arena.root(); + assert!( + tc.collect_full_device_indices(leaf, tc.arena.node(root).id) + .equal(&Tensor::from_slice(&[10i64, 11, 12, 13, 14])) + ); + assert!( + tc.collect_full_device_indices(leaf, tc.arena.node(parent).id) + .equal(&Tensor::from_slice(&[13i64, 14])) + ); + assert_eq!( + tc.collect_full_device_indices(tc.arena.node(root).id, tc.arena.node(root).id) + .numel(), + 0 + ); +} + +#[test] +#[should_panic(expected = "value: Full/device slot has no value")] +fn collect_full_device_indices_panics_on_an_evicted_path() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + tc.insert(&insert_params(&vec![1, 2, 3, 4, 5], &[20, 21, 22, 13, 14])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 3, 4, 5])) + .best_match_node_id; + let parent = tc.arena.node(tc.arena.resolve(leaf)).parent(); + let _ = tc.arena.take_device_value(parent, FULL); + let root = tc.arena.root(); + let _ = tc.collect_full_device_indices(leaf, tc.arena.node(root).id); +} + +#[test] +fn all_values_flatten_spans_namespaces() { + let mut tc = core(); + assert_eq!(tc.all_values_flatten().numel(), 0); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + tc.insert(&insert_params(&vec![1, 2, 3, 4, 5], &[20, 21, 22, 13, 14])); + tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some("chat"), None), + ..insert_params(&vec![7, 8], &[20, 21]) + }); + let (sorted, _) = tc.all_values_flatten().sort(0, /* descending = */ false); + assert!(sorted.equal(&Tensor::from_slice(&[10i64, 11, 12, 13, 14, 20, 21]))); +} + +#[test] +fn collect_all_nodes_visits_every_root_subtree() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some("chat"), None), + ..insert_params(&vec![7, 8], &[20, 21]) + }); + let mut nodes = tc.collect_all_nodes_(); + nodes.sort(); + assert_eq!(nodes.len(), tc.arena.len()); + assert_eq!(nodes, vec![NodeIdx_(0), NodeIdx_(1), NodeIdx_(2)]); +} + +#[test] +fn pretty_format_renders_every_namespace_and_component() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + tc.insert(&insert_params(&vec![1, 2, 3, 4, 5], &[20, 21, 22, 13, 14])); + tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some("chat"), None), + ..insert_params(&vec![7], &[30]) + }); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 3, 4, 5])) + .best_match_node_id; + let _ = tc.arena.take_device_value(tc.arena.resolve(leaf), FULL); + tc.register_component_(Arc::new(SwaComponentForTest)); + tc.arena.node_mut(NodeIdx_(1)).values[SWA.idx()].value = Some(Tensor::from_slice(&[0i64])); + // Sibling render order follows HashMap iteration, so pin the line set. + let mut lines: Vec = tc.pretty_format_().lines().map(str::to_string).collect(); + lines.sort(); + let mut expected: Vec = [ + " [0] 0 full_lock=1 Full=no Swa=no", + " [3] 1 full_lock=0 Full=yes Swa=no", + " [1] 3 full_lock=0 Full=yes Swa=yes", + " [2] 2 full_lock=0 Full=no Swa=no", + ] + .map(str::to_string) + .to_vec(); + expected.sort(); + assert_eq!(lines, expected); +} + +// A healthy multi-namespace tree with a split for the sanity pins. +fn sane_tree() -> UnifiedTreeCore> { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + tc.insert(&insert_params(&vec![1, 2, 3, 4, 5], &[20, 21, 22, 13, 14])); + tc.insert(&insert_params(&vec![1, 2, 9], &[30, 31, 39])); + tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some("chat"), None), + ..insert_params(&vec![7, 8], &[40, 41]) + }); + tc +} + +#[test] +fn sanity_check_passes_on_a_healthy_tree() { + let mut tc = sane_tree(); + tc.sanity_check(&[], &[]); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + tc.inc_lock_ref(leaf); + tc.sanity_check(&[(1, leaf)], &[(2, leaf)]); + tc.dec_lock_ref( + tc.arena.node(tc.arena.resolve(leaf)).id, + /* params = */ None, + /* skip_swa = */ false, + ); + tc.sanity_check(&[], &[]); +} + +#[test] +fn sanity_check_passes_after_the_eviction_walk() { + let mut tc = sane_tree(); + let mut tracker = HashMap::from([(FULL, 0)]); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.evict_device_start(FULL, /* request_cnt = */ 100); + loop { + let (leaf, step) = tc.evict_device_next_node(FULL, &tracker); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + let Some(leaf) = leaf else { break }; + let (_, step) = tc.evict_device_leaf(leaf, /* is_write_back = */ false); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + } + tc.evict_device_end(FULL); + tc.sanity_check(&[], &[]); + // The emptied "chat" namespace leaves nothing behind; only the root survives. + assert_eq!(tc.arena.len(), 1); + assert!(!tc.arena.namespace_exists(Some("chat"))); +} + +#[test] +#[should_panic(expected = "D-leaf missing")] +fn sanity_check_detects_a_missing_device_leaf() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + tc.evictable_device_leaves.discard(tc.arena.resolve(leaf)); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "D-leaf extra")] +fn sanity_check_detects_an_extra_device_leaf() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + let parent = tc.arena.node(tc.arena.resolve(leaf)).parent(); + tc.evictable_device_leaves.add(parent); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "[Size]")] +fn sanity_check_detects_size_drift() { + let mut tc = sane_tree(); + tc.component_state_mut(FULL).evictable_size = 999; + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "dead: no Full device and no Full host")] +fn sanity_check_detects_a_dead_node() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + let _ = tc.arena.take_device_value(tc.arena.resolve(leaf), FULL); + tc.sanity_check(&[], &[]); +} + +#[test] +fn try_sanity_check_returns_a_dead_node_error() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + let leaf_idx = tc.arena.resolve(leaf); + let _ = tc.arena.take_device_value(leaf_idx, FULL); + + let error = tc.try_sanity_check(&[], &[]).unwrap_err(); + assert!(error.starts_with("Sanity check FAILED")); + assert!(error.contains(&format!( + "node {leaf_idx} dead: no Full device and no Full host" + ))); +} + +#[test] +#[should_panic(expected = "device present but parent")] +fn sanity_check_detects_an_evicted_parent_prefix() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + let parent = tc.arena.node(tc.arena.resolve(leaf)).parent(); + let _ = tc.arena.take_device_value(parent, FULL); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "evicted but lock_ref")] +fn sanity_check_detects_a_locked_tombstone() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + tc.inc_lock_ref(leaf); + let _ = tc.arena.take_device_value(tc.arena.resolve(leaf), FULL); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "Full device LRU not empty")] +fn sanity_check_detects_full_lru_pollution() { + let mut tc = sane_tree(); + tc.device_lru_list_mut(FULL).insert_mru(NodeIdx_(0)); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "device LRU mismatch at node")] +fn sanity_check_detects_an_aux_lru_mismatch() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + tc.register_component_(Arc::new(SwaComponentForTest)); + tc.arena.node_mut(tc.arena.resolve(leaf)).values[SWA.idx()].value = + Some(Tensor::from_slice(&[0i64, 0, 0])); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "write_through node 7 not in tree")] +fn sanity_check_detects_an_untracked_ongoing_node() { + let tc = sane_tree(); + tc.sanity_check(&[(7, 999)], &[]); +} + +#[test] +#[should_panic(expected = "load_back node 8 lock_ref=0")] +fn sanity_check_detects_an_unlocked_ongoing_node() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + tc.sanity_check(&[], &[(8, leaf)]); +} + +#[test] +#[should_panic(expected = "[Root] root 0 holds a Full device value")] +fn sanity_check_detects_a_valued_root() { + let mut tc = sane_tree(); + tc.arena.node_mut(NodeIdx_(0)).values[FULL.idx()].value = Some(Tensor::from_slice(&[0i64])); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "[Root] root 0 Full lock_ref=0")] +fn sanity_check_detects_an_unlocked_root() { + let mut tc = sane_tree(); + tc.arena.node_mut(NodeIdx_(0)).values[FULL.idx()].lock_ref = 0; + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "[Root] root 0 has a parent pointer")] +fn sanity_check_detects_a_parented_root() { + let mut tc = sane_tree(); + tc.arena.node_mut(NodeIdx_(0)).parent = Some(NodeIdx_(1)); + tc.sanity_check(&[], &[]); +} + +// A tree with the Swa stub registered and the root's Swa lock backfilled. +fn swa_locked_roots_tree() -> UnifiedTreeCore> { + let mut tc = sane_tree(); + tc.register_component_(Arc::new(SwaComponentForTest)); + tc.arena.node_mut(tc.arena.root()).values[SWA.idx()].lock_ref = 1; + tc +} + +#[test] +#[should_panic(expected = "[Root] root 0 holds a Swa device value")] +fn sanity_check_detects_an_aux_device_valued_root() { + let mut tc = swa_locked_roots_tree(); + tc.arena.node_mut(NodeIdx_(0)).values[SWA.idx()].value = Some(Tensor::from_slice(&[0i64])); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "[Root] root 0 holds a Swa host value")] +fn sanity_check_detects_an_aux_host_valued_root() { + let mut tc = swa_locked_roots_tree(); + tc.arena + .node_mut(NodeIdx_(0)) + .state_mut_(ValueSlotIdx::host(SWA)) + .value = Some(Tensor::from_slice(&[0i64])); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "[Root] root 0 Swa lock_ref=0")] +fn sanity_check_detects_an_unlocked_aux_root() { + let mut tc = swa_locked_roots_tree(); + tc.arena.node_mut(NodeIdx_(0)).values[SWA.idx()].lock_ref = 0; + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "[Tree] child")] +fn sanity_check_detects_a_broken_parent_pointer() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + let leaf_idx = tc.arena.resolve(leaf); + tc.arena.node_mut(leaf_idx).parent = Some(NodeIdx_(0)); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "device present but Full.value=None")] +fn sanity_check_detects_aux_device_without_full() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + tc.register_component_(Arc::new(SwaComponentForTest)); + tc.arena.node_mut(tc.arena.resolve(leaf)).values[SWA.idx()].value = + Some(Tensor::from_slice(&[0i64])); + let _ = tc.arena.take_device_value(tc.arena.resolve(leaf), FULL); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "host present but Full.host_value=None")] +fn sanity_check_detects_aux_host_without_full_host() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + tc.register_component_(Arc::new(SwaComponentForTest)); + tc.arena + .node_mut(tc.arena.resolve(leaf)) + .state_mut_(ValueSlotIdx::host(SWA)) + .value = Some(Tensor::from_slice(&[0i64])); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "backed up but parent")] +fn sanity_check_detects_an_unbacked_parent_prefix() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + tc.arena + .node_mut(tc.arena.resolve(leaf)) + .state_mut_(ValueSlotIdx::host(FULL)) + .value = Some(Tensor::from_slice(&[30i64])); + tc.sanity_check(&[], &[]); +} + +#[test] +fn sanity_check_accepts_a_write_back_child_backed_up_before_its_parent() { + // Write-back backs up leaf-first, so an unbacked parent is legal. + let params = CacheInitParams { + is_write_back: true, + ..Default::default() + }; + let mut tc: UnifiedTreeCore> = UnifiedTreeCore::new(params, vec![FULL]); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + tc.insert(&insert_params(&vec![1, 2, 3, 4, 5], &[20, 21, 22, 13, 14])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 3, 4, 5])) + .best_match_node_id; + tc.arena + .node_mut(tc.arena.resolve(leaf)) + .state_mut_(ValueSlotIdx::host(FULL)) + .value = Some(Tensor::from_slice(&[13i64, 14])); + // Register the host value set directly by the test. + tc.update_full_coexisting_host_tracking_(tc.arena.resolve(leaf)); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "full_lock=0 < Swa_lock=5")] +fn sanity_check_detects_an_aux_lock_above_full() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + tc.register_component_(Arc::new(SwaComponentForTest)); + tc.arena.node_mut(tc.arena.resolve(leaf)).values[SWA.idx()].lock_ref = 5; + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "H-leaf missing")] +fn sanity_check_detects_a_missing_host_leaf() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + let parent = tc.arena.node(tc.arena.resolve(leaf)).parent(); + let _ = tc.arena.take_device_value(tc.arena.resolve(leaf), FULL); + tc.arena + .node_mut(tc.arena.resolve(leaf)) + .state_mut_(ValueSlotIdx::host(FULL)) + .value = Some(Tensor::from_slice(&[30i64])); + tc.arena + .node_mut(parent) + .state_mut_(ValueSlotIdx::host(FULL)) + .value = Some(Tensor::from_slice(&[10i64, 11])); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "H-leaf extra")] +fn sanity_check_detects_an_extra_host_leaf() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + tc.evictable_host_leaves.add(tc.arena.resolve(leaf)); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "in both sets")] +fn sanity_check_detects_a_leaf_in_both_sets() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + tc.evictable_host_leaves.add(tc.arena.resolve(leaf)); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "stale nodes in device_leaves")] +fn sanity_check_detects_a_stale_device_leaf() { + let mut tc = sane_tree(); + tc.evictable_device_leaves.add(NodeIdx_(999)); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "stale nodes in host_leaves")] +fn sanity_check_detects_a_stale_host_leaf() { + let mut tc = sane_tree(); + tc.evictable_host_leaves.add(NodeIdx_(999)); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "Full host LRU not empty")] +fn sanity_check_detects_full_host_lru_pollution() { + let mut tc = sane_tree(); + tc.host_lru_list_mut(FULL).insert_mru(NodeIdx_(0)); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "host LRU mismatch at node")] +fn sanity_check_detects_an_aux_host_lru_mismatch() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + tc.register_component_(Arc::new(SwaComponentForTest)); + let leaf_idx2 = tc.arena.resolve(leaf); + tc.host_lru_list_mut(SWA).insert_mru(leaf_idx2); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "in both device and host LRU")] +fn sanity_check_detects_an_aux_node_in_both_lrus() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + tc.register_component_(Arc::new(SwaComponentForTest)); + tc.arena.node_mut(tc.arena.resolve(leaf)).values[SWA.idx()].value = + Some(Tensor::from_slice(&[0i64, 0, 0])); + let leaf_idx = tc.arena.resolve(leaf); + tc.device_lru_list_mut(SWA).insert_mru(leaf_idx); + tc.host_lru_list_mut(SWA).insert_mru(leaf_idx); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "device LRU: tree=0 != lru=1")] +fn sanity_check_detects_device_lru_length_drift() { + let mut tc = sane_tree(); + tc.register_component_(Arc::new(SwaComponentForTest)); + tc.device_lru_list_mut(SWA).insert_mru(NodeIdx_(999)); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "host LRU: tree=0 != lru=1")] +fn sanity_check_detects_host_lru_length_drift() { + let mut tc = sane_tree(); + tc.register_component_(Arc::new(SwaComponentForTest)); + tc.host_lru_list_mut(SWA).insert_mru(NodeIdx_(999)); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "[device][Swa] list=0 != len=1")] +fn sanity_check_wires_the_device_list_integrity_walk() { + let mut tc = sane_tree(); + tc.register_component_(Arc::new(SwaComponentForTest)); + tc.device_lru_list_mut(SWA).bump_len_for_test(); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "[host][Swa] list=0 != len=1")] +fn sanity_check_wires_the_host_list_integrity_walk() { + let mut tc = sane_tree(); + tc.register_component_(Arc::new(SwaComponentForTest)); + tc.host_lru_list_mut(SWA).bump_len_for_test(); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "[device][Full] list=0 != len=1")] +fn sanity_check_wires_the_full_device_list_integrity_walk() { + let mut tc = sane_tree(); + tc.device_lru_list_mut(FULL).bump_len_for_test(); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "protected=999")] +fn sanity_check_detects_protected_size_drift() { + let mut tc = sane_tree(); + tc.component_state_mut(FULL).protected_size = 999; + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "write_through node 9 lock_ref=0")] +fn sanity_check_detects_an_unlocked_write_through() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + tc.sanity_check(&[(9, leaf)], &[]); +} + +#[test] +#[should_panic(expected = "load_back node 10 not in tree")] +fn sanity_check_detects_an_untracked_load_back() { + let tc = sane_tree(); + tc.sanity_check(&[], &[(10, 999)]); +} + +#[test] +fn match_prefix_on_an_unknown_namespace_allocates_nothing() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + let arena_len = tc.arena.len(); + let result = tc.match_prefix(&MatchPrefixParams { + key: &vec![1, 2, 3], + namespace: KeyNamespaceRef::new(Some("ghost"), None), + }); + assert_eq!(result.device_indices.numel(), 0); + assert_eq!(tc.arena.len(), arena_len); + // The empty result anchors at the default root (the namespace has no root). + let default_root = tc.arena.root(); + assert_eq!(result.best_match_node_id, tc.arena.node(default_root).id); +} + +#[test] +fn refresh_dispatches_fire_per_walk_phase_in_a_namespace() { + let mut tc = core(); + let recorder = Arc::new(RecordingComponentForTest::default()); + tc.register_component_(recorder.clone()); + tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some("chat"), None), + ..insert_params(&vec![7, 8], &[40, 41]) + }); + // The deeper insert walks down through the existing [7,8] node. + tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some("chat"), None), + ..insert_params(&vec![7, 8, 9], &[40, 41, 42]) + }); + let leaf = tc + .match_prefix(&MatchPrefixParams { + key: &vec![7, 8], + namespace: KeyNamespaceRef::new(Some("chat"), None), + }) + .best_match_node_id; + let refreshes = recorder.refreshes.lock().unwrap(); + assert!(!refreshes.is_empty()); + assert!( + refreshes + .iter() + .any(|&(phase, node)| phase == LRURefreshPhase::Walkdown + && node == tc.arena.resolve(leaf)) + ); + assert!(refreshes.iter().any( + |&(phase, node)| phase == LRURefreshPhase::InsertEnd && node == tc.arena.resolve(leaf) + )); + assert!( + refreshes + .iter() + .any(|&(phase, node)| phase == LRURefreshPhase::MatchEnd + && node == tc.arena.resolve(leaf)) + ); +} + +#[test] +#[should_panic(expected = "orphaned live nodes")] +fn sanity_check_detects_an_orphaned_node() { + let mut tc = sane_tree(); + let root = tc.arena.root(); + // A parented node missing from its parent's child map is unreachable. + tc.new_node_( + /* key = */ vec![99], + root, + /* priority = */ 0, + /* hit_count = */ 0, + /* creation_counter = */ None, + /* extra_key = */ None, + ); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "not mapped under its own child key")] +fn sanity_check_detects_a_reverse_map_mismatch() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + let leaf_idx3 = tc.arena.resolve(leaf); + let parent = tc.arena.node(leaf_idx3).parent(); + let key = tc.arena.node(leaf_idx3).key.child_key(1); + let parent_node = tc.arena.node_mut(parent); + parent_node.children.remove(&(KeyNamespace::default(), key)); + parent_node + .children + .insert((KeyNamespace::default(), vec![99]), leaf_idx3); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "Full value length")] +fn sanity_check_detects_a_value_length_mismatch() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + tc.arena.node_mut(tc.arena.resolve(leaf)).values[FULL.idx()].value = + Some(Tensor::from_slice(&[7i64, 8, 9, 10])); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "Full host value length")] +fn sanity_check_detects_a_host_value_length_mismatch() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + let parent = tc.arena.node(tc.arena.resolve(leaf)).parent(); + tc.arena + .node_mut(parent) + .state_mut_(ValueSlotIdx::host(FULL)) + .value = Some(Tensor::from_slice(&[10i64, 11])); + tc.arena + .node_mut(tc.arena.resolve(leaf)) + .state_mut_(ValueSlotIdx::host(FULL)) + .value = Some(Tensor::from_slice(&[7i64, 8, 9, 10])); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "has an empty key")] +fn sanity_check_detects_an_empty_key() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + tc.arena.node_mut(tc.arena.resolve(leaf)).key = vec![]; + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "key is not page-aligned")] +fn sanity_check_detects_an_unaligned_key() { + let mut tc = page2_core(); + tc.insert(&insert_params(&vec![1, 2, 3, 4], &[10, 11, 12, 13])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 3, 4])) + .best_match_node_id; + tc.arena.node_mut(tc.arena.resolve(leaf)).key = vec![1]; + tc.sanity_check(&[], &[]); +} + +// Corrupt the [1,2,9] leaf's child map to point back at its own parent. +fn cyclic_child_map_tree() -> UnifiedTreeCore> { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + let parent = tc.arena.node(tc.arena.resolve(leaf)).parent(); + tc.arena + .node_mut(tc.arena.resolve(leaf)) + .children + .insert((KeyNamespace::default(), vec![50]), parent); + tc +} + +#[test] +fn collect_all_nodes_terminates_on_a_cyclic_child_map() { + let tc = cyclic_child_map_tree(); + let nodes = tc.collect_all_nodes_(); + assert_eq!(nodes.len(), tc.arena.len()); +} + +#[test] +#[should_panic(expected = "[Tree] child")] +fn sanity_check_reports_a_cyclic_child_map_without_hanging() { + let tc = cyclic_child_map_tree(); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "host LRU mismatch")] +fn sanity_check_detects_a_host_locked_value_missing_from_the_lru() { + let mut tc = sane_tree(); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + let parent = tc.arena.node(tc.arena.resolve(leaf)).parent(); + tc.register_component_(Arc::new(SwaComponentForTest)); + // The arena was built Full-only; give the root the stub's lock too. + tc.arena.node_mut(tc.arena.root()).values[SWA.idx()].lock_ref = 1; + tc.arena + .node_mut(parent) + .state_mut_(ValueSlotIdx::host(FULL)) + .value = Some(Tensor::from_slice(&[10i64, 11])); + let leaf_node = tc.arena.node_mut(tc.arena.resolve(leaf)); + leaf_node.state_mut_(ValueSlotIdx::host(FULL)).value = Some(Tensor::from_slice(&[30i64])); + leaf_node.state_mut_(ValueSlotIdx::host(SWA)).value = Some(Tensor::from_slice(&[30i64])); + leaf_node.state_mut_(ValueSlotIdx::host(SWA)).lock_ref = 1; + tc.sanity_check(&[], &[]); +} + +// A backed-up leaf whose unlocked Swa value is host-only (no device value). +fn host_only_aux_leaf(tc: &mut UnifiedTreeCore>) -> NodeIdx_ { + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 9])) + .best_match_node_id; + let parent = tc.arena.node(tc.arena.resolve(leaf)).parent(); + tc.register_component_(Arc::new(SwaComponentForTest)); + // The arena was built Full-only; give the root the stub's lock too. + tc.arena.node_mut(tc.arena.root()).values[SWA.idx()].lock_ref = 1; + tc.arena + .node_mut(parent) + .state_mut_(ValueSlotIdx::host(FULL)) + .value = Some(Tensor::from_slice(&[10i64, 11])); + let leaf_node = tc.arena.node_mut(tc.arena.resolve(leaf)); + leaf_node.state_mut_(ValueSlotIdx::host(FULL)).value = Some(Tensor::from_slice(&[30i64])); + leaf_node.state_mut_(ValueSlotIdx::host(SWA)).value = Some(Tensor::from_slice(&[30i64])); + // Register the host values set directly by the test. + tc.update_full_coexisting_host_tracking_(parent); + tc.update_full_coexisting_host_tracking_(tc.arena.resolve(leaf)); + tc.arena.resolve(leaf) +} + +#[test] +fn sanity_check_accepts_an_unlocked_host_only_value_in_the_lru() { + let mut tc = sane_tree(); + let leaf = host_only_aux_leaf(&mut tc); + tc.host_lru_list_mut(SWA).insert_mru(leaf); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "host LRU mismatch")] +fn sanity_check_detects_a_host_only_value_missing_from_the_lru() { + let mut tc = sane_tree(); + host_only_aux_leaf(&mut tc); + tc.sanity_check(&[], &[]); +} + +#[test] +#[should_panic(expected = "EvictLayer::All is not a single layer")] +fn for_each_component_lru_rejects_the_all_layer() { + let mut tc = core(); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.register_component_(Arc::new(SwaComponentForTest)); + tc.for_each_component_lru_( + a, + &mut |_, _| {}, + EvictLayer::All, + /* skip_existing = */ false, + ); +} + +#[test] +fn insert_unevicts_a_tombstoned_deep_node() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + tc.insert(&insert_params(&vec![1, 2, 3, 4, 5], &[20, 21, 22, 13, 14])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2, 3, 4, 5])) + .best_match_node_id; + let _ = tc.arena.take_device_value(tc.arena.resolve(leaf), FULL); + tc.component_state_mut(FULL).evictable_size = 3; + tc.evictable_device_leaves.discard(tc.arena.resolve(leaf)); + let result = tc.insert(&insert_params(&vec![1, 2, 3, 4, 5], &[30, 31, 32, 33, 34])); + assert_eq!(result.prefix_len, 5); + // The revived leaf takes its own span of the fresh KV, not the key head. + assert!( + tc.arena + .device_value(tc.arena.resolve(leaf), FULL) + .equal(&Tensor::from_slice(&[33i64, 34])) + ); + let [CacheAction::FreeDeviceKV(freed)] = result.cache_actions.as_slice() else { + panic!( + "expected one FreeDeviceKV action, got {:?}", + action_kinds(&result.cache_actions) + ); + }; + assert!(freed[0].equal(&Tensor::from_slice(&[30i64, 31, 32]))); +} + +#[test] +fn insert_ragged_key_onto_an_existing_prefix_page_size_two() { + let mut tc = page2_core(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + let result = tc.insert(&insert_params(&vec![1, 2, 3], &[20, 21, 22])); + assert_eq!(result.prefix_len, 2); + // The ragged tail never enters: no new leaf, the aligned span is duplicate. + assert_eq!(tc.arena.len(), 2); + let [CacheAction::FreeDeviceKV(freed)] = result.cache_actions.as_slice() else { + panic!( + "expected one FreeDeviceKV action, got {:?}", + action_kinds(&result.cache_actions) + ); + }; + assert!(freed[0].equal(&Tensor::from_slice(&[20i64, 21]))); +} + +#[test] +fn insert_ragged_key_traverses_a_node_before_the_tail_page_size_two() { + let mut tc = page2_core(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + tc.insert(&insert_params(&vec![1, 2, 3, 4], &[20, 21, 12, 13])); + let result = tc.insert(&insert_params(&vec![1, 2, 3, 4, 5], &[30, 31, 32, 33, 34])); + assert_eq!(result.prefix_len, 4); + assert_eq!(tc.arena.len(), 3); + let [ + CacheAction::FreeDeviceKV(freed_head), + CacheAction::FreeDeviceKV(freed_tail), + ] = result.cache_actions.as_slice() + else { + panic!( + "expected one FreeDeviceKV per walked node, got {:?}", + action_kinds(&result.cache_actions) + ); + }; + assert!(freed_head[0].equal(&Tensor::from_slice(&[30i64, 31]))); + assert!(freed_tail[0].equal(&Tensor::from_slice(&[32i64, 33]))); +} + +#[test] +fn match_ragged_query_stops_at_the_aligned_window_page_size_two() { + let mut tc = page2_core(); + tc.insert(&insert_params(&vec![1, 2, 3, 4], &[10, 11, 12, 13])); + // The node key runs past the query's aligned window; the ragged atom matches too. + let result = tc.match_prefix(&match_params(&vec![1, 2, 3])); + assert!( + result + .device_indices + .equal(&Tensor::from_slice(&[10i64, 11])) + ); +} + +#[test] +fn begin_insert_empty_key_completes_in_one_step() { + let mut tc = core(); + let step = tc.begin_insert(&insert_params(&vec![], &[])); + assert!(step.actions.is_empty()); + let result = step.result.expect("an empty insert completes immediately"); + assert_eq!(result.prefix_len, 0); + assert!(result.mamba_exist); + assert!(!tc.has_ongoing_insert()); +} + +#[test] +fn deferrable_dup_frees_ride_the_final_step_without_suspension() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + let step = tc.begin_insert(&insert_params(&vec![1, 2, 3, 4], &[20, 21, 22, 13])); + assert_eq!(action_kinds(&step.actions), vec!["FreeDeviceKV"]); + let result = step + .result + .expect("a deferrable-only walk completes in one step"); + assert_eq!(result.prefix_len, 3); + assert!(result.cache_actions.is_empty()); + assert!(!tc.has_ongoing_insert()); + assert!(tc.end_insert().is_empty()); +} + +fn suspended_walk_core() -> (UnifiedTreeCore>, NodeIdx_, InsertStepResult) { + let mut tc: UnifiedTreeCore> = UnifiedTreeCore::new( + CacheInitParams { + write_through_threshold: 2, + ..Default::default() + }, + vec![FULL], + ); + tc.set_hicache_enabled(); + tc.insert(&insert_params(&vec![1, 2, 3], &[10, 11, 12])); + let a = tc + .match_prefix(&match_params(&vec![1, 2, 3])) + .best_match_node_id; + let step = tc.begin_insert(&insert_params(&vec![1, 2, 3, 4, 5], &[20, 21, 22, 13, 14])); + let a_idx = tc.arena.resolve(a); + (tc, a_idx, step) +} + +#[test] +fn walk_backup_crossing_suspends_then_resume_completes() { + let (mut tc, a, step) = suspended_walk_core(); + assert!(step.result.is_none()); + assert!(tc.has_ongoing_insert()); + assert_eq!( + action_kinds(&step.actions), + vec!["FreeDeviceKV", "BackupKV"] + ); + let CacheAction::BackupKV(backup) = &step.actions[1] else { + unreachable!(); + }; + assert_eq!(backup.node_ids, vec![tc.arena.node(a).id]); + + let done = tc.resume_insert(); + assert!(done.actions.is_empty()); + let result = done.result.expect("the resumed walk completes"); + assert_eq!(result.prefix_len, 3); + assert!(!tc.has_ongoing_insert()); + assert!(tc.end_insert().is_empty()); + tc.sanity_check(&[], &[]); +} + +#[test] +fn try_begin_insert_rejects_a_concurrent_walk() { + let (mut tc, _, step) = suspended_walk_core(); + assert!(step.result.is_none()); + assert!(matches!( + tc.try_begin_insert(&insert_params(&vec![9], &[90])), + Err(TreeCoreRuntimeError::ConcurrentInsertWalk) + )); + assert!(matches!( + tc.try_insert(&insert_params(&vec![9], &[90])), + Err(TreeCoreRuntimeError::ConcurrentInsertWalk) + )); +} + +#[test] +fn try_resume_insert_rejects_a_missing_walk() { + let mut tc = core(); + assert!(matches!( + tc.try_resume_insert(), + Err(TreeCoreRuntimeError::NoInFlightInsert) + )); +} + +#[test] +fn end_insert_aborts_the_suspended_walk() { + let (mut tc, _, step) = suspended_walk_core(); + assert!(step.result.is_none()); + // The barrier flushed everything pending; the abort drain is empty. + assert!(tc.end_insert().is_empty()); + assert!(!tc.has_ongoing_insert()); + assert!(tc.end_insert().is_empty()); + // The single-flight slot is clear: a fresh insert starts normally. + let result = tc.insert(&insert_params(&vec![9], &[90])); + assert_eq!(result.prefix_len, 0); +} + +#[test] +fn resume_insert_completes_after_an_on_path_host_leaf_is_evicted() { + let mut tc: UnifiedTreeCore> = UnifiedTreeCore::new( + CacheInitParams { + write_through_threshold: 2, + ..Default::default() + }, + vec![FULL], + ); + tc.set_hicache_enabled(); + tc.insert(&insert_params(&vec![1, 2, 3, 4], &[10, 11, 12, 13])); + let top = tc + .match_prefix(&match_params(&vec![1, 2, 3, 4])) + .best_match_node_id; + tc.insert(&insert_params( + &vec![1, 2, 3, 4, 5, 6, 7, 8], + &[20, 21, 22, 23, 24, 25, 26, 27], + )); + let h_leaf = tc + .match_prefix(&match_params(&vec![1, 2, 3, 4, 5, 6, 7, 8])) + .best_match_node_id; + let h_leaf_idx = tc.arena.resolve(h_leaf); + tc.commit_backup( + h_leaf, + Tensor::from_slice(&[104i64, 105, 106, 107]), + HashMap::new(), + ); + demote_node(&mut tc, h_leaf_idx); + let step = tc.begin_insert(&insert_params( + &vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], + &[20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31], + )); + assert!(step.result.is_none()); + assert_eq!( + action_kinds(&step.actions), + vec!["FreeDeviceKV", "BackupKV"] + ); + // The barrier's backup host-evicts the on-path H-leaf before committing. + let mut tracker = HashMap::from([(FULL, 0)]); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.evict_host_leaf_(tc.arena.resolve(h_leaf), &mut tracker, &mut df, &mut hf); + assert_eq!(tracker[&FULL], 4); + tc.commit_backup( + top, + Tensor::from_slice(&[100i64, 101, 102, 103]), + HashMap::new(), + ); + let done = tc.resume_insert(); + let result = done.result.expect("the resumed walk completes"); + assert_eq!(result.prefix_len, 4); + assert!(!tc.has_ongoing_insert()); + assert!(tc.arena.try_resolve(h_leaf).is_none()); + // The recreated suffix is top's single child, spanning the whole gap. + let top_idx = tc.arena.resolve(top); + assert_eq!(tc.arena.node(top_idx).children.len(), 1); + let suffix = *tc.arena.node(top_idx).children.values().next().unwrap(); + assert_eq!(tc.arena.node(suffix).key, vec![5, 6, 7, 8, 9, 10, 11, 12]); + tc.sanity_check(&[], &[]); +} + +#[test] +fn aborted_barrier_crossing_refires_on_the_next_insert() { + let (mut tc, a, step) = suspended_walk_core(); + assert!(step.result.is_none()); + assert!(tc.end_insert().is_empty()); + assert!(!tc.has_ongoing_insert()); + // The abort never committed the backup, so the same crossing fires again. + let step = tc.begin_insert(&insert_params(&vec![1, 2, 3, 4, 5], &[20, 21, 22, 13, 14])); + assert!(step.result.is_none()); + let backups: Vec<_> = step + .actions + .iter() + .filter_map(|action| match action { + CacheAction::BackupKV(backup) => Some(backup.node_ids.clone()), + _ => None, + }) + .collect(); + assert_eq!(backups, vec![vec![tc.arena.node(a).id]]); + tc.commit_backup( + tc.arena.node(a).id, + Tensor::from_slice(&[100i64, 101, 102]), + HashMap::new(), + ); + let done = tc.resume_insert(); + assert_eq!( + done.result.expect("the resumed walk completes").prefix_len, + 3 + ); + tc.sanity_check(&[], &[]); +} + +#[test] +fn one_insert_walk_fires_two_crossings_around_a_backuped_middle() { + let mut tc = core(); + tc.set_hicache_enabled(); + tc.insert(&insert_params(&vec![1, 2, 3, 4], &[10, 11, 12, 13])); + let top = tc + .match_prefix(&match_params(&vec![1, 2, 3, 4])) + .best_match_node_id; + // Build a backed-up, device-evicted middle below the unbacked top. A + // write-through host refill below an unbacked parent is now rejected. + tc.insert(&insert_params( + &vec![1, 2, 3, 4, 5, 6, 7, 8], + &[20, 21, 22, 23, 24, 25, 26, 27], + )); + let middle = tc + .match_prefix(&match_params(&vec![1, 2, 3, 4, 5, 6, 7, 8])) + .best_match_node_id; + let middle_idx = tc.arena.resolve(middle); + tc.commit_backup( + middle, + Tensor::from_slice(&[104i64, 105, 106, 107]), + HashMap::new(), + ); + demote_node(&mut tc, middle_idx); + + // The device insert restores the middle and adds the unbacked deep leaf. + tc.insert(&insert_params( + &vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], + &[20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31], + )); + let deep = tc + .match_prefix(&match_params(&vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])) + .best_match_node_id; + tc.write_through_threshold = 2; + let result = tc.insert(&insert_params( + &vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], + &[ + 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, + ], + )); + // Both crossings fire in walk order; the backuped middle joins neither chain. + let backups: Vec<_> = result + .cache_actions + .iter() + .filter_map(|action| match action { + CacheAction::BackupKV(backup) => Some(backup.node_ids.clone()), + _ => None, + }) + .collect(); + assert_eq!(backups, vec![vec![top], vec![deep]]); + assert!(tc.arena.node(tc.arena.resolve(middle)).backuped()); +} + +#[test] +fn full_kv_hit_length_counts_the_split_fragment() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2, 3, 4], &[10, 11, 12, 13])); + let result = tc.match_prefix(&match_params(&vec![1, 2, 99])); + // The mid-node partial match splits the node; the fragment still counts. + assert_eq!(result.full_kv_hit_length, 2); + assert_eq!(result.device_indices.size()[0], 2); +} + +#[test] +fn dec_evictable_size_updates_only_the_addressed_component() { + let mut tc = core(); + tc.component_state_mut(FULL).evictable_size = 5; + tc.component_state_mut(SWA).evictable_size = 7; + tc.dec_evictable_size(FULL, 2); + assert_eq!(tc.evictable_size_(FULL), 3); + assert_eq!(tc.evictable_size_(SWA), 7); +} + +#[test] +fn dec_evictable_size_to_exactly_zero() { + let mut tc = core(); + tc.component_state_mut(FULL).evictable_size = 4; + tc.dec_evictable_size(FULL, 4); + assert_eq!(tc.evictable_size_(FULL), 0); +} + +#[test] +#[should_panic(expected = "dec_evictable_size: Full evictable size underflow")] +fn dec_evictable_size_panics_on_underflow() { + let mut tc = core(); + tc.component_state_mut(FULL).evictable_size = 1; + tc.dec_evictable_size(FULL, 2); +} + +#[test] +fn size_helpers_move_tokens_for_the_addressed_component() { + let mut tc = core(); + tc.inc_evictable_size(FULL, 4); + tc.inc_protected_size(FULL, 3); + assert_eq!(tc.evictable_size_(FULL), 4); + assert_eq!(tc.protected_size_(FULL), 3); + assert_eq!(tc.evictable_size_(SWA), 0); + assert_eq!(tc.protected_size_(SWA), 0); + tc.dec_protected_size(FULL, 2); + assert_eq!(tc.protected_size_(FULL), 1); +} + +#[test] +#[should_panic(expected = "dec_protected_size: Full protected size underflow")] +fn dec_protected_size_panics_on_underflow() { + let mut tc = core(); + tc.dec_protected_size(FULL, 1); +} + +#[test] +fn component_state_accessors_address_the_given_component() { + let mut tc = core(); + tc.component_state_mut(FULL).evictable_size = 5; + assert_eq!(tc.component_state(FULL).evictable_size, 5); + assert_eq!(tc.component_state(SWA).evictable_size, 0); +} + +#[test] +fn evict_walk_lifecycle_tracks_the_bookkeeping() { + let mut tc = core(); + tc.set_evict_device_start(FULL, /* request_cnt = */ 7); + assert!(tc.component_state(FULL).is_evict_device_ongoing); + assert_eq!(tc.component_state(FULL).evict_device_request_cnt, 7); + assert_eq!(tc.component_state(FULL).evict_device_cursor, None); + tc.set_evict_device_end(FULL); + assert!(!tc.component_state(FULL).is_evict_device_ongoing); +} + +#[test] +#[should_panic(expected = "Full device eviction already in progress")] +fn set_evict_device_start_panics_when_already_ongoing() { + let mut tc = core(); + tc.set_evict_device_start(FULL, /* request_cnt = */ 1); + tc.set_evict_device_start(FULL, /* request_cnt = */ 1); +} + +#[test] +#[should_panic(expected = "Full device eviction not started")] +fn set_evict_device_end_panics_before_a_walk() { + let mut tc = core(); + tc.set_evict_device_end(FULL); +} + +#[test] +fn lru_list_accessors_address_the_given_component_lists() { + let mut tc = core(); + let root = tc.arena.root(); + let a = tc + .arena + .alloc_child( + root, + /* key = */ vec![1], + /* priority = */ 0, + /* extra_key = */ None, + ) + .unwrap(); + tc.device_lru_list_mut(FULL).insert_mru(a); + assert!(tc.device_lru_list(FULL).in_list(Some(a))); + assert!(!tc.host_lru_list(FULL).in_list(Some(a))); + // A disabled component's list exists but stays empty and independent. + assert!(!tc.device_lru_list(SWA).in_list(Some(a))); + assert_eq!(tc.device_lru_list(SWA).len(), 0); + // The mutable host accessor addresses the same per-component list. + tc.host_lru_list_mut(SWA).insert_mru(a); + assert!(tc.host_lru_list(SWA).in_list(Some(a))); + assert!(!tc.host_lru_list(FULL).in_list(Some(a))); +} + +#[test] +fn reset_invalidates_every_prior_handle() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + let old_root = tc.root_node_handle(/* extra_key = */ None); + let old_leaf = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + tc.reset(); + // Handles are never re-minted, so pre-reset ones miss instead of aliasing. + assert!(tc.arena.try_resolve(old_root).is_none()); + assert!(tc.arena.try_resolve(old_leaf).is_none()); + let new_root = tc.root_node_handle(/* extra_key = */ None); + assert_ne!(new_root, old_root); + assert_eq!(tc.arena.resolve(new_root), tc.arena.root()); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + assert_eq!( + tc.match_prefix(&match_params(&vec![1, 2])) + .device_indices + .size()[0], + 2 + ); +} + +#[test] +#[should_panic(expected = "is not enabled")] +fn component_has_host_value_only_panics_on_a_disabled_component() { + let tc = core(); + let root = tc.root_node_handle(/* extra_key = */ None); + tc.component_has_host_value_only(root, SWA); +} + +#[test] +fn reset_clears_an_ongoing_evict_walk() { + let mut tc = core(); + tc.insert(&insert_params(&vec![1], &[10])); + tc.evict_device_start(FULL, /* request_cnt = */ 4); + assert!(tc.component_state(FULL).is_evict_device_ongoing); + tc.reset(); + // Reset drops the walk bookkeeping with the tree; a fresh walk starts clean. + assert!(!tc.component_state(FULL).is_evict_device_ongoing); + tc.insert(&insert_params(&vec![1], &[10])); + tc.evict_device_start(FULL, /* request_cnt = */ 4); + tc.evict_device_end(FULL); +} + +// Deterministic xorshift so the sequence test needs no rand dependency. +fn xorshift64(state: &mut u64) -> u64 { + let mut x = *state; + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + *state = x; + x +} + +// Fresh-KV insert params for the sequence test; donates a mamba slot when asked. +fn sequence_insert_params<'k>( + key: &'k Vec, + prev_prefix_len: usize, + kv_next: &mut i64, + mamba_next: &mut i64, + mamba: bool, +) -> InsertParams<'k, Vec> { + let depth = key.len(); + let kv: Vec = (0..depth).map(|i| *kv_next + i as i64).collect(); + *kv_next += depth as i64; + let mamba_value = mamba.then(|| { + *mamba_next += 1; + Tensor::from_slice(&[*mamba_next]) + }); + InsertParams { + key, + namespace: Default::default(), + value: Tensor::from_slice(&kv), + mamba_value, + prev_prefix_len, + swa_evicted_seqlen: 0, + chunked: false, + priority: 0, + track_adopted_ranges: false, + } +} + +// Randomized op sequence with a per-step sanity_check; `page` sizes the keys, +// `mamba` donates one state slot per insert. +fn run_random_op_sequence(mut tc: UnifiedTreeCore>, page: usize, mamba: bool) { + let cts: Vec = tc.components.iter().map(|c| c.component_type()).collect(); + let mut rng = 0x9E3779B97F4A7C15u64; + let mut kv_next = 1000i64; + let mut mamba_next = 1i64; + for step in 0..400 { + let depth = page * (1 + (xorshift64(&mut rng) % 3) as usize); + let key: Vec = (0..depth) + .map(|_| 1 + (xorshift64(&mut rng) % 5) as i64) + .collect(); + match xorshift64(&mut rng) % 4 { + 0 => { + tc.insert(&sequence_insert_params( + &key, + 0, + &mut kv_next, + &mut mamba_next, + mamba, + )); + } + 1 => { + // Consecutive matches of the same key agree (idempotency). + let first = tc.match_prefix(&match_params(&key)).device_indices.numel(); + let second = tc.match_prefix(&match_params(&key)).device_indices.numel(); + assert_eq!(first, second); + } + 2 => { + // Balanced lock round trip on whatever the key matches. + let anchor = tc.match_prefix(&match_params(&key)).best_match_node_id; + let lock = tc.inc_lock_ref(anchor); + let params = DecLockRefParams { + swa_uuid_for_lock: lock.swa_uuid_for_lock, + swa_uuid_for_host_lock: lock.swa_uuid_for_host_lock, + skip_lock_node_ids: lock.skip_lock_node_ids, + }; + tc.dec_lock_ref(anchor, Some(¶ms), /* skip_swa = */ false); + } + _ => { + // Insert-while-locked churn, the cache_finished_req shape. + let matched = tc.match_prefix(&match_params(&key)); + let anchor = matched.best_match_node_id; + let matched_len = matched.device_indices.numel() as usize; + let lock = tc.inc_lock_ref(anchor); + tc.insert(&sequence_insert_params( + &key, + matched_len, + &mut kv_next, + &mut mamba_next, + mamba, + )); + let params = DecLockRefParams { + swa_uuid_for_lock: lock.swa_uuid_for_lock, + swa_uuid_for_host_lock: lock.swa_uuid_for_host_lock, + skip_lock_node_ids: lock.skip_lock_node_ids, + }; + tc.dec_lock_ref(anchor, Some(¶ms), /* skip_swa = */ false); + } + } + if step % 8 == 7 { + for &ct in &cts { + let mut tracker: HashMap = + cts.iter().map(|&c| (c, 0)).collect(); + let mut device_frees = HashMap::new(); + let mut host_frees = HashMap::new(); + tc.evict_device_start(ct, /* request_cnt = */ 3); + loop { + let (next, step_result) = tc.evict_device_next_node(ct, &tracker); + accumulate_step( + step_result, + &mut tracker, + &mut device_frees, + &mut host_frees, + ); + let Some(leaf) = next else { break }; + let (_, evict_result) = + tc.evict_device_leaf(leaf, /* is_write_back = */ false); + accumulate_step( + evict_result, + &mut tracker, + &mut device_frees, + &mut host_frees, + ); + } + tc.evict_device_end(ct); + } + } + tc.sanity_check(&[], &[]); + } +} + +#[test] +fn random_op_sequence_holds_the_sanity_invariants() { + run_random_op_sequence( + swa_match_core(/* window = */ 4), + 1, + /* mamba = */ false, + ); +} + +#[test] +fn random_op_sequence_holds_on_a_mamba_core() { + let tc = UnifiedTreeCore::>::new( + CacheInitParams { + page_size: 1, + mamba_cache_chunk_size: Some(256), + ..CacheInitParams::default() + }, + vec![FULL, MAMBA], + ); + run_random_op_sequence(tc, 1, /* mamba = */ true); +} + +#[test] +fn random_op_sequence_holds_on_a_paged_swa_core() { + let tc = UnifiedTreeCore::>::new( + CacheInitParams { + page_size: 2, + swa_sliding_window_size: Some(8), + ..CacheInitParams::default() + }, + vec![FULL, SWA], + ); + run_random_op_sequence(tc, 2, /* mamba = */ false); +} + +// Drain every evictable FULL device leaf, as the orchestrator's evict loop does. +fn drain_full_device(tc: &mut UnifiedTreeCore>) { + let mut tracker = HashMap::from([(FULL, 0)]); + let (mut df, mut hf) = (HashMap::new(), HashMap::new()); + tc.evict_device_start(FULL, /* request_cnt = */ 1_000_000); + loop { + let (leaf, step) = tc.evict_device_next_node(FULL, &tracker); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + let Some(leaf) = leaf else { break }; + let (_, step) = tc.evict_device_leaf(leaf, /* is_write_back = */ false); + accumulate_step(step, &mut tracker, &mut df, &mut hf); + } + tc.evict_device_end(FULL); +} + +#[test] +fn an_emptied_namespace_leaves_nothing_behind() { + let mut tc = core(); + tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some("salted"), None), + ..insert_params(&vec![1, 2], &[10, 11]) + }); + let top = tc + .match_prefix(&MatchPrefixParams { + namespace: KeyNamespaceRef::new(Some("salted"), None), + ..match_params(&vec![1, 2]) + }) + .best_match_node_id; + drain_full_device(&mut tc); + // The namespace's nodes evict like any others; its edge map drops with them. + assert!(!tc.arena.namespace_exists(Some("salted"))); + assert!(tc.arena.try_resolve(top).is_none()); + assert_eq!(tc.arena.len(), 1); + tc.sanity_check(&[], &[]); + // A later insert respins the namespace from scratch. + tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some("salted"), None), + ..insert_params(&vec![1, 2], &[10, 11]) + }); + assert!(tc.arena.namespace_exists(Some("salted"))); + tc.sanity_check(&[], &[]); +} + +#[test] +fn namespaces_do_not_accumulate_across_salts() { + let mut tc = core(); + for salt in 0..64 { + let salt = format!("session-{salt}"); + tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some(&salt), None), + ..insert_params(&vec![1, 2], &[10, 11]) + }); + drain_full_device(&mut tc); + } + assert_eq!(tc.arena.len(), 1); + assert!(!tc.arena.namespace_exists(Some("session-0"))); + tc.sanity_check(&[], &[]); +} + +#[test] +fn a_zero_length_match_anchors_at_the_root() { + let mut tc = core(); + tc.insert(&InsertParams { + namespace: KeyNamespaceRef::new(Some("salted"), None), + ..insert_params(&vec![1, 2], &[10, 11]) + }); + let anchor = tc + .match_prefix(&MatchPrefixParams { + namespace: KeyNamespaceRef::new(Some("salted"), None), + ..match_params(&vec![9]) + }) + .best_match_node_id; + assert_eq!(anchor, tc.root_node_handle(Some("salted"))); + // The root handle stays valid across a full namespace eviction. + tc.inc_lock_ref(anchor); + drain_full_device(&mut tc); + tc.dec_lock_ref( + anchor, /* params = */ None, /* skip_swa = */ false, + ); + assert!(tc.arena.try_resolve(anchor).is_some()); + tc.sanity_check(&[], &[]); +} diff --git a/rust/mem-cache/src/unified_lru_list.rs b/rust/mem-cache/src/unified_lru_list.rs new file mode 100644 index 000000000..064937ad1 --- /dev/null +++ b/rust/mem-cache/src/unified_lru_list.rs @@ -0,0 +1,501 @@ +//! Self-contained LRU order over `NodeIdx_`s: MRU at the head side, LRU at the +//! tail side. Node semantics stay with callers through predicates; the reset +//! walks read parent links from the arena. + +use std::collections::HashSet; + +use crate::node::ChildKeyType; +use crate::node::Node; +use crate::node::NodeArena; +use crate::node::{NodeIdx_, ValueSlotIdx}; + +/// Index into the cell table; distinct from `NodeIdx_` so shifted and unshifted +/// ids cannot be mixed. +#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)] +struct CellId(usize); + +/// Head sentinel cell. +const HEAD: CellId = CellId(0); +/// Tail sentinel cell. +const TAIL: CellId = CellId(1); +/// Table offset: node ids map to cells after the sentinels. +const OFFSET: usize = 2; + +/// One doubly-linked cell; a node's cell lives at `NodeIdx_ + OFFSET`. +#[derive(Clone, Copy, Default)] +struct Cell { + prev: CellId, + next: CellId, + in_list: bool, +} + +/// LRU list over `NodeIdx_`s, with head/tail sentinel cells keeping the link +/// operations branchless. External APIs take `NodeIdx_`s; internal (`_`-suffixed) +/// interfaces work on `CellId`s. +pub struct UnifiedLRUList { + /// The (component × tier) value slot whose lock gates this list's walkers. + slot: ValueSlotIdx, + /// Cell table indexed by `NodeIdx_ + OFFSET`; cells 0/1 are the sentinels. + cells: Vec, + /// Number of member cells, excluding the sentinels. + len: usize, +} + +impl UnifiedLRUList { + pub fn new(slot: ValueSlotIdx) -> Self { + UnifiedLRUList { + slot, + cells: vec![ + // Sentinels link to each other and stay permanently flagged so + // the gated cell accessors admit them. + Cell { + prev: HEAD, + next: TAIL, + in_list: true, + }, + Cell { + prev: HEAD, + next: TAIL, + in_list: true, + }, + ], + len: 0, + } + } + + // ==== List operations ==== + + fn add_node_after_(&mut self, prev: CellId, cell: CellId) { + self.new_cell_(cell); + let next = self.cell_(prev).next; + self.connect_(cell, next); + self.connect_(prev, cell); + } + + fn add_node_(&mut self, cell: CellId) { + self.add_node_after_(HEAD, cell); + } + + fn remove_node_(&mut self, cell: CellId) { + let Cell { prev, next, .. } = *self.cell_(cell); + self.connect_(prev, next); + // Unflag the cell; the stale prev/next are never read while unlisted. + self.cell_mut_(cell).in_list = false; + self.len -= 1; + } + + // ==== NodeIdx_ <-> CellId plumbing ==== + + /// The node's cell slot; the only `NodeIdx_` -> `CellId` crossing. + fn cell_of_(node_id: NodeIdx_) -> CellId { + CellId(node_id.0 + OFFSET) + } + + /// The cell's node; the only `CellId` -> `NodeIdx_` crossing. + fn node_of_(cell: CellId) -> NodeIdx_ { + NodeIdx_(cell.0 - OFFSET) + } + + /// The cell, asserting it is linked (sentinels always are). + #[track_caller] + fn cell_(&self, id: CellId) -> &Cell { + let cell = &self.cells[id.0]; + assert!( + cell.in_list, + "node {} not in the LRU list", + Self::node_of_(id) + ); + cell + } + + #[track_caller] + fn cell_mut_(&mut self, id: CellId) -> &mut Cell { + let cell = &mut self.cells[id.0]; + assert!( + cell.in_list, + "node {} not in the LRU list", + Self::node_of_(id) + ); + cell + } + + /// Admit an unlisted cell: grow the table to cover it, then flag and count + /// it before any connections. + fn new_cell_(&mut self, cell: CellId) { + if cell.0 >= self.cells.len() { + self.cells.resize(cell.0 + 1, Cell::default()); + } + assert!( + !self.cells[cell.0].in_list, + "new_cell_: cell {cell:?} already in the LRU list" + ); + self.cells[cell.0].in_list = true; + self.len += 1; + } + + /// Whether the cell is linked into the list; safe on cells beyond the table. + fn in_list_(&self, cell: CellId) -> bool { + self.cells.get(cell.0).is_some_and(|cell| cell.in_list) + } + + /// Link `a -> b`. + fn connect_(&mut self, a: CellId, b: CellId) { + self.cell_mut_(a).next = b; + self.cell_mut_(b).prev = a; + } + + /// Insert a node as the most-recently-used; panics if already a member. + pub fn insert_mru(&mut self, node_id: NodeIdx_) { + self.add_node_(Self::cell_of_(node_id)); + } + + /// Remove a member node, resetting its cell; panics if not a member. + pub fn remove_node(&mut self, node_id: NodeIdx_) { + self.remove_node_(Self::cell_of_(node_id)); + } + + /// Move a member node back to the most-recently-used position. + pub fn reset_node_mru(&mut self, node_id: NodeIdx_) { + let cell = Self::cell_of_(node_id); + self.remove_node_(cell); + self.add_node_(cell); + } + + /// Re-rank the `should_include` nodes from `node_id` up to its root + /// (exclusive) as the MRU run, deepest first. + pub fn reset_node_and_parents_mru( + &mut self, + node_id: NodeIdx_, + arena: &NodeArena, + mut should_include: impl FnMut(&Node) -> bool, + ) { + let mut prev = HEAD; + let mut cur = node_id; + loop { + let node = arena.node(cur); + let Some(parent) = node.try_parent() else { + break; + }; + if should_include(node) { + let cell = Self::cell_of_(cur); + self.remove_node_(cell); + self.add_node_after_(prev, cell); + prev = cell; + } + cur = parent; + } + } + + /// Like `reset_node_and_parents_mru`, stopping once `window_size` atoms + /// are covered; excluded ancestors consume the window too. + pub fn reset_node_and_window_ancestors_mru( + &mut self, + node_id: NodeIdx_, + window_size: usize, + arena: &NodeArena, + mut should_include: impl FnMut(&Node) -> bool, + ) { + let mut prev = HEAD; + let mut accumulated = 0; + let mut cur = node_id; + while accumulated < window_size { + let node = arena.node(cur); + let Some(parent) = node.try_parent() else { + break; + }; + if should_include(node) { + let cell = Self::cell_of_(cur); + self.remove_node_(cell); + self.add_node_after_(prev, cell); + prev = cell; + } + accumulated += node.key.atom_len(); + cur = parent; + } + } + + /// Whether the node is a member (`None` is never a member). + pub fn in_list(&self, node_id: Option) -> bool { + node_id.is_some_and(|id| self.in_list_(Self::cell_of_(id))) + } + + /// The nearest predecessor of `cell` satisfying `pred`, walking toward the + /// head; `cell` itself is excluded. + fn get_prev_where_( + &self, + cell: CellId, + mut pred: impl FnMut(NodeIdx_) -> bool, + ) -> Option { + let mut cell = self.cell_(cell).prev; + while cell != HEAD { + let node = Self::node_of_(cell); + if pred(node) { + return Some(node); + } + cell = self.cell_(cell).prev; + } + None + } + + /// The nearest predecessor of a member satisfying `pred`; panics if + /// `node_id` is not a member. + pub fn get_prev_where( + &self, + node_id: NodeIdx_, + pred: impl FnMut(NodeIdx_) -> bool, + ) -> Option { + self.get_prev_where_(Self::cell_of_(node_id), pred) + } + + /// The least-recent member whose lock on the list's own slot is free. + pub fn get_lru_no_lock(&self, arena: &NodeArena) -> Option { + self.get_lru_where(|id| arena.node(id).lock_ref_(self.slot) == 0) + } + + /// The nearest more-recent member whose lock on the list's own slot is + /// free, from `node_id`. + pub fn get_prev_no_lock( + &self, + node_id: NodeIdx_, + arena: &NodeArena, + ) -> Option { + self.get_prev_where(node_id, |id| arena.node(id).lock_ref_(self.slot) == 0) + } + + /// The least-recently-used member satisfying `pred`. + pub fn get_lru_where(&self, pred: impl FnMut(NodeIdx_) -> bool) -> Option { + self.get_prev_where_(TAIL, pred) + } + + /// Number of member cells, excluding the sentinels. + pub fn len(&self) -> usize { + self.len + } + + /// Materialize the current members from most to least recent. + /// + /// Inspection callers need an owned snapshot across the Python boundary; + /// the linked-list iterator itself never escapes the Rust core. + pub(crate) fn snapshot_node_ids(&self) -> Vec { + let mut node_ids = Vec::with_capacity(self.len); + let mut cell = self.cell_(HEAD).next; + while cell != TAIL { + node_ids.push(Self::node_of_(cell)); + cell = self.cell_(cell).next; + } + node_ids + } + + // ==== Test-only conveniences ==== + + /// The members, MRU to LRU. + #[cfg(test)] + pub fn iter(&self) -> impl Iterator + '_ { + let mut cell = self.cell_(HEAD).next; + std::iter::from_fn(move || { + if cell == TAIL { + return None; + } + let node = Self::node_of_(cell); + cell = self.cell_(cell).next; + Some(node) + }) + } + + /// Panics if the links, membership flags, or member counter are inconsistent. + /// Reads cells raw: it inspects possibly-inconsistent state that the gated + /// accessors would reject. + #[cfg(test)] + pub fn validate(&self) { + let mut count = 0; + let mut prev = HEAD; + let mut cell = self.cells[HEAD.0].next; + while cell != TAIL { + assert!( + cell.0 >= OFFSET && cell.0 < self.cells.len(), + "validate: cell {cell:?} out of bounds" + ); + assert_eq!( + self.cells[cell.0].prev, prev, + "validate: broken prev link at cell {cell:?}" + ); + assert!( + self.cells[cell.0].in_list, + "validate: membership mismatch at cell {cell:?}" + ); + count += 1; + assert!(count <= self.len, "validate: cycle detected"); + prev = cell; + cell = self.cells[cell.0].next; + } + assert_eq!(self.cells[TAIL.0].prev, prev, "validate: broken tail link"); + assert_eq!(count, self.len, "validate: length mismatch"); + let flagged = self + .cells + .iter() + .skip(OFFSET) + .filter(|cell| cell.in_list) + .count(); + assert_eq!(flagged, self.len, "validate: membership mismatch"); + } + /// Test-only: desynchronize `len` to force integrity errors. + #[cfg(test)] + pub(crate) fn bump_len_for_test(&mut self) { + self.len += 1; + } + + /// Walk a LRU doubly-linked list, collect integrity errors. + pub(crate) fn check_linked_list_(&self, label: &str, errors: &mut Vec) { + let mut visited: HashSet = HashSet::new(); + let mut prev = HEAD; + let mut x = self.cells[HEAD.0].next; + while x != TAIL { + if x.0 < OFFSET { + errors.push(format!("{label} broken chain: link points at a sentinel")); + break; + } + let Some(cell) = self.cells.get(x.0) else { + errors.push(format!("{label} broken chain: cell {} out of bounds", x.0)); + break; + }; + if cell.prev != prev { + errors.push(format!("{label} broken prev at node {}", Self::node_of_(x))); + } + if !cell.in_list { + errors.push(format!( + "{label} node {} in list not flagged", + Self::node_of_(x) + )); + } + if !visited.insert(x.0) { + errors.push(format!("{label} cycle at node {}", Self::node_of_(x))); + break; + } + prev = x; + x = cell.next; + } + // The tail backlink closes the list onto the last visited member. + if x == TAIL && self.cells[TAIL.0].prev != prev { + errors.push(format!("{label} broken tail backlink")); + } + // Every flagged member cell must be reachable from the head. + for (idx, cell) in self.cells.iter().enumerate().skip(OFFSET) { + if cell.in_list && !visited.contains(&idx) { + errors.push(format!( + "{label} node {} flagged but unreachable", + idx - OFFSET + )); + } + } + if visited.len() != self.len { + errors.push(format!( + "{label} list={} != len={}", + visited.len(), + self.len + )); + } + } +} + +// Eviction priority keys. + +/// Eviction-priority key, ordered lexicographically; lower evicts first. +#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug)] +pub struct PriorityKey(pub i64, pub i64); + +/// Ranks nodes for eviction; lower priority evicts first. +pub trait EvictionStrategy { + /// The node's eviction priority. + fn get_priority(&self, node: &Node) -> PriorityKey; +} + +/// Least-recently-used. +pub struct LruStrategy; + +impl EvictionStrategy for LruStrategy { + fn get_priority(&self, node: &Node) -> PriorityKey { + PriorityKey(node.last_access_counter, 0) + } +} + +/// Least-frequently-used; LRU within a hit count. +pub struct LfuStrategy; + +impl EvictionStrategy for LfuStrategy { + fn get_priority(&self, node: &Node) -> PriorityKey { + PriorityKey(node.hit_count, node.last_access_counter) + } +} + +/// First-in-first-out over creation order. +pub struct FifoStrategy; + +impl EvictionStrategy for FifoStrategy { + fn get_priority(&self, node: &Node) -> PriorityKey { + PriorityKey(node.creation_counter, 0) + } +} + +/// Most-recently-used first. +pub struct MruStrategy; + +impl EvictionStrategy for MruStrategy { + fn get_priority(&self, node: &Node) -> PriorityKey { + PriorityKey(-node.last_access_counter, 0) + } +} + +/// First-in-last-out over creation order. +pub struct FiloStrategy; + +impl EvictionStrategy for FiloStrategy { + fn get_priority(&self, node: &Node) -> PriorityKey { + PriorityKey(-node.creation_counter, 0) + } +} + +/// Priority-aware: lower node priority evicts first, LRU within a priority. +pub struct PriorityStrategy; + +impl EvictionStrategy for PriorityStrategy { + fn get_priority(&self, node: &Node) -> PriorityKey { + PriorityKey(node.priority, node.last_access_counter) + } +} + +/// Segmented LRU: probationary nodes (hits below the threshold) evict before +/// protected ones, LRU within a segment. +pub struct SlruStrategy { + pub protected_threshold: i64, +} + +impl EvictionStrategy for SlruStrategy { + fn get_priority(&self, node: &Node) -> PriorityKey { + PriorityKey( + (node.hit_count >= self.protected_threshold) as i64, + node.last_access_counter, + ) + } +} + +/// The strategy for an eviction-policy name. +pub fn get_eviction_strategy(policy: &str) -> Box + Send> { + match policy.to_lowercase().as_str() { + "lru" => Box::new(LruStrategy), + "lfu" => Box::new(LfuStrategy), + "fifo" => Box::new(FifoStrategy), + "mru" => Box::new(MruStrategy), + "filo" => Box::new(FiloStrategy), + "priority" => Box::new(PriorityStrategy), + "slru" => Box::new(SlruStrategy { + protected_threshold: 2, + }), + other => panic!( + "Unknown eviction policy: {other}. Supported policies: \ + 'lru', 'lfu', 'fifo', 'mru', 'filo', 'priority', 'slru'." + ), + } +} +#[cfg(test)] +#[path = "tests/unified_lru_list.rs"] +mod tests; diff --git a/rust/mem-cache/src/unified_tree_core.rs b/rust/mem-cache/src/unified_tree_core.rs new file mode 100644 index 000000000..c7e8d65e6 --- /dev/null +++ b/rust/mem-cache/src/unified_tree_core.rs @@ -0,0 +1,4811 @@ +//! The radix prefix tree of cached KV. +#![allow(unused_variables)] + +use std::cmp::Reverse; +use std::collections::{BinaryHeap, HashMap, HashSet}; +use std::sync::Arc; + +use sha2::{Digest, Sha256}; +use tch::{Device, Kind, Tensor}; + +use crate::components::{self, FullComponent, MambaComponent, SwaComponent, TreeComponent}; +use crate::components::{ + BASE_COMPONENT_TYPE, ComponentType, FULL, MAMBA, NUM_COMPONENT_TYPES, SWA, +}; +use crate::node::EvictableNodeSet; +use crate::node::Node; +use crate::node::NodeArena; +use crate::node::{ChildKeyType, HashDigest, KeyNamespace, KeyNamespaceRef}; +use crate::node::{NUM_VALUE_SLOTS, NodeId, NodeIdx_, TreeCoreRuntimeError, ValueSlotIdx}; +use crate::unified_lru_list::UnifiedLRUList; +use crate::unified_lru_list::{EvictionStrategy, PriorityKey, get_eviction_strategy}; + +// A 42-bit mask keeps digest multiplication by 1_000_003 within i64. +const COEXIST_RECLAIM_DIGEST_MULTIPLIER: i64 = 1_000_003; +const COEXIST_RECLAIM_DIGEST_MASK: i64 = (1 << 42) - 1; + +fn next_coexist_reclaim_digest(current: i64, node_id: NodeId, component_idx: usize) -> i64 { + let event = (node_id as i64 + 1) * NUM_COMPONENT_TYPES as i64 + component_idx as i64; + (current * COEXIST_RECLAIM_DIGEST_MULTIPLIER + event) & COEXIST_RECLAIM_DIGEST_MASK +} + +// ---- interface types ---- + +/// Result of `inc_lock_ref`, handed back to the matching `dec_lock_ref`. +#[derive(Default)] +pub struct IncLockRefResult { + /// Tokens newly protected (moved out of evictable) by this lock. + pub delta: Option, + /// SWA lock-window uuid minted/reused by the device lock walk. + pub swa_uuid_for_lock: Option, + /// SWA lock-window uuid minted/reused by the host lock walk. + pub swa_uuid_for_host_lock: Option, + /// Per-component nodes that were tombstones at acquire time; replayed at + /// release so the unlock skips them. + pub skip_lock_node_ids: HashMap>, +} + +/// Params for `dec_lock_ref`. +#[derive(Default)] +pub struct DecLockRefParams { + /// SWA lock-window uuid the device unlock stops at, from the matching acquire. + pub swa_uuid_for_lock: Option, + /// SWA lock-window uuid the host unlock stops at, from the matching acquire. + pub swa_uuid_for_host_lock: Option, + /// Per-component nodes the unlock walk skips (from the matching acquire). + pub skip_lock_node_ids: HashMap>, +} + +/// Result of `dec_lock_ref`. +#[derive(Default)] +pub struct DecLockRefResult {} + +/// Result of a prefix match. +pub struct MatchResult { + /// Device KV indices matched by the common prefix. + pub device_indices: Tensor, + /// Last matched node still resident on device. + pub last_device_node_id: NodeId, + /// Last matched node on host; equals `last_device_node_id` without HiCache. + pub last_host_node_id: NodeId, + /// Deepest node accepted by all component validators; anchors host->device load-back. + pub best_match_node_id: NodeId, + /// Full-KV tokens that hit on host and must be loaded back to device. + pub host_hit_length: usize, + /// SWA tokens that hit on host (within the sliding window) and will be + /// loaded back into the SWA device pool. + pub swa_host_hit_length: usize, + /// Mamba slots that hit on host and will be loaded back; 0 or 1. + pub mamba_host_hit_length: usize, + /// The longest chunk-aligned position that could have hit if a mamba state existed. + pub mamba_branching_seqlen: Option, + /// Longest Full-KV prefix available on either device or host, independent + /// of other components. + pub full_kv_hit_length: usize, + /// Actions for the controller to apply. + pub cache_actions: Vec, +} + +/// Params for a prefix match; the key is borrowed from the caller. +pub struct MatchPrefixParams<'k, K: ChildKeyType> { + /// The query key (already page-typed; bigram conversion happens at the boundary). + pub key: &'k K, + /// Namespace of the query; picks the matching subtree root. + pub namespace: KeyNamespaceRef<'k>, +} + +/// Params for an insert; the key is borrowed from the caller. +pub struct InsertParams<'k, K: ChildKeyType> { + /// The insert key (already page-typed; bigram conversion happens at the boundary). + pub key: &'k K, + /// Namespace of the insert; picks the matching subtree root. + pub namespace: KeyNamespaceRef<'k>, + /// Device KV indices covering the key, one row per atom. + pub value: Tensor, + /// Tokens of this request already cached before the insert (the duplicate + /// window starts past them). + pub prev_prefix_len: usize, + /// The request's SWA-evicted prefix boundary; SWA data below it stays tombstoned. + pub swa_evicted_seqlen: usize, + /// The donated mamba slot for the insert target leaf; None on non-mamba trees. + pub mamba_value: Option, + /// Whether this is a chunked-prefill insert (no hit-count bump). + pub chunked: bool, + /// Eviction priority floor applied along the walked path. + pub priority: i64, + /// Whether the result should report which incoming ranges the tree retained. + pub track_adopted_ranges: bool, +} + +/// Result of an insert. +#[derive(Default)] +pub struct InsertResult { + /// Tokens of the insert key that overlapped existing nodes. + pub prefix_len: usize, + /// The inserted key's full (page-aligned) length. + pub total_len: usize, + /// The device-resident node at the end of the inserted path. + pub last_device_node_id: Option, + /// Whether the cache holds Mamba state covering the inserted sequence; + /// vacuously true for an empty insert. + pub mamba_exist: bool, + /// The deepest host-backed node an insert_host attached or matched. + pub inserted_host_node: Option, + /// Whether write-through rejected a host suffix below an unbacked parent. + pub host_insert_dropped: bool, + /// Incoming ranges retained by each component, in key-relative atom offsets. + pub adopted_ranges: Option>>, + /// Actions for the controller to apply. + pub cache_actions: Vec, +} + +impl InsertResult { + pub fn record_adopted_range( + &mut self, + component_type: ComponentType, + start: usize, + end: usize, + ) { + let Some(adopted_ranges) = self.adopted_ranges.as_mut() else { + return; + }; + if start >= end { + return; + } + let ranges = adopted_ranges.entry(component_type).or_default(); + if let Some((previous_start, previous_end)) = ranges.last_mut() + && start <= *previous_end + { + *previous_start = (*previous_start).min(start); + *previous_end = (*previous_end).max(end); + } else { + ranges.push((start, end)); + } + } +} + +/// One step of a resumable insert: the Controller executes `actions`, then +/// resumes while `result` is None; `result` is set on the final step. +pub struct InsertStepResult { + pub actions: Vec, + pub result: Option, +} + +// WALK (one node per step) -> COMMIT (leaf + commit hooks) -> TAIL (refresh + backup). +pub enum InsertPhase { + Walk, + Commit, + Tail, +} + +/// In-flight resumable-insert state persisted across step barriers; owns its +/// key/value/params snapshot so the walk survives across boundary calls. +pub struct InsertWalkState { + phase: InsertPhase, + node_id: NodeIdx_, + /// The full page-aligned insert key; `total_prefix_length` is the walk cursor. + key: K, + aligned_key_len: usize, + value: Tensor, + namespace: KeyNamespace, + prev_prefix_len: usize, + swa_evicted_seqlen: usize, + mamba_value: Option, + chunked: bool, + priority: i64, + track_adopted_ranges: bool, + total_prefix_length: usize, + is_new_leaf: bool, + target_node_id: Option, + result: Option, + /// Emitted actions awaiting the next barrier flush (or the final step). + pending_actions: Vec, +} + +/// Result of a KV-canary walk: parallel per-slot rows over the tree's FULL device slots. +pub struct KvCanaryWalkResult { + /// Device slot index of each emitted token. + pub slot_indices: Vec, + /// Token depth from the root for each emitted slot. + pub positions: Vec, + /// The preceding device slot on the path (-1 at a chain start). + pub prev_slot_indices: Vec, +} + +/// A queued cache IO action. +pub enum CacheAction { + /// Duplicate device KV slices the cache frees after the insert. + FreeDeviceKV(Vec), + /// Free the full side only, for a tombstoned node whose SWA peers are gone; + /// FreeDeviceKV would release the SWA side twice. + FreeDeviceKVFullOnly(Vec), + /// A device->host backup work item (the write-through threshold fired). + BackupKV(BackupKV), + /// Replace the pending write-through node on a node split: + /// + /// parent -> node => parent -> new_node -> new_child + /// + /// old_node_id (the pre-split node) is replaced by new_node_id + new_child_node_id. + ReplaceWriteThroughOnNodeSplit { + ack_id: usize, + old_node_id: NodeId, + new_node_id: NodeId, + new_child_node_id: NodeId, + }, + /// Per-path Mamba state-cap eviction from the tail's root path; applied at + /// the insert's commit barrier, after the walk-time backups whose + /// write-through locks shield the backed-up chain. + MambaEvictExcessPathStates { tail_node_id: NodeId }, + /// Free only the given component's device KV slots. + FreeComponentDeviceSlot { + component_type: ComponentType, + indices: Vec, + }, + /// Free the given component's host KV pages. + FreeComponentHostSlot { + component_type: ComponentType, + host_indices: Vec, + }, + /// Rebuild the SWA allocator's full->swa index mapping for loaded chunks. + RebuildFullToSwaMapping { + full_indices: Vec, + swa_indices: Vec, + }, + /// Recover an SWA tombstone whose full is locked: keep the locked full, remap + /// it onto the incoming full's SWA translation, and free only the incoming full. + RecoverSwaWithLockedFull { + node_id: NodeId, + kept_full: Tensor, + incoming_full: Tensor, + }, + /// Rebuild a node's SWA value by translating its source full value, then store it. + SwaRebuild { + node_id: NodeId, + source_value: Tensor, + }, +} + +/// A HiCache pool transfer descriptor. +#[derive(Default)] +pub struct PoolTransfer { + /// The pool this transfer targets. + pub name: PoolName, + /// Host-side indices for the device<->host path. + pub host_indices: Option, + /// Device-side indices, filled in once the transfer lands on device. + pub device_indices: Option, + /// Per-page storage keys for the host<->storage path. + pub keys: Option>, + /// How storage prefix-matches this pool's pages. + pub hit_policy: PoolHitPolicy, + /// The nodes a load-back restores, ancestors first (external handles). + pub nodes_to_load: Option>, +} + +/// Hit policy for storage's per-pool prefix matching. +#[derive(Copy, Clone, PartialEq, Eq, Debug, Default)] +pub enum PoolHitPolicy { + /// Every page in the hit range must exist. + #[default] + AllPages, + /// Only the last N pages must exist (window/state pools). + TrailingPages, +} + +impl PoolHitPolicy { + /// The python PoolHitPolicy enum value. + pub fn as_str(self) -> &'static str { + match self { + PoolHitPolicy::AllPages => "all_pages", + PoolHitPolicy::TrailingPages => "trailing_pages", + } + } +} + +/// Well-known pool names used as PoolTransfer identifiers. +#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug, Default)] +pub enum PoolName { + #[default] + Kv, + Mamba, + Swa, + Indexer, + DeepseekV4C4, + DeepseekV4C4Indexer, + DeepseekV4C128, + DeepseekV4C4State, + DeepseekV4C4IndexerState, + DeepseekV4C128State, + Draft, + DraftIndexer, + DraftSwa, +} + +/// Result of a HiCache pool transfer. +#[derive(Default)] +pub struct PoolTransferResult { + /// Pages of the KV pool the storage transfer completed. + pub kv_hit_pages: usize, + /// Completed pages per auxiliary pool. + pub extra_pool_hit_pages: HashMap, +} + +/// A device->host backup work item for the cache to execute. +#[derive(Default)] +pub struct BackupKV { + /// Backup these nodes device->host in order, stopping at the first failure; the + /// caller orders them parent-before-child for write-through and child-first for + /// write-back. External handles: the list crosses to the orchestrator. + pub node_ids: Vec, +} + +/// A device->storage backup spec. +#[derive(Default)] +pub struct StorageBackupSpec { + /// The node's FULL host value (the storage write's source indices). + pub host_value: Tensor, + /// Raw token ids spanned by the node's key. + pub token_ids: Vec, + /// The node's per-page hash chain (the storage keys). + pub hash_value: Option>, + /// Ancestor-chain hashes, root-to-parent, when requested. + pub prefix_keys: Option>, + /// Auxiliary per-component transfers riding the same storage write. + pub comp_xfers: HashMap>, +} + +/// Immutable metadata for a queued buffer-only backup. +pub struct BufferBackupSnapshot { + pub node_id: NodeId, + pub parent_node_id: NodeId, + pub parent_is_root: bool, + pub parent_last_hash: Option, + pub token_ids: Vec, + pub extra_key: Option, + pub cache_salt: Option, + pub is_bigram: bool, + pub hash_values: Vec, + pub prefix_keys: Option>, +} + +pub struct BufferBackupState { + pub parent_node_id: NodeId, + pub parent_is_root: bool, + pub parent_last_hash: Option, +} + +/// Which storage layer(s) an eviction targets. +#[derive(Copy, Clone, PartialEq, Eq, Debug)] +pub enum EvictLayer { + Device, + Host, + All, +} + +impl EvictLayer { + /// Whether this target includes `layer` (the Python IntFlag `in` membership). + pub fn contains(self, layer: EvictLayer) -> bool { + self == EvictLayer::All || self == layer + } +} + +/// The request fields load-back planning reads. +#[derive(Default)] +pub struct Req { + /// Mamba pool slot backing the request, when one is assigned. + pub mamba_pool_idx: Option, +} + +/// When the LRU is refreshed during a tree walk. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub enum LRURefreshPhase { + Walkdown, + MatchEnd, + InsertEnd, +} + +/// Direction of a HiCache transfer. +#[derive(Copy, Clone, PartialEq, Eq, Debug)] +pub enum CacheTransferPhase { + BackupHost, + LoadBack, + BackupStorage, + Prefetch, +} + +/// Per-component tree-wide bookkeeping (device-tier accounting and walk state). +#[derive(Default)] +pub struct ComponentState { + /// Evictable device token count. + pub(crate) evictable_size: usize, + /// Locked (protected) device token count. + pub(crate) protected_size: usize, + /// Whether a device-eviction walk is between start and end. + pub(crate) is_evict_device_ongoing: bool, + /// The walk's resume point, captured at return time since the returned + /// leaf may be freed: the leaf's parent for Full, the LRU predecessor for + /// SWA and Mamba. + pub(crate) evict_device_cursor: Option, + /// Token budget for the current eviction walk. + pub(crate) evict_device_request_cnt: usize, +} + +/// Tree-construction parameters: the tree-consumed slice of the cache's init params. +pub struct CacheInitParams { + /// Eviction-policy name resolved into the tree's strategy. + pub eviction_policy: String, + /// Atoms per radix page; children are keyed by their key's first page. + pub page_size: usize, + /// Whether the cache runs the write-back (vs write-through) policy. + pub is_write_back: bool, + /// Whether the host tier (HiCache) is wired. + pub enable_hicache: bool, + /// Hit count at which a node earns a host write-through backup. + pub write_through_threshold: i64, + /// Device the KV indices live on. + pub device: Device, + /// SWA sliding window size in tokens; None when SWA is disabled. + pub swa_sliding_window_size: Option, + /// Whether the cache wired a host SWA pool (HiCache). + pub has_swa_host_pool: bool, + /// Whether tree mutations emit BlockStored/BlockRemoved events. + pub enable_kv_cache_events: bool, + /// Chunk alignment for the mamba branching seqlen; None when Mamba is disabled. + pub mamba_cache_chunk_size: Option, + /// Per-root-path cap on cached Mamba states; None means unlimited. + pub mamba_max_states_per_path: Option, +} + +impl Default for CacheInitParams { + fn default() -> Self { + CacheInitParams { + eviction_policy: "lru".to_string(), + page_size: 1, + is_write_back: false, + enable_hicache: false, + write_through_threshold: 256, + device: Device::Cpu, + swa_sliding_window_size: None, + has_swa_host_pool: false, + enable_kv_cache_events: false, + mamba_cache_chunk_size: None, + mamba_max_states_per_path: None, + } + } +} + +/// Radix tree of cached token prefixes; each node carries its KV per component. +/// A single eviction step's outputs: this step's per-component evicted +/// counts (deltas, for the Controller to accumulate) and freed tensors. +#[derive(Default, Debug)] +pub struct EvictionStepResult { + pub tracker: HashMap, + pub device_frees: HashMap>, + pub host_frees: HashMap>, +} + +/// The radix tree mechanism: owns the tree structure, per-node values, the +/// per-component LRUs, the size/leaf bookkeeping, and the component drivers, +/// plus `reset()`. +pub struct UnifiedTreeCore { + pub(crate) arena: NodeArena, + /// Ordered component registry; each driver reports its own type. + components: Vec + Send + Sync>>, + /// Prebuilt per-type driver lookup, indexed by `ComponentType::idx`. + components_by_type: [Option + Send + Sync>>; NUM_COMPONENT_TYPES], + /// Per-component bookkeeping, indexed by `ComponentType::idx`. + pub(crate) component_states: [ComponentState; NUM_COMPONENT_TYPES], + /// Nodes currently eligible for device eviction (D-leaves). + pub(crate) evictable_device_leaves: EvictableNodeSet, + /// Nodes currently eligible for host eviction (H-leaves). + pub(crate) evictable_host_leaves: EvictableNodeSet, + /// Full has no device LRU, so track nodes whose device and host values coexist. + pub(crate) full_coexisting_host_nodes: EvictableNodeSet, + pub(crate) write_back_coexist_reclaim_digest: i64, + /// Per-slot LRU lists, indexed by `ValueSlotIdx::idx`. + pub(crate) lru_lists: [UnifiedLRUList; NUM_VALUE_SLOTS], + /// Device-eviction candidates; the lowest priority is popped first. + pub(crate) full_evict_device_heap: BinaryHeap>, + /// Eviction-priority strategy; lower priority evicts first. + pub(crate) eviction_strategy: Box + Send>, + /// Atoms per radix page; children are keyed by their key's first page. + pub(crate) page_size: usize, + /// Whether the cache runs the write-back (vs write-through) policy. + pub(crate) is_write_back: bool, + /// Whether the host tier (HiCache) is wired. + pub(crate) enable_hicache: bool, + /// Whether the storage tier (L3) is wired; gates page-hash computation. + pub(crate) enable_storage: bool, + /// Whether the cache wired a host SWA pool (HiCache). + pub(crate) has_swa_host_pool: bool, + /// Whether tree mutations emit BlockStored/BlockRemoved events. + pub(crate) enable_kv_cache_events: bool, + /// Queued placement events, drained by take_events. + pub(crate) kv_event_queue: Vec>, + /// Namespace-aware event hashes, populated only for salted nodes whose + /// placement events are requested. Storage hashes remain on the nodes. + pub(crate) salted_event_hashes: HashMap>, + /// Hit count at which a node earns a host write-through backup. + pub(crate) write_through_threshold: i64, + + /// Monotonic source for SWA lock-window uuids. + pub(crate) swa_uuid_counter: i64, + /// Device the KV indices live on. + pub(crate) device: Device, + /// Shared empty device-index tensor (an empty match's indices). + pub(crate) empty_device_indices: Tensor, + /// The single in-flight resumable insert, if suspended at a barrier. + ongoing_insert_walk_state: Option>, +} + +impl UnifiedTreeCore { + /// Build a tree core for the given component types with a fresh arena. + /// Fresh per-slot LRU lists. + pub(crate) fn new_lru_lists() -> [UnifiedLRUList; NUM_VALUE_SLOTS] { + std::array::from_fn(|i| UnifiedLRUList::new(ValueSlotIdx::from_idx(i))) + } + + /// The component's device-tier LRU list. + pub(crate) fn device_lru_list(&self, component_type: ComponentType) -> &UnifiedLRUList { + self.lru_list_(ValueSlotIdx::device(component_type)) + } + + /// The component's device-tier LRU list, mutable. + pub(crate) fn device_lru_list_mut( + &mut self, + component_type: ComponentType, + ) -> &mut UnifiedLRUList { + self.lru_list_mut_(ValueSlotIdx::device(component_type)) + } + + /// The component's host-tier LRU list. + pub(crate) fn host_lru_list(&self, component_type: ComponentType) -> &UnifiedLRUList { + self.lru_list_(ValueSlotIdx::host(component_type)) + } + + /// The component's host-tier LRU list, mutable. + pub(crate) fn host_lru_list_mut( + &mut self, + component_type: ComponentType, + ) -> &mut UnifiedLRUList { + self.lru_list_mut_(ValueSlotIdx::host(component_type)) + } + + /// The LRU list gated by the slot's lock. + pub(crate) fn lru_list_(&self, slot: ValueSlotIdx) -> &UnifiedLRUList { + &self.lru_lists[slot.idx()] + } + + /// The LRU list gated by the slot's lock, mutable. + pub(crate) fn lru_list_mut_(&mut self, slot: ValueSlotIdx) -> &mut UnifiedLRUList { + &mut self.lru_lists[slot.idx()] + } + + /// The component's device LRU list, mutable, paired with the arena the + /// reset walks read. + pub(crate) fn device_lru_list_mut_with_arena( + &mut self, + component_type: ComponentType, + ) -> (&mut UnifiedLRUList, &NodeArena) { + ( + &mut self.lru_lists[ValueSlotIdx::device(component_type).idx()], + &self.arena, + ) + } + + /// The component's tree-wide bookkeeping state. + pub(crate) fn component_state(&self, component_type: ComponentType) -> &ComponentState { + &self.component_states[component_type.idx()] + } + + /// The component's mutable tree-wide bookkeeping state. + pub(crate) fn component_state_mut( + &mut self, + component_type: ComponentType, + ) -> &mut ComponentState { + &mut self.component_states[component_type.idx()] + } + + /// The component's evictable device-token count. + pub(crate) fn evictable_size_(&self, component_type: ComponentType) -> usize { + self.component_state(component_type).evictable_size + } + + /// The component's protected (locked) device-token count. + pub(crate) fn protected_size_(&self, component_type: ComponentType) -> usize { + self.component_state(component_type).protected_size + } + + /// Begin the component's device-eviction bookkeeping for up to + /// `request_cnt` tokens; panics if a walk is already in progress. + pub(crate) fn set_evict_device_start( + &mut self, + component_type: ComponentType, + request_cnt: usize, + ) { + let state = self.component_state_mut(component_type); + assert!( + !state.is_evict_device_ongoing, + "{component_type:?} device eviction already in progress" + ); + state.is_evict_device_ongoing = true; + state.evict_device_request_cnt = request_cnt; + state.evict_device_cursor = None; + } + + /// Finish the component's device-eviction bookkeeping; panics if no walk + /// is in progress. + pub(crate) fn set_evict_device_end(&mut self, component_type: ComponentType) { + let state = self.component_state_mut(component_type); + assert!( + state.is_evict_device_ongoing, + "{component_type:?} device eviction not started" + ); + state.is_evict_device_ongoing = false; + state.evict_device_cursor = None; + } + + /// Add newly evictable device tokens to the component's evictable size. + pub(crate) fn inc_evictable_size(&mut self, component_type: ComponentType, tokens: usize) { + self.component_state_mut(component_type).evictable_size += tokens; + } + + /// Subtract freed device tokens from the component's evictable size. + pub(crate) fn dec_evictable_size(&mut self, component_type: ComponentType, tokens: usize) { + let state = self.component_state_mut(component_type); + state.evictable_size = state.evictable_size.checked_sub(tokens).unwrap_or_else(|| { + panic!("dec_evictable_size: {component_type:?} evictable size underflow") + }); + } + + /// Add newly locked device tokens to the component's protected size. + pub(crate) fn inc_protected_size(&mut self, component_type: ComponentType, tokens: usize) { + self.component_state_mut(component_type).protected_size += tokens; + } + + /// Subtract unlocked device tokens from the component's protected size. + pub(crate) fn dec_protected_size(&mut self, component_type: ComponentType, tokens: usize) { + let state = self.component_state_mut(component_type); + state.protected_size = state.protected_size.checked_sub(tokens).unwrap_or_else(|| { + panic!("dec_protected_size: {component_type:?} protected size underflow") + }); + } + + pub fn new(params: CacheInitParams, component_types: Vec) -> Self { + assert!( + !component_types.is_empty(), + "at least one component type is required" + ); + assert!( + component_types.contains(&BASE_COMPONENT_TYPE), + "the base (Full) component is required" + ); + assert!(params.page_size >= 1, "page_size must be at least 1"); + let arena = NodeArena::new(component_types.clone(), params.page_size); + let mut tree_core = UnifiedTreeCore { + arena, + components: Vec::new(), + components_by_type: Default::default(), + component_states: Default::default(), + evictable_device_leaves: EvictableNodeSet::new(), + evictable_host_leaves: EvictableNodeSet::new(), + full_coexisting_host_nodes: EvictableNodeSet::new(), + write_back_coexist_reclaim_digest: 0, + // Disabled components keep harmless empty lists, like component_states. + lru_lists: Self::new_lru_lists(), + full_evict_device_heap: BinaryHeap::new(), + eviction_strategy: get_eviction_strategy(¶ms.eviction_policy), + page_size: params.page_size, + is_write_back: params.is_write_back, + enable_hicache: params.enable_hicache, + enable_storage: false, + has_swa_host_pool: params.has_swa_host_pool, + enable_kv_cache_events: params.enable_kv_cache_events, + kv_event_queue: Vec::new(), + salted_event_hashes: HashMap::new(), + write_through_threshold: params.write_through_threshold, + swa_uuid_counter: 1, + device: params.device, + empty_device_indices: Tensor::empty([0], (Kind::Int64, params.device)), + ongoing_insert_walk_state: None, + }; + for ct in &component_types { + let component: Arc + Send + Sync> = match ct { + ComponentType::Full => Arc::new(FullComponent), + ComponentType::Swa => Arc::new(SwaComponent::new(¶ms)), + ComponentType::Mamba => Arc::new(MambaComponent::new(¶ms)), + }; + tree_core.register_component_(component); + } + tree_core + } + + /// Rebuild the root, LRUs, sizes, evictable-leaf sets, and the empty + /// match result. + pub fn reset(&mut self) { + self.arena.reset(); + self.component_states = Default::default(); + self.evictable_device_leaves = EvictableNodeSet::new(); + self.evictable_host_leaves = EvictableNodeSet::new(); + self.full_coexisting_host_nodes = EvictableNodeSet::new(); + self.write_back_coexist_reclaim_digest = 0; + self.lru_lists = Self::new_lru_lists(); + self.full_evict_device_heap.clear(); + self.salted_event_hashes.clear(); + self.ongoing_insert_walk_state = None; + } + + /// Create a keyed, parented node not yet in its parent's child map; + /// `creation_counter` None keeps the fresh allocation stamp. + pub fn new_node_( + &mut self, + key: K, + parent_id: NodeIdx_, + priority: i64, + hit_count: i64, + creation_counter: Option, + extra_key: Option<&str>, + ) -> NodeIdx_ { + self.new_node_in_namespace_( + key, + parent_id, + priority, + hit_count, + creation_counter, + KeyNamespaceRef::new(extra_key, /* cache_salt = */ None), + ) + } + + pub fn new_node_in_namespace_( + &mut self, + key: K, + parent_id: NodeIdx_, + priority: i64, + hit_count: i64, + creation_counter: Option, + namespace: KeyNamespaceRef<'_>, + ) -> NodeIdx_ { + let new_node_id = self.arena.alloc_detached(priority); + // Root children adopt the op namespace; deeper nodes inherit the parent's. + let ns = if self.arena.node(parent_id).is_root() { + namespace.to_owned() + } else { + self.arena.node(parent_id).namespace.clone() + }; + let new_node = self.arena.node_mut(new_node_id); + new_node.key = key; + new_node.parent = Some(parent_id); + new_node.namespace = ns; + new_node.hit_count = hit_count; + if let Some(creation_counter) = creation_counter { + new_node.creation_counter = creation_counter; + } + new_node_id + } + + /// Mint the next SWA lock-window uuid. + pub(crate) fn next_swa_uuid_(&mut self) -> i64 { + self.swa_uuid_counter += 1; + self.swa_uuid_counter + } + + /// Bump the reference count on a node's component locks. + pub fn inc_lock_ref(&mut self, node_id: NodeId) -> IncLockRefResult { + self.inc_lock_ref_with_skip(node_id, &[]) + } + + /// Bump component locks, leaving explicitly skipped target components evictable. + pub fn inc_lock_ref_with_skip( + &mut self, + node_id: NodeId, + skip_lock_components: &[ComponentType], + ) -> IncLockRefResult { + let node_id = self.arena.resolve(node_id); + let node = self.arena.node(node_id); + let node_handle = node.id; + let is_root = node.is_root(); + let mut result = IncLockRefResult::default(); + for i in 0..self.components.len() { + let component_type = self.components[i].component_type(); + if skip_lock_components.contains(&component_type) { + if !is_root { + result + .skip_lock_node_ids + .entry(component_type) + .or_default() + .insert(node_handle); + } + continue; + } + let component = Arc::clone(&self.components[i]); + result = component + .acquire_component_lock(self, node_id, result, /* lock_host = */ false); + } + self.update_evictable_leaf_sets_(node_id); + result + } + + /// Decrease the reference count on a node's component locks. + pub fn dec_lock_ref( + &mut self, + node_id: NodeId, + params: Option<&DecLockRefParams>, + skip_swa: bool, + ) -> DecLockRefResult { + let node_id = self.arena.resolve(node_id); + for i in 0..self.components.len() { + if skip_swa && self.components[i].component_type() == SWA { + continue; + } + let component = Arc::clone(&self.components[i]); + component.release_component_lock(self, node_id, params, /* lock_host = */ false); + } + self.update_evictable_leaf_sets_(node_id); + // TODO: delta is not aggregated from components; no caller uses it yet. + DecLockRefResult::default() + } + + /// Early-release the SWA portion of a request's tree lock, plus any + /// strictly-lower-priority locks (e.g. Mamba) co-located on the node. + pub fn dec_swa_lock_only( + &mut self, + node_id: NodeId, + swa_uuid_for_lock: Option, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + ) { + self.dec_swa_lock_only_with_skip( + node_id, + swa_uuid_for_lock, + /* skip_lock_node_ids = */ None, + device_frees, + host_frees, + ); + } + + /// Skip-aware variant used when an acquire deliberately omitted a component. + pub fn dec_swa_lock_only_with_skip( + &mut self, + node_id: NodeId, + swa_uuid_for_lock: Option, + skip_lock_node_ids: Option<&HashMap>>, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + ) { + let node_id = self.arena.resolve(node_id); + let Some(swa) = self.try_component_by_type_(SWA) else { + return; + }; + swa.release_window_lock(self, node_id, swa_uuid_for_lock, device_frees, host_frees); + + // Drop strictly-lower-priority locks (e.g. Mamba) co-located on the node. + let swa_priority = swa.eviction_priority(/* is_leaf = */ false); + let dec_params = DecLockRefParams { + swa_uuid_for_lock, + skip_lock_node_ids: skip_lock_node_ids.cloned().unwrap_or_default(), + ..Default::default() + }; + for i in 0..self.components.len() { + let component = Arc::clone(&self.components[i]); + if component.eviction_priority(/* is_leaf = */ false) < swa_priority { + component.release_component_lock( + self, + node_id, + Some(&dec_params), + /* lock_host = */ false, + ); + } + } + } + + /// Evict shallow Mamba device checkpoints beyond the per-path cap on the + /// tail's root path; the mamba component drives the walk. + pub fn evict_excess_path_states(&mut self, tail_node_id: NodeId) -> EvictionStepResult { + let tail_node_id = self.arena.resolve(tail_node_id); + let mut result = EvictionStepResult::default(); + let component = self.component_by_type_(MAMBA); + component.evict_excess_path_states( + self, + tail_node_id, + &mut result.device_frees, + &mut result.host_frees, + ); + result + } + + /// Bump the reference count on a node's host-side component locks. + pub fn inc_host_lock_ref(&mut self, node_id: NodeId) -> IncLockRefResult { + let node_id = self.arena.resolve(node_id); + let mut result = IncLockRefResult::default(); + for i in 0..self.components.len() { + let component = Arc::clone(&self.components[i]); + result = component + .acquire_component_lock(self, node_id, result, /* lock_host = */ true); + } + self.update_evictable_leaf_sets_(node_id); + result + } + + /// Decrease the reference count on a node's host-side component locks. + pub fn dec_host_lock_ref( + &mut self, + node_id: NodeId, + params: Option<&DecLockRefParams>, + ) -> DecLockRefResult { + let node_id = self.arena.resolve(node_id); + for i in 0..self.components.len() { + let component = Arc::clone(&self.components[i]); + component.release_component_lock(self, node_id, params, /* lock_host = */ true); + } + self.update_evictable_leaf_sets_(node_id); + DecLockRefResult::default() + } + + /// Match a key against the tree; returns device indices + boundary NodeIds. + pub fn match_prefix(&mut self, params: &MatchPrefixParams<'_, K>) -> MatchResult { + // Bigram view conversion happens at the boundary; the key arrives typed. + let aligned_key_len = params.key.atom_len() / self.page_size * self.page_size; + if aligned_key_len == 0 { + return self.empty_match_result(); + } + // The walk reads only [0, aligned_key_len); the ragged tail never enters. + let key = params.key; + + let root_id = self.arena.root(); + let ( + value, + best_match_node_id, + best_match_device_node_id, + best_match_device_value_len, + full_kv_hit_length, + action, + ) = self.match_prefix_helper_(root_id, params.namespace, key, aligned_key_len); + self.match_post_processor_( + params, + root_id, + value, + best_match_node_id, + best_match_device_node_id, + best_match_device_value_len, + full_kv_hit_length, + action, + ) + } + + /// Walk the tree for `key`; returns matched value chunks, the best match, + /// the best device-resident match, its device value length, and any split action. + pub fn match_prefix_helper_( + &mut self, + root_id: NodeIdx_, + namespace: KeyNamespaceRef<'_>, + key: &K, + aligned_key_len: usize, + ) -> ( + Vec, + NodeIdx_, + NodeIdx_, + usize, + usize, + Option, + ) { + // Non-HiCache mode has only device-resident matches, so the scheduler + // device anchor follows the best match. In HiCache mode, host-backed + // nodes can also match, so we separately track the best device-resident + // match for scheduler prefix indices and locking. + let mut node_id = root_id; + // Walk cursor: atoms of `key` already matched. + let mut offset = 0; + let mut value: Vec = Vec::new(); + let mut best_match_node_id = node_id; + let mut best_match_device_node_id = node_id; + let mut best_match_device_value_len = 0; + let mut full_kv_hit_length = 0; + let mut action: Option = None; + let separate_device_match = self.enable_hicache; + let mut validators = Vec::with_capacity(self.components.len()); + let mut device_validators = if separate_device_match { + Vec::with_capacity(self.components.len()) + } else { + Vec::new() + }; + for i in 0..self.components.len() { + let component = Arc::clone(&self.components[i]); + if separate_device_match { + validators.push( + component.create_match_validator(self, /* match_device_only = */ false), + ); + device_validators + .push(component.create_match_validator(self, /* match_device_only = */ true)); + } else { + validators + .push(component.create_match_validator(self, /* match_device_only = */ true)); + } + } + + fn update_best_if_valid( + tree: &UnifiedTreeCore, + node_id: NodeIdx_, + value_len: usize, + separate_device_match: bool, + validators: &mut [Box, NodeIdx_) -> bool>], + device_validators: &mut [Box, NodeIdx_) -> bool>], + best_match_node_id: &mut NodeIdx_, + best_match_device_node_id: &mut NodeIdx_, + best_match_device_value_len: &mut usize, + ) { + // Every validator observes every node (stateful validators need the full walk). + let matched = validators + .iter_mut() + .fold(true, |acc, validator| validator(tree, node_id) & acc); + if matched { + *best_match_node_id = node_id; + } + if !separate_device_match { + if matched { + *best_match_device_value_len = value_len; + *best_match_device_node_id = node_id; + } + return; + } + if device_validators + .iter_mut() + .fold(true, |acc, validator| validator(tree, node_id) & acc) + { + *best_match_device_value_len = value_len; + *best_match_device_node_id = node_id; + } + } + + while offset < aligned_key_len { + let Some(child_id) = self.arena.child_on_page_in_namespace( + node_id, + namespace, + key.page_at(offset, self.page_size), + ) else { + break; + }; + let child = self.arena.node(child_id); + // HiCache: a dead node (evicted and not backuped) stops the traversal. + if child.evicted() && !child.backuped() { + break; + } + let prefix_len = key.match_len(offset, &child.key, self.page_size); + full_kv_hit_length += prefix_len; + if prefix_len < child.key.atom_len() { + let (split_node_id, split_action) = self.split_node_(child_id, prefix_len); + node_id = split_node_id; + action = split_action; + let node = self.arena.node(node_id); + if !node.evicted() { + value.push(node.device_value(FULL).shallow_clone()); + } + update_best_if_valid( + self, + node_id, + value.len(), + separate_device_match, + &mut validators, + &mut device_validators, + &mut best_match_node_id, + &mut best_match_device_node_id, + &mut best_match_device_value_len, + ); + break; + } + + if !child.evicted() { + value.push(child.device_value(FULL).shallow_clone()); + } + node_id = child_id; + update_best_if_valid( + self, + node_id, + value.len(), + separate_device_match, + &mut validators, + &mut device_validators, + &mut best_match_node_id, + &mut best_match_device_node_id, + &mut best_match_device_value_len, + ); + offset += prefix_len; + } + + ( + value, + best_match_node_id, + best_match_device_node_id, + best_match_device_value_len, + full_kv_hit_length, + action, + ) + } + + /// Assemble the MatchResult from the walk outputs. + pub fn match_post_processor_( + &mut self, + params: &MatchPrefixParams<'_, K>, + root_id: NodeIdx_, + value: Vec, + best_match_node_id: NodeIdx_, + best_match_device_node_id: NodeIdx_, + best_match_device_value_len: usize, + full_kv_hit_length: usize, + action: Option, + ) -> MatchResult { + for i in 0..self.components.len() { + // Full uses last_access ticks, not LRU. + if self.components[i].component_type() == BASE_COMPONENT_TYPE { + continue; + } + let component = Arc::clone(&self.components[i]); + component.refresh_lru(self, LRURefreshPhase::MatchEnd, best_match_node_id); + } + + // Re-stamp the matched path with fresh ticks, newest leaf-ward. + let mut path = Vec::new(); + let mut cur = Some(best_match_node_id); + while let Some(id) = cur { + path.push(id); + cur = self.arena.node(id).try_parent(); + } + let newest_tick = self + .arena + .get_and_batch_bump_access_counter(path.len() as i64); + for (i, id) in path.iter().enumerate() { + self.arena.node_mut(*id).last_access_counter = newest_tick - i as i64; + } + + // last_host_node will be used as the starting node for the subsequent + // `prefetch_from_storage` flow. We directly use best_match_node here, + // because best_match_node represents the node where all components + // have reached consensus on both device & host availability. + let last_host_node_id = if self.enable_hicache { + best_match_node_id + } else { + best_match_device_node_id + }; + + let device_indices = if best_match_device_value_len > 0 { + Tensor::cat(&value[..best_match_device_value_len], 0) + } else { + self.empty_device_indices.shallow_clone() + }; + let mut result = MatchResult { + device_indices, + last_device_node_id: self.arena.node(best_match_device_node_id).id, + last_host_node_id: self.arena.node(last_host_node_id).id, + best_match_node_id: self.arena.node(best_match_node_id).id, + host_hit_length: 0, + mamba_host_hit_length: 0, + mamba_branching_seqlen: None, + swa_host_hit_length: 0, + full_kv_hit_length, + cache_actions: Vec::new(), + }; + for i in 0..self.components.len() { + let component = Arc::clone(&self.components[i]); + result = component.finalize_match_result_in_tree_core( + self, + result, + params, + &value, + best_match_device_value_len, + ); + } + result.cache_actions = action.into_iter().collect(); + result + } + + /// An empty match: no device indices, every boundary anchored at the root. + pub fn empty_match_result(&self) -> MatchResult { + let root_id = self.arena.node(self.arena.root()).id; + MatchResult { + device_indices: self.empty_device_indices.shallow_clone(), + last_device_node_id: root_id, + last_host_node_id: root_id, + best_match_node_id: root_id, + host_hit_length: 0, + swa_host_hit_length: 0, + full_kv_hit_length: 0, + mamba_host_hit_length: 0, + mamba_branching_seqlen: None, + cache_actions: Vec::new(), + } + } + + /// Whether the node's FULL device value has been evicted. + pub fn is_full_device_evicted(&self, node_id: NodeId) -> bool { + let node_id = self.arena.resolve(node_id); + self.arena.node(node_id).evicted() + } + + /// Concatenate FULL device values from ``from_node`` up to (exclusive) + /// ``until_node``, in root order; empty tensor if the path is empty. + pub fn collect_full_device_indices( + &self, + from_node_id: NodeId, + until_node_id: NodeId, + ) -> Tensor { + let from_node_id = self.arena.resolve(from_node_id); + let until_node_id = self.arena.resolve(until_node_id); + let mut prefix_chunks: Vec = Vec::new(); + let mut node_id = from_node_id; + while node_id != until_node_id { + let node = self.arena.node(node_id); + prefix_chunks.push(node.device_value(FULL).shallow_clone()); + node_id = node.parent(); + } + if prefix_chunks.is_empty() { + return self.empty_device_indices.shallow_clone(); + } + prefix_chunks.reverse(); + Tensor::cat(&prefix_chunks, 0) + } + + /// Refresh a node's access tick and component LRU positions. + pub fn touch_node_(&mut self, node_id: NodeIdx_) { + let tick = self.arena.get_and_bump_access_counter(); + let node = self.arena.node_mut(node_id); + node.last_access_counter = tick; + if node.is_root() { + return; + } + for i in 0..self.components.len() { + // Full uses leaf sets, not LRU. + if self.components[i].component_type() == BASE_COMPONENT_TYPE { + continue; + } + let component = Arc::clone(&self.components[i]); + component.refresh_lru(self, LRURefreshPhase::Walkdown, node_id); + } + } + + /// Increment hit count; check whether a write backup should be fired. + pub fn inc_hit_count_and_check_(&mut self, node_id: NodeIdx_, chunked: bool) -> bool { + let node = self.arena.node_mut(node_id); + if node.evicted() || chunked { + return false; + } + if self.is_write_back { + return false; + } + node.hit_count += 1; + self.enable_hicache && !node.backuped() && node.hit_count >= self.write_through_threshold + } + + /// Insert device values to the tree per the provided key. + pub fn insert(&mut self, params: &InsertParams<'_, K>) -> InsertResult { + self.try_insert(params) + .unwrap_or_else(|error| panic!("{error}")) + } + + /// Fallible variant of [`Self::insert`]. + pub fn try_insert( + &mut self, + params: &InsertParams<'_, K>, + ) -> Result { + // Single-shot pump over the resumable walk: run every step inline and + // fold the step actions into the result for the caller to apply. + let mut actions = Vec::new(); + let mut step = self.try_begin_insert(params)?; + loop { + actions.append(&mut step.actions); + if let Some(mut result) = step.result { + result.cache_actions = actions; + return Ok(result); + } + step = self.try_resume_insert()?; + } + } + + /// Start the insert, running to its first barrier or completion. + pub fn begin_insert(&mut self, params: &InsertParams<'_, K>) -> InsertStepResult { + self.try_begin_insert(params) + .unwrap_or_else(|error| panic!("{error}")) + } + + /// Fallible variant of [`Self::begin_insert`]. + pub fn try_begin_insert( + &mut self, + params: &InsertParams<'_, K>, + ) -> Result { + // Insert walks are single-flight; a live walk means re-entrancy. + if self.ongoing_insert_walk_state.is_some() { + return Err(TreeCoreRuntimeError::ConcurrentInsertWalk); + } + // Bigram view conversion happens at the boundary; the key arrives typed. + let aligned_key_len = params.key.atom_len() / self.page_size * self.page_size; + if aligned_key_len == 0 { + // An empty insert still touches the root. + let root_id = self.arena.root(); + self.touch_node_(root_id); + { + let node = self.arena.node_mut(root_id); + node.priority = node.priority.max(params.priority); + } + return Ok(InsertStepResult { + actions: Vec::new(), + result: Some(InsertResult { + prefix_len: 0, + total_len: 0, + last_device_node_id: Some(self.arena.node(root_id).id), + inserted_host_node: None, + host_insert_dropped: false, + mamba_exist: true, + adopted_ranges: None, + cache_actions: Vec::new(), + }), + }); + } + let root_id = self.arena.root(); + self.touch_node_(root_id); + { + let node = self.arena.node_mut(root_id); + node.priority = node.priority.max(params.priority); + } + // The walk reads only [0, aligned_key_len); the ragged tail never enters. + self.ongoing_insert_walk_state = Some(InsertWalkState { + phase: InsertPhase::Walk, + node_id: root_id, + key: K::from(params.key.as_ref()[..aligned_key_len].to_vec()), + aligned_key_len, + value: params.value.narrow(0, 0, aligned_key_len as i64), + namespace: params.namespace.to_owned(), + prev_prefix_len: params.prev_prefix_len, + swa_evicted_seqlen: params.swa_evicted_seqlen, + mamba_value: params.mamba_value.as_ref().map(Tensor::shallow_clone), + chunked: params.chunked, + priority: params.priority, + track_adopted_ranges: params.track_adopted_ranges, + total_prefix_length: 0, + is_new_leaf: false, + target_node_id: None, + result: Some(InsertResult { + adopted_ranges: params.track_adopted_ranges.then(HashMap::new), + ..InsertResult::default() + }), + pending_actions: Vec::new(), + }); + Ok(self.advance_insert_()) + } + + /// Continue the suspended insert after its step actions were executed. + pub fn resume_insert(&mut self) -> InsertStepResult { + self.try_resume_insert() + .unwrap_or_else(|error| panic!("{error}")) + } + + /// Fallible variant of [`Self::resume_insert`]. + pub fn try_resume_insert(&mut self) -> Result { + if self.ongoing_insert_walk_state.is_none() { + return Err(TreeCoreRuntimeError::NoInFlightInsert); + } + Ok(self.advance_insert_()) + } + + /// Whether an insert walk is suspended at a barrier. + pub fn has_ongoing_insert(&self) -> bool { + self.ongoing_insert_walk_state.is_some() + } + + /// Finish the insert (idempotent); returns still-pending actions to drain. + pub fn end_insert(&mut self) -> Vec { + self.ongoing_insert_walk_state + .take() + .map(|state| state.pending_actions) + .unwrap_or_default() + } + + /// Run the in-flight insert to its next barrier or to completion. + fn advance_insert_(&mut self) -> InsertStepResult { + // The state moves out of self while steps run (they borrow the tree mutably). + let mut state = self + .ongoing_insert_walk_state + .take() + .expect("no in-flight insert"); + loop { + let flushed_len = state.pending_actions.len(); + match state.phase { + InsertPhase::Walk => self.insert_walk_step_(&mut state), + InsertPhase::Commit => self.insert_commit_step_(&mut state), + InsertPhase::Tail => { + self.insert_tail_step_(&mut state); + return InsertStepResult { + actions: state.pending_actions, + result: state.result, + }; + } + } + let new_actions = &state.pending_actions[flushed_len..]; + // Suspend only when a step emitted a non-deferrable action. + if !new_actions.is_empty() && !new_actions.iter().all(Self::is_deferrable_action_) { + let flushed = std::mem::take(&mut state.pending_actions); + self.ongoing_insert_walk_state = Some(state); + return InsertStepResult { + actions: flushed, + result: None, + }; + } + } + } + + /// Fire-and-forget actions safe to batch until the next barrier. + fn is_deferrable_action_(action: &CacheAction) -> bool { + matches!( + action, + CacheAction::FreeDeviceKV(_) + | CacheAction::FreeDeviceKVFullOnly(_) + | CacheAction::ReplaceWriteThroughOnNodeSplit { .. } + ) + } + + /// Process one walked node, appending its barrier actions to the state. + fn insert_walk_step_(&mut self, state: &mut InsertWalkState) { + // Walk cursor: atoms of `key` already matched (also the running prefix length). + let cursor = state.total_prefix_length; + let child_id = if cursor < state.aligned_key_len { + self.arena.child_on_page_in_namespace( + state.node_id, + state.namespace.as_ref(), + state.key.page_at(cursor, self.page_size), + ) + } else { + None + }; + let Some(child_id) = child_id else { + state.phase = InsertPhase::Commit; + return; + }; + let mut node_id = child_id; + self.touch_node_(node_id); + let node = self.arena.node(node_id); + let prefix_len = state.key.match_len(cursor, &node.key, self.page_size); + if prefix_len < node.key.atom_len() { + let (split_node_id, action) = self.split_node_(node_id, prefix_len); + node_id = split_node_id; + if let Some(action) = action { + state.pending_actions.push(action); + } + } + { + let node = self.arena.node_mut(node_id); + node.priority = node.priority.max(state.priority); + } + + let params = InsertParams { + key: &state.key, + namespace: state.namespace.as_ref(), + value: state.value.shallow_clone(), + prev_prefix_len: state.prev_prefix_len, + swa_evicted_seqlen: state.swa_evicted_seqlen, + mamba_value: state.mamba_value.as_ref().map(Tensor::shallow_clone), + chunked: state.chunked, + priority: state.priority, + track_adopted_ranges: state.track_adopted_ranges, + }; + if self.arena.node(node_id).evicted() { + self.unevict_node_on_insert_( + node_id, + &state.value.narrow(0, cursor as i64, prefix_len as i64), + ); + state + .result + .as_mut() + .expect("insert result exists during the walk") + .record_adopted_range(BASE_COMPONENT_TYPE, cursor, cursor + prefix_len); + // FULL was restored from the request's fresh KV. Aux + // components (e.g. SWA) may still hold tombstones and need + // to rebuild their value from the same slice. + for i in 0..self.components.len() { + if self.components[i].component_type() == BASE_COMPONENT_TYPE { + continue; + } + let component = Arc::clone(&self.components[i]); + component.recover_after_unevict( + self, + node_id, + prefix_len, + cursor, + ¶ms, + state + .result + .as_mut() + .expect("insert result exists during the walk"), + &mut state.pending_actions, + ); + } + } else { + let value_slice = state.value.narrow(0, cursor as i64, prefix_len as i64); + let mut consumed_from = prefix_len; + // Let each component claim ownership of overlapping KV slots. + for i in 0..self.components.len() { + let component = Arc::clone(&self.components[i]); + let comp_consumed_from = component.update_component_on_insert_overlap( + self, + node_id, + prefix_len, + cursor, + value_slice.shallow_clone(), + ¶ms, + state + .result + .as_mut() + .expect("insert result exists during the walk"), + &mut state.pending_actions, + ); + consumed_from = consumed_from.min(comp_consumed_from); + } + + let dup_start = state.prev_prefix_len.saturating_sub(cursor); + if dup_start < consumed_from { + state + .pending_actions + .push(CacheAction::FreeDeviceKV(vec![value_slice.narrow( + 0, + dup_start as i64, + (consumed_from - dup_start) as i64, + )])); + } + } + + if self.inc_hit_count_and_check_(node_id, state.chunked) { + let backup = self + .build_backup_kv_action_(self.arena.node(node_id), /* write_back = */ false); + state.pending_actions.push(CacheAction::BackupKV(backup)); + } + state.node_id = node_id; + state.total_prefix_length += prefix_len; + } + + /// Create the tail leaf and run the component commit hooks. + fn insert_commit_step_(&mut self, state: &mut InsertWalkState) { + // Create new leaf for remaining suffix. A leaf survives on its Full + // value alone; auxiliary components (SWA, Mamba) may legitimately hold + // only a tombstone for this span (e.g. the whole leaf is outside the SWA + // window). Materialize it anyway so the Full KV stays cacheable. + let target_node_id = if state.total_prefix_length < state.aligned_key_len { + state.is_new_leaf = true; + state + .result + .as_mut() + .expect("insert result exists during commit") + .record_adopted_range( + BASE_COMPONENT_TYPE, + state.total_prefix_length, + state.aligned_key_len, + ); + // The walk's only owned key: the unmatched suffix backing the new leaf. + let leaf_value = state.value.narrow( + 0, + state.total_prefix_length as i64, + (state.aligned_key_len - state.total_prefix_length) as i64, + ); + self.add_new_node_in_namespace_( + state.node_id, + K::from( + state.key.as_ref()[state.total_prefix_length..state.aligned_key_len].to_vec(), + ), + &leaf_value, + state.priority, + state.namespace.as_ref(), + ) + } else { + state.node_id + }; + state.target_node_id = Some(target_node_id); + + // Finalize: let each component attach its data to the target node. + // e.g. Mamba attaches mamba_value to the leaf node + // All hooks run before their emitted actions execute; an action failure + // fail-stops the process, so partial-commit state is never observed. + let result = state + .result + .as_mut() + .expect("insert result exists during commit"); + result.prefix_len = state.total_prefix_length; + result.last_device_node_id = Some(self.arena.node(target_node_id).id); + let params = InsertParams { + key: &state.key, + namespace: state.namespace.as_ref(), + value: state.value.shallow_clone(), + prev_prefix_len: state.prev_prefix_len, + swa_evicted_seqlen: state.swa_evicted_seqlen, + mamba_value: state.mamba_value.as_ref().map(Tensor::shallow_clone), + chunked: state.chunked, + priority: state.priority, + track_adopted_ranges: state.track_adopted_ranges, + }; + for i in 0..self.components.len() { + let component = Arc::clone(&self.components[i]); + component.commit_insert_component_data( + self, + target_node_id, + state.is_new_leaf, + ¶ms, + result, + &mut state.pending_actions, + ); + } + state.phase = InsertPhase::Tail; + } + + /// Whether an auxiliary component has new device data missing from Host. + fn needs_incremental_component_backup_(&self, node_id: NodeIdx_) -> bool { + self.components.iter().any(|component| { + component.component_type() != BASE_COMPONENT_TYPE + && component.needs_incremental_backup(self, node_id) + }) + } + + /// Check whether the insert target needs a Host backup. + fn should_backup_after_insert_( + &mut self, + state: &InsertWalkState, + target_node_id: NodeIdx_, + ) -> bool { + if state.is_new_leaf { + return self.inc_hit_count_and_check_(target_node_id, state.chunked); + } + + let node = self.arena.node(target_node_id); + self.enable_hicache + && !self.is_write_back + && node.backuped() + && node.write_through_pending_id.is_none() + && self.needs_incremental_component_backup_(target_node_id) + } + + /// Refresh the LRUs and append terminal backup actions. + fn insert_tail_step_(&mut self, state: &mut InsertWalkState) { + let target_node_id = state + .target_node_id + .expect("the commit step sets the target"); + if !self.arena.node(target_node_id).is_root() { + for i in 0..self.components.len() { + // Full uses leaf sets, not LRU. + if self.components[i].component_type() == BASE_COMPONENT_TYPE { + continue; + } + let component = Arc::clone(&self.components[i]); + component.refresh_lru(self, LRURefreshPhase::InsertEnd, target_node_id); + } + } + + if self.should_backup_after_insert_(state, target_node_id) { + let backup = self.build_backup_kv_action_( + self.arena.node(target_node_id), + /* write_back = */ false, + ); + state.pending_actions.push(CacheAction::BackupKV(backup)); + } + } + + /// Split `child` at `split_len`; returns the new prefix node and any split action. + pub fn split_node_( + &mut self, + child_id: NodeIdx_, + split_len: usize, + ) -> (NodeIdx_, Option) { + assert!( + split_len > 0 && split_len.is_multiple_of(self.page_size), + "split_node_: split_len {split_len} must be a nonzero page multiple" + ); + let page_size = self.page_size; + + // The new node takes the child's prefix, link position, and stats. + let child = self.arena.node(child_id); + let parent_id = child.parent(); + let child_namespace = child.namespace.clone(); + let (key_head, key_tail) = child.key.split_at(split_len); + // key_head keeps the original key's first page, which keys the parent's child map. + let parent_map_key = key_head.child_key(page_size); + let new_node_id = self.new_node_in_namespace_( + key_head, + parent_id, + child.priority, + child.hit_count, + Some(child.creation_counter), + child_namespace.as_ref(), + ); + self.arena.node_mut(new_node_id).children.insert( + (child_namespace.clone(), key_tail.child_key(page_size)), + child_id, + ); + + // The child's aux LRU cells detach while it is re-linked. + self.for_each_component_lru_( + child_id, + &mut |lru, node_id| lru.remove_node(node_id), + EvictLayer::Device, + /* skip_existing = */ false, + ); + + let child = self.arena.node_mut(child_id); + child.parent = Some(new_node_id); + child.key = key_tail; + let (new_node_hash, child_hash) = + crate::node::split_node_hash_value(child.hash_value.take(), split_len, self.page_size); + child.hash_value = child_hash; + self.arena.node_mut(new_node_id).hash_value = new_node_hash; + let child_handle = self.arena.node(child_id).id; + if let Some(mut parent_event_hashes) = self.salted_event_hashes.remove(&child_handle) { + let child_event_hashes = parent_event_hashes.split_off(split_len / self.page_size); + parent_event_hashes.shrink_to_fit(); + let new_node_handle = self.arena.node(new_node_id).id; + self.salted_event_hashes + .insert(new_node_handle, parent_event_hashes); + self.salted_event_hashes + .insert(child_handle, child_event_hashes); + } + + for i in 0..self.components.len() { + let component = Arc::clone(&self.components[i]); + component.redistribute_on_node_split(self, new_node_id, child_id); + } + let replaced = self + .arena + .insert_child_edge(parent_id, parent_map_key, new_node_id); + assert_eq!( + replaced, + Some(child_id), + "split_node_: the parent's page entry must map to the split child" + ); + + // Preserve the load-back pin across a split. + self.arena.node_mut(new_node_id).load_back_pending_id = + self.arena.node(child_id).load_back_pending_id; + + // A split of a backuped node tells the cache to fix its publish list. + let action = if let Some(ack_id) = self.arena.node(child_id).write_through_pending_id { + self.arena.node_mut(new_node_id).write_through_pending_id = Some(ack_id); + Some(CacheAction::ReplaceWriteThroughOnNodeSplit { + ack_id, + old_node_id: self.arena.node(child_id).id, + new_node_id: self.arena.node(new_node_id).id, + new_child_node_id: self.arena.node(child_id).id, + }) + } else { + None + }; + + self.for_each_component_lru_( + new_node_id, + &mut |lru, node_id| lru.insert_mru(node_id), + EvictLayer::Device, + /* skip_existing = */ true, + ); + self.for_each_component_lru_( + child_id, + &mut |lru, node_id| lru.insert_mru(node_id), + EvictLayer::Device, + /* skip_existing = */ true, + ); + let tick = self.arena.get_and_bump_access_counter(); + self.arena.node_mut(child_id).last_access_counter = tick; + + self.update_evictable_leaf_sets_(new_node_id); + self.update_evictable_leaf_sets_(child_id); + self.update_full_coexisting_host_tracking_(new_node_id); + (new_node_id, action) + } + + /// Create a leaf holding `value` under `parent`. + pub fn add_new_node_( + &mut self, + parent_id: NodeIdx_, + key: K, + value: &Tensor, + priority: i64, + extra_key: Option<&str>, + ) -> NodeIdx_ { + self.add_new_node_in_namespace_( + parent_id, + key, + value, + priority, + KeyNamespaceRef::new(extra_key, /* cache_salt = */ None), + ) + } + + pub fn add_new_node_in_namespace_( + &mut self, + parent_id: NodeIdx_, + key: K, + value: &Tensor, + priority: i64, + namespace: KeyNamespaceRef<'_>, + ) -> NodeIdx_ { + let page_size = self.page_size; + let child_map_key = key.child_key(page_size); + let new_node_id = self.new_node_in_namespace_( + key, parent_id, priority, /* hit_count = */ 0, /* creation_counter = */ None, + namespace, + ); + self.arena.set_device_value(new_node_id, FULL, value.copy()); + let displaced = self + .arena + .insert_child_edge(parent_id, child_map_key, new_node_id); + assert!( + displaced.is_none(), + "add_new_node_: parent {parent_id} already has a child on the new node's page" + ); + self.inc_evictable_size(FULL, value.size()[0] as usize); + if self.enable_storage { + let hash_values = self.arena.compute_node_hash_values(new_node_id, page_size); + self.arena.node_mut(new_node_id).hash_value = Some(hash_values); + } + + self.update_evictable_leaf_sets_(new_node_id); + self.update_evictable_leaf_sets_(parent_id); + self.record_store_event_(new_node_id, StorageMedium::Gpu); + new_node_id + } + + /// Restore an evicted node's Full device value from fresh KV indices + /// during insert. + pub fn unevict_node_on_insert_(&mut self, node_id: NodeIdx_, fresh_value: &Tensor) { + self.arena + .set_device_value(node_id, FULL, fresh_value.copy()); + self.inc_evictable_size(FULL, fresh_value.size()[0] as usize); + self.update_evictable_leaf_sets_(node_id); + self.update_full_coexisting_host_tracking_(node_id); + if let Some(parent_id) = self.arena.node(node_id).try_parent() { + self.update_evictable_leaf_sets_(parent_id); + } + self.record_store_event_(node_id, StorageMedium::Gpu); + } + + /// Update both device and host leaf sets for a node. + pub(crate) fn update_evictable_leaf_sets_(&mut self, node_id: NodeIdx_) { + let node = self.arena.node(node_id); + let is_evictable_device_leaf = self.is_evictable_device_leaf_(node); + let is_evictable_host_leaf = self.is_evictable_host_leaf_(node); + if is_evictable_device_leaf { + self.evictable_device_leaves.add(node_id); + } else { + self.evictable_device_leaves.discard(node_id); + } + if is_evictable_host_leaf { + self.evictable_host_leaves.add(node_id); + } else { + self.evictable_host_leaves.discard(node_id); + } + } + + /// Refresh Full's lazily maintained device/host coexistence registry. + pub(crate) fn update_full_coexisting_host_tracking_(&mut self, node_id: NodeIdx_) { + if self.is_settled_full_coexisting_host_node_(self.arena.node(node_id)) { + self.full_coexisting_host_nodes.add(node_id); + } else { + self.full_coexisting_host_nodes.discard(node_id); + } + } + + fn is_settled_full_coexisting_host_node_(&self, node: &Node) -> bool { + !node.is_root() + && node.has_device_value(FULL) + && node.has_host_value(FULL) + && node.write_through_pending_id.is_none() + && !node.is_load_back_pending() + } + + /// Apply lru_op to each aux component's LRU that has data on this node. + /// If skip_existing=True, skip components already in the target LRU list. + pub(crate) fn for_each_component_lru_( + &mut self, + node_id: NodeIdx_, + lru_op: &mut dyn FnMut(&mut UnifiedLRUList, NodeIdx_), + target: EvictLayer, + skip_existing: bool, + ) { + assert!( + target != EvictLayer::All, + "for_each_component_lru_: EvictLayer::All is not a single layer" + ); + for i in 0..self.components.len() { + let ct = self.components[i].component_type(); + // Full uses leaf sets, not LRU. + if ct == BASE_COMPONENT_TYPE { + continue; + } + let node = self.arena.node(node_id); + let slot = if target == EvictLayer::Host { + ValueSlotIdx::host(ct) + } else { + ValueSlotIdx::device(ct) + }; + if !node.has_value_(slot) { + continue; + } + let lru = self.lru_list_mut_(slot); + if skip_existing && lru.in_list(Some(node_id)) { + continue; + } + lru_op(lru, node_id); + } + } + + /// Register a component driver into the ordered fan-out list and the + /// by-type lookup slot; rejects duplicates. + pub(crate) fn register_component_( + &mut self, + component: Arc + Send + Sync>, + ) { + let component_type = component.component_type(); + let slot = &mut self.components_by_type[component_type.idx()]; + assert!( + slot.is_none(), + "duplicate component type {component_type:?}" + ); + *slot = Some(Arc::clone(&component)); + self.components.push(component); + } + + /// Panics if the component is not enabled, matching the python KeyError. + fn assert_component_enabled_(&self, component_type: ComponentType) { + let _ = self.component_by_type_(component_type); + } + + /// The component driver for `component_type`; panics if not enabled. + fn component_by_type_( + &self, + component_type: ComponentType, + ) -> Arc + Send + Sync> { + self.try_component_by_type_(component_type) + .unwrap_or_else(|| panic!("{component_type:?} component is not enabled")) + } + + /// The component driver for `component_type`, or None when not enabled. + fn try_component_by_type_( + &self, + component_type: ComponentType, + ) -> Option + Send + Sync>> { + // Cloning the Arc hands out an owned driver, leaving the registry unborrowed. + self.components_by_type[component_type.idx()].clone() + } + + /// Begin a component's device-eviction walk for up to request_cnt tokens. + pub fn evict_device_start(&mut self, component_type: ComponentType, request_cnt: usize) { + self.component_by_type_(component_type) + .evict_device_start(self, request_cnt); + } + + /// Advance one component eviction step and report whether it progressed. + pub fn evict_device_next_node( + &mut self, + component_type: ComponentType, + baseline: &HashMap, + ) -> (Option, EvictionStepResult) { + let mut tracker = baseline.clone(); + // The walk gates on the walked component's entry, so seed it. + tracker.entry(component_type).or_insert(0); + let mut result = EvictionStepResult::default(); + let node_id = self + .component_by_type_(component_type) + .evict_device_next_node( + self, + &mut tracker, + &mut result.device_frees, + &mut result.host_frees, + ); + for (ct, total) in tracker { + let delta = total - baseline.get(&ct).copied().unwrap_or(0); + if delta > 0 { + result.tracker.insert(ct, delta); + } + } + (node_id.map(|idx| self.arena.node(idx).id), result) + } + + /// Finish a component's device-eviction walk. + pub fn evict_device_end(&mut self, component_type: ComponentType) { + self.component_by_type_(component_type) + .evict_device_end(self); + } + + /// Evict one device leaf (demote if backuped, delete if write-through); + /// for an unbacked write-back node, return the BackupKV for the cache to + /// execute and then demote, else None. + pub fn evict_device_leaf( + &mut self, + node_id: NodeId, + is_write_back: bool, + ) -> (Option, EvictionStepResult) { + let node_id = self.arena.resolve(node_id); + let mut result = EvictionStepResult::default(); + { + let node = self.arena.node(node_id); + assert!( + self.is_evictable_device_leaf_(node), + "node {node_id} is not a D-leaf" + ); + } + if self.arena.node(node_id).backuped() { + self.demote_( + node_id, + &mut result.tracker, + &mut result.device_frees, + &mut result.host_frees, + ); + return (None, result); + } + if is_write_back { + let backup = self + .build_backup_kv_action_(self.arena.node(node_id), /* write_back = */ true); + return (Some(backup), result); + } + + // Write-through: node has no backup, delete entirely. + self.delete_unbacked_device_leaf_( + node_id, + &mut result.tracker, + &mut result.device_frees, + &mut result.host_frees, + ); + (None, result) + } + + /// Write-back fallback when a D-leaf's D->H backup fails under host + /// memory pressure: drop the subtree rooted at the unbacked leaf so + /// device eviction keeps making progress instead of leaving its KV + /// unevictable until host space frees up. + pub fn drop_subtree_no_host(&mut self, node_id: NodeId) -> (bool, EvictionStepResult) { + let node_id = self.arena.resolve(node_id); + let mut result = EvictionStepResult::default(); + { + let node = self.arena.node(node_id); + assert!( + self.is_evictable_device_leaf_(node), + "node {node_id} is not a D-leaf" + ); + // A failed backup never issues the D->H copy, so the subtree root has + // no host state and no in-flight DMA reading its device slots. + assert!(!node.backuped() && node.write_through_pending_id.is_none()); + if node.is_host_locked() { + return (false, result); + } + } + let mut descendants: Vec = Vec::new(); + let mut stack: Vec = self + .arena + .node(node_id) + .children + .values() + .copied() + .collect(); + while let Some(cur_id) = stack.pop() { + let cur = self.arena.node(cur_id); + if cur.is_device_locked() || cur.is_host_locked() { + return (false, result); + } + descendants.push(cur_id); + stack.extend(cur.children.values().copied()); + } + for &desc_id in descendants.iter().rev() { + { + let desc = self.arena.node(desc_id); + // Host-only by construction: a device descendant would contradict + // this node being a D-leaf, and D-leaves evict before ancestors. + assert!( + desc.evicted() && desc.backuped(), + "node {desc_id} not host-only" + ); + assert!(desc.write_through_pending_id.is_none()); + } + self.release_all_component_layers_( + desc_id, + StorageMedium::Cpu, + &mut result.tracker, + &mut result.device_frees, + &mut result.host_frees, + ); + self.remove_leaf_from_parent_(desc_id); + } + self.delete_unbacked_device_leaf_( + node_id, + &mut result.tracker, + &mut result.device_frees, + &mut result.host_frees, + ); + (true, result) + } + + /// Free every component layer on the node and detach it from the LRU + /// lists and evictable leaf sets. + pub fn release_all_component_layers_( + &mut self, + node_id: NodeIdx_, + medium: StorageMedium, + tracker: &mut HashMap, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + ) { + self.record_remove_event_(node_id, medium); + for i in 0..self.components.len() { + let component_type = self.components[i].component_type(); + self.evict_component_and_detach_lru_( + node_id, + component_type, + device_frees, + host_frees, + EvictLayer::All, + Some(tracker), + ); + } + self.evictable_device_leaves.discard(node_id); + self.evictable_host_leaves.discard(node_id); + } + + /// Delete a device leaf that has no host backup, freeing all layers. + pub fn delete_unbacked_device_leaf_( + &mut self, + node_id: NodeIdx_, + tracker: &mut HashMap, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + ) { + self.release_all_component_layers_( + node_id, + StorageMedium::Gpu, + tracker, + device_frees, + host_frees, + ); + let parent = self.arena.node(node_id).parent(); + self.remove_leaf_from_parent_(node_id); + self.update_evictable_leaf_sets_(parent); + self.iteratively_delete_tombstone_leaf_(parent, tracker, device_frees, host_frees); + } + + /// Evict a component's host-side resources. + pub fn drive_host_eviction( + &mut self, + component_type: ComponentType, + num_tokens: usize, + ) -> EvictionStepResult { + let mut result = EvictionStepResult::default(); + if let Some(component) = self.try_component_by_type_(component_type) { + // The drive gates on the driven component's entry, so seed it. + result.tracker.insert(component_type, 0); + if self.is_write_back { + component.reclaim_coexisting_host_values( + self, + num_tokens, + &mut result.tracker, + &mut result.device_frees, + &mut result.host_frees, + ); + } + component.drive_host_eviction( + self, + num_tokens, + &mut result.tracker, + &mut result.device_frees, + &mut result.host_frees, + ); + } + result + } + + pub(crate) fn can_reclaim_coexisting_host_value_( + &self, + node_id: NodeIdx_, + component_type: ComponentType, + ) -> bool { + let node = self.arena.node(node_id); + !node.is_root() + && node.has_device_value(component_type) + && node.has_host_value(component_type) + && node.write_through_pending_id.is_none() + && !node.is_load_back_pending() + && node.host_lock_ref(component_type) == 0 + } + + /// Free one component's host value while its device value remains resident. + pub(crate) fn release_coexisting_host_value_( + &mut self, + node_id: NodeIdx_, + component_type: ComponentType, + tracker: &mut HashMap, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + ) { + assert!( + self.can_reclaim_coexisting_host_value_(node_id, component_type), + "cannot reclaim coexisting {component_type:?} host value from node {node_id}" + ); + if component_type == BASE_COMPONENT_TYPE { + // BlockRemoved tracks Full host residency, not auxiliary slices. + self.record_remove_event_(node_id, StorageMedium::Cpu); + } + self.evict_component_and_detach_lru_( + node_id, + component_type, + device_frees, + host_frees, + EvictLayer::Host, + Some(tracker), + ); + let victim_id = self.arena.node(node_id).id; + self.write_back_coexist_reclaim_digest = next_coexist_reclaim_digest( + self.write_back_coexist_reclaim_digest, + victim_id, + component_type.idx(), + ); + } + + /// Atomically evict all components on a host leaf. + /// + /// All freed tokens are accumulated into *tracker*. + pub fn evict_host_leaf_( + &mut self, + node_id: NodeIdx_, + tracker: &mut HashMap, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + ) { + assert!( + self.is_evictable_host_leaf_(self.arena.node(node_id)), + "node {node_id} is not an H-leaf" + ); + self.record_remove_event_(node_id, StorageMedium::Cpu); + for i in 0..self.components.len() { + let component_type = self.components[i].component_type(); + let (_, host_freed) = self.evict_component_and_detach_lru_( + node_id, + component_type, + device_frees, + host_frees, + EvictLayer::All, + None, + ); + *tracker.entry(component_type).or_insert(0) += host_freed; + } + self.evictable_host_leaves.discard(node_id); + let parent = self.arena.node(node_id).parent(); + self.remove_leaf_from_parent_(node_id); + self.iteratively_delete_tombstone_leaf_(parent, tracker, device_frees, host_frees); + } + + /// Release a node's device KV once its host copy exists; the node stays in the + /// tree, now host-only. + pub fn demote(&mut self, node_id: NodeId) -> EvictionStepResult { + let node_id = self.arena.resolve(node_id); + let mut result = EvictionStepResult::default(); + // Skip a deferred demote when a load-back now pins the device indices. + if self.arena.node(node_id).is_load_back_pending() { + return result; + } + self.demote_( + node_id, + &mut result.tracker, + &mut result.device_frees, + &mut result.host_frees, + ); + result + } + + /// Fallible variant of [`Self::demote`]. + pub fn try_demote( + &mut self, + node_id: NodeId, + ) -> Result { + let node_id = self.try_resolve_node_handle_(node_id)?; + let mut result = EvictionStepResult::default(); + // Skip a deferred demote when a load-back now pins the device indices. + if self.arena.node(node_id).is_load_back_pending() { + return Ok(result); + } + { + let node = self.arena.node(node_id); + if node.evicted() || !node.backuped() { + return Err(TreeCoreRuntimeError::InvalidDemoteState { + node_id: node.id, + evicted: node.evicted(), + backuped: node.backuped(), + }); + } + } + self.demote_( + node_id, + &mut result.tracker, + &mut result.device_frees, + &mut result.host_frees, + ); + Ok(result) + } + + /// Drop a backed-up node's device value, keeping the host copy. + pub fn demote_( + &mut self, + node_id: NodeIdx_, + tracker: &mut HashMap, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + ) { + { + let node = self.arena.node(node_id); + assert!(!node.evicted() && node.backuped()); + } + self.evict_component_and_detach_lru_( + node_id, + BASE_COMPONENT_TYPE, + device_frees, + host_frees, + EvictLayer::Device, + Some(tracker), + ); + self.cascade_evict_( + node_id, + BASE_COMPONENT_TYPE, + tracker, + device_frees, + host_frees, + EvictLayer::Device, + ); + self.record_remove_event_(node_id, StorageMedium::Gpu); + + // after device eviction, insert aux components into host LRU. + self.for_each_component_lru_( + node_id, + &mut UnifiedLRUList::insert_mru, + EvictLayer::Host, + /* skip_existing = */ true, + ); + let parent = self.arena.node(node_id).parent(); + self.update_evictable_leaf_sets_(parent); + } + + /// Cascade eviction from trigger to lower-or-equal priority components. + pub fn cascade_evict_( + &mut self, + node_id: NodeIdx_, + trigger_component_type: ComponentType, + tracker: &mut HashMap, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + target: EvictLayer, + ) { + let is_leaf = match target { + EvictLayer::Device => self.evictable_device_leaves.contains(node_id), + EvictLayer::Host => self.evictable_host_leaves.contains(node_id), + EvictLayer::All => panic!("cascade_evict_: EvictLayer::All is not a single layer"), + }; + + let trigger_component = self.component_by_type_(trigger_component_type); + let trigger_priority = trigger_component.eviction_priority(is_leaf); + let trigger_internal_priority = + trigger_component.eviction_priority(/* is_leaf = */ false); + + for i in 0..self.components.len() { + let component = Arc::clone(&self.components[i]); + let ct = component.component_type(); + let should_evict = self + .should_cascade_evict_component_( + node_id, + trigger_component_type, + component.as_ref(), + target, + is_leaf, + trigger_priority, + trigger_internal_priority, + ) + .unwrap_or_else(|message| panic!("{message}")); + if !should_evict { + continue; + } + self.evict_component_and_detach_lru_( + node_id, + ct, + device_frees, + host_frees, + target, + Some(tracker), + ); + } + + // Now that all components (including SWA which depends on Full.value) + // have been freed, we can safely tombstone Full.value. + // This is deferred from evict_component because free_swa needs it. + if target == EvictLayer::Device && trigger_component_type == BASE_COMPONENT_TYPE { + let _ = self.arena.take_device_value(node_id, FULL); + } + + self.update_evictable_leaf_sets_(node_id); + } + + /// Decide whether one component participates in a cascade eviction. + /// + /// Lock violations are returned instead of panicking so the inspection + /// binding can translate them to Python ``AssertionError``. Production + /// cascade eviction converts the same error back into its existing panic. + fn should_cascade_evict_component_( + &self, + node_id: NodeIdx_, + trigger_component_type: ComponentType, + component: &dyn TreeComponent, + target: EvictLayer, + is_leaf: bool, + trigger_priority: i64, + trigger_internal_priority: i64, + ) -> Result { + let component_type = component.component_type(); + if component.eviction_priority(is_leaf) > trigger_priority + || component_type == trigger_component_type + { + return Ok(false); + } + + let node = self.arena.node(node_id); + let has_target_data = match target { + EvictLayer::Device => node.has_device_value(component_type), + EvictLayer::Host | EvictLayer::All => node.has_host_value(component_type), + }; + if !has_target_data { + return Ok(false); + } + + let lock_ref = node.device_lock_ref(component_type); + let host_lock_ref = node.host_lock_ref(component_type); + // A component whose true internal priority outranks the trigger is + // present only because leaf-collapse flattened priorities. Its lock is + // a legitimate pin; a lower-priority component's lock is a strand. + if component.eviction_priority(/* is_leaf = */ false) >= trigger_internal_priority { + if target.contains(EvictLayer::Device) && lock_ref != 0 { + return Ok(false); + } + if target.contains(EvictLayer::Host) && host_lock_ref != 0 { + return Ok(false); + } + } + if target.contains(EvictLayer::Device) && lock_ref != 0 { + return Err(format!( + "cascade_evict_: a {component_type:?} device lock strands node {node_id}" + )); + } + if target.contains(EvictLayer::Host) && host_lock_ref != 0 { + return Err(format!( + "cascade_evict_: a {component_type:?} host lock strands node {node_id}" + )); + } + Ok(true) + } + + /// Unlink a leaf from its parent. + pub fn remove_leaf_from_parent_(&mut self, node_id: NodeIdx_) { + // Arena slots are reused, so discard tracking before freeing the node. + self.full_coexisting_host_nodes.discard(node_id); + self.salted_event_hashes + .remove(&self.arena.node(node_id).id); + // The arena is the registry: freeing detaches by page key and recycles the slot. + self.arena + .free_leaf(node_id) + .expect("remove_leaf_from_parent_: a deletable leaf"); + } + + /// Evict one component on the node and detach its LRU entries. + pub fn evict_component_and_detach_lru_( + &mut self, + node_id: NodeIdx_, + component_type: ComponentType, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + target: EvictLayer, + tracker: Option<&mut HashMap>, + ) -> (usize, usize) { + let component = self.component_by_type_(component_type); + let (device_freed, host_freed) = + component.evict_component(self, node_id, device_frees, host_frees, target); + if let Some(tracker) = tracker { + let freed = if target.contains(EvictLayer::Device) { + device_freed + } else { + host_freed + }; + *tracker.entry(component_type).or_insert(0) += freed; + } + + // Detach from the targeted LRU list(s). + if target.contains(EvictLayer::Device) { + let lru = self.device_lru_list_mut(component_type); + if lru.in_list(Some(node_id)) { + lru.remove_node(node_id); + } + } + if target.contains(EvictLayer::Host) { + let lru = self.host_lru_list_mut(component_type); + if lru.in_list(Some(node_id)) { + lru.remove_node(node_id); + } + } + (device_freed, host_freed) + } + + /// Walk up from *deleted_node* and cascade-delete childless ancestors. + /// + /// Only the Full (base) component decides whether a node survives: + /// - Full device present → keep as D-leaf + /// - Full host present → keep as H-leaf + /// - neither → evict all remaining data, delete, continue up + pub fn iteratively_delete_tombstone_leaf_( + &mut self, + deleted_node_parent_id: NodeIdx_, + tracker: &mut HashMap, + device_frees: &mut HashMap>, + host_frees: &mut HashMap>, + ) { + let ct = BASE_COMPONENT_TYPE; + let mut cur = deleted_node_parent_id; + loop { + let node = self.arena.node(cur); + if node.is_root() || !node.children.is_empty() { + break; + } + if node.is_device_locked() || node.is_host_locked() { + break; + } + let has_device = node.values[ct.idx()].value.is_some(); + let has_host = node.has_host_value(ct); + + if has_device { + self.update_evictable_leaf_sets_(cur); + break; + } + + // Full device absent — clean up orphaned aux device data. + for i in 0..self.components.len() { + let component = Arc::clone(&self.components[i]); + if self.arena.has_device_value(cur, component.component_type()) { + self.evict_component_and_detach_lru_( + cur, + component.component_type(), + device_frees, + host_frees, + EvictLayer::Device, + Some(tracker), + ); + } + } + + if has_host { + self.update_evictable_leaf_sets_(cur); + break; + } + + // Full absent on both layers — evict remaining host data, delete. + for i in 0..self.components.len() { + let component = Arc::clone(&self.components[i]); + if self.arena.has_host_value(cur, component.component_type()) { + self.evict_component_and_detach_lru_( + cur, + component.component_type(), + device_frees, + host_frees, + EvictLayer::Host, + Some(tracker), + ); + } + } + + self.evictable_host_leaves.discard(cur); + let parent = self.arena.node(cur).parent(); + self.remove_leaf_from_parent_(cur); + self.update_evictable_leaf_sets_(parent); + cur = parent; + } + } + + /// Whether the node is an evictable Full device leaf. + pub(crate) fn is_evictable_device_leaf_(&self, node: &Node) -> bool { + if node.is_root() || node.evicted() { + return false; + } + if node.is_device_locked() { + return false; + } + if node.is_load_back_pending() { + return false; + } + if node + .children + .values() + .any(|&child_id| self.arena.has_device_value(child_id, FULL)) + { + return false; + } + true + } + + /// Whether the node is an evictable Full host leaf. + fn is_evictable_host_leaf_(&self, node: &Node) -> bool { + if node.is_root() || !node.evicted() { + return false; + } + if !node.backuped() { + return false; + } + if node.is_load_back_pending() { + return false; + } + if node.is_host_locked() { + return false; + } + if !node.children.is_empty() { + return false; + } + true + } + + /// Mark the host tier (HiCache) as wired. + pub fn set_hicache_enabled(&mut self) { + self.enable_hicache = true; + } + + /// Whether the storage tier (L3) is wired; storage attaches after tree construction. + pub fn set_enable_storage(&mut self, value: bool) { + self.enable_storage = value; + } + + // ==== KV cache placement events ==== + + /// Append an event, coalescing it with a compatible queue tail. + fn enqueue_kv_event_(&mut self, event: KvCacheEvent) { + match (self.kv_event_queue.last_mut(), event) { + ( + Some(KvCacheEvent::BlockRemoved { + block_hashes: tail_hashes, + medium: tail_medium, + }), + KvCacheEvent::BlockRemoved { + mut block_hashes, + medium, + }, + ) if *tail_medium == medium => tail_hashes.append(&mut block_hashes), + ( + Some(KvCacheEvent::BlockStored { + block_hashes: tail_hashes, + token_ids: tail_token_ids, + block_size: tail_block_size, + medium: tail_medium, + cache_salt: tail_cache_salt, + .. + }), + KvCacheEvent::BlockStored { + mut block_hashes, + parent_block_hash, + mut token_ids, + block_size, + medium, + cache_salt, + }, + ) if *tail_medium == medium + && *tail_block_size == block_size + && *tail_cache_salt == cache_salt + && !tail_hashes.is_empty() + && parent_block_hash == tail_hashes.last().copied() => + { + tail_hashes.append(&mut block_hashes); + tail_token_ids.append(&mut token_ids); + } + (_, event) => self.kv_event_queue.push(event), + } + } + + /// Fill the salted external event-hash chain through `node_id`. + fn ensure_salted_event_hashes_(&mut self, node_id: NodeIdx_) { + let cache_salt = self + .arena + .node(node_id) + .namespace + .cache_salt_arc() + .expect("salted event hashing requires cache_salt"); + let node_handle = self.arena.node(node_id).id; + if self.salted_event_hashes.contains_key(&node_handle) { + return; + } + + let mut missing = Vec::new(); + let mut cursor = Some(node_id); + let mut prior = None; + while let Some(id) = cursor { + let node = self.arena.node(id); + if node.is_root() || node.key.atom_len() == 0 { + break; + } + assert_eq!( + node.namespace.cache_salt(), + Some(cache_salt.as_ref()), + "radix path contains mismatched cache_salt values" + ); + if let Some(hashes) = self.salted_event_hashes.get(&node.id) { + prior = hashes.last().copied(); + break; + } + missing.push(id); + cursor = node.try_parent(); + } + + let mut prior = prior.unwrap_or_else(|| { + let mut hasher = Sha256::new(); + hasher.update(b"sglang-cache-salt-v1\0"); + hasher.update(cache_salt.as_bytes()); + hasher.finalize().into() + }); + for id in missing.into_iter().rev() { + let (handle, hashes) = { + let node = self.arena.node(id); + ( + node.id, + crate::node::get_hash_digests::( + node.key.as_ref(), + Some(&prior), + self.page_size, + ), + ) + }; + if let Some(last) = hashes.last() { + prior = *last; + } + self.salted_event_hashes.insert(handle, hashes); + } + } + + /// Build one BlockStored per page and coalesce compatible queue neighbors. + fn record_store_event_(&mut self, node_id: NodeIdx_, medium: StorageMedium) { + if !self.enable_kv_cache_events { + return; + } + if self.arena.node(node_id).hash_value.is_none() { + let hash_values = self.arena.compute_node_hash_values(node_id, self.page_size); + self.arena.node_mut(node_id).hash_value = Some(hash_values); + } + let cache_salt = self.arena.node(node_id).namespace.cache_salt_arc(); + if cache_salt.is_some() { + self.ensure_salted_event_hashes_(node_id); + } + let events = { + let node = self.arena.node(node_id); + let mut parent_block_hash = node.parent.and_then(|parent_id| { + let parent = self.arena.node(parent_id); + if cache_salt.is_some() { + self.salted_event_hashes + .get(&parent.id) + .and_then(|hashes| hashes.last()) + .map(crate::node::hash_digest_to_int64) + } else { + parent + .get_last_hash_value() + .map(crate::node::hash_str_to_int64) + } + }); + let num_pages = node.key.atom_len().div_ceil(self.page_size); + let mut events = Vec::with_capacity(num_pages); + let mut append_event = |page: &[K::Atom], block_hash| { + events.push(KvCacheEvent::BlockStored { + block_hashes: vec![block_hash], + parent_block_hash, + token_ids: page.to_vec(), + block_size: page.len(), + medium, + cache_salt: cache_salt.clone(), + }); + parent_block_hash = Some(block_hash); + }; + if cache_salt.is_some() { + let hashes = &self.salted_event_hashes[&node.id]; + assert!( + hashes.len() >= num_pages, + "store event: {} page hashes for {num_pages} pages", + hashes.len() + ); + for (page, digest) in node.key.as_ref().chunks(self.page_size).zip(hashes) { + append_event(page, crate::node::hash_digest_to_int64(digest)); + } + } else { + let hashes = node.hash_value.as_ref().expect("hashed above"); + assert!( + hashes.len() >= num_pages, + "store event: {} page hashes for {num_pages} pages", + hashes.len() + ); + for (page, hash) in node.key.as_ref().chunks(self.page_size).zip(hashes) { + append_event(page, crate::node::hash_str_to_int64(hash)); + } + } + events + }; + for event in events { + self.enqueue_kv_event_(event); + } + } + + /// Queue one BlockRemoved carrying all the node's page hashes; hashes lazily if needed. + fn record_remove_event_(&mut self, node_id: NodeIdx_, medium: StorageMedium) { + if !self.enable_kv_cache_events { + return; + } + if self.arena.node(node_id).hash_value.is_none() { + let hash_values = self.arena.compute_node_hash_values(node_id, self.page_size); + self.arena.node_mut(node_id).hash_value = Some(hash_values); + } + let cache_salt = self.arena.node(node_id).namespace.cache_salt_arc(); + if cache_salt.is_some() { + self.ensure_salted_event_hashes_(node_id); + } + let node = self.arena.node(node_id); + let num_pages = node.key.atom_len().div_ceil(self.page_size); + let block_hashes: Vec = if cache_salt.is_some() { + self.salted_event_hashes[&node.id][..num_pages] + .iter() + .map(crate::node::hash_digest_to_int64) + .collect() + } else { + node.hash_value.as_ref().expect("hashed above")[..num_pages] + .iter() + .map(|hash| crate::node::hash_str_to_int64(hash)) + .collect() + }; + if !block_hashes.is_empty() { + self.enqueue_kv_event_(KvCacheEvent::BlockRemoved { + block_hashes, + medium, + }); + } + } + + /// Queue the all-cleared marker. + pub fn record_all_cleared_event(&mut self) { + if self.enable_kv_cache_events { + self.enqueue_kv_event_(KvCacheEvent::AllBlocksCleared); + } + } + + /// Take all queued events, leaving the queue empty. + pub fn take_events(&mut self) -> Vec> { + std::mem::take(&mut self.kv_event_queue) + } + + /// Mark the SWA host pool as wired; pools attach after tree construction. + pub fn set_has_swa_host_pool(&mut self) { + self.has_swa_host_pool = true; + } + + /// Insert a host-side (backuped) tree path descending from the given node. + pub fn insert_host( + &mut self, + node_id: NodeId, + extra_key: Option<&str>, + key: K, + host_value: Tensor, + hash_value: Vec, + ) -> InsertResult { + self.insert_host_in_namespace( + node_id, + KeyNamespaceRef::new(extra_key, /* cache_salt = */ None), + key, + host_value, + hash_value, + ) + } + + pub fn insert_host_in_namespace( + &mut self, + node_id: NodeId, + namespace: KeyNamespaceRef<'_>, + key: K, + host_value: Tensor, + hash_value: Vec, + ) -> InsertResult { + self.try_insert_host_in_namespace(node_id, namespace, key, host_value, hash_value) + .unwrap_or_else(|error| panic!("{error}")) + } + + pub fn try_insert_host_in_namespace( + &mut self, + node_id: NodeId, + namespace: KeyNamespaceRef<'_>, + key: K, + host_value: Tensor, + hash_value: Vec, + ) -> Result { + let total_len = key.atom_len(); + let mut node_id = self.arena.resolve(node_id); + let anchor = self.arena.node(node_id); + if !anchor.is_root() && anchor.namespace.as_ref() != namespace { + return Err(TreeCoreRuntimeError::InsertHostNamespaceMismatch { node_id: anchor.id }); + } + self.touch_node_(node_id); + if total_len == 0 { + return Ok(InsertResult { + prefix_len: 0, + total_len: 0, + last_device_node_id: None, + inserted_host_node: None, + host_insert_dropped: false, + mamba_exist: true, + adopted_ranges: None, + cache_actions: Vec::new(), + }); + } + + // Walk cursor: atoms of `key` already matched (also the running prefix length). + let mut matched_length = 0; + let mut cache_actions: Vec = Vec::new(); + while matched_length < total_len { + let Some(child_id) = self.arena.child_on_page_in_namespace( + node_id, + namespace, + key.page_at(matched_length, self.page_size), + ) else { + break; + }; + node_id = child_id; + self.touch_node_(node_id); + let node = self.arena.node(node_id); + let prefix_len = key.match_len(matched_length, &node.key, self.page_size); + let node_key_len = node.key.atom_len(); + matched_length += prefix_len; + + if prefix_len < node_key_len { + let (split_node_id, action) = self.split_node_(node_id, prefix_len); + node_id = split_node_id; + if let Some(action) = action { + cache_actions.push(action); + } + } + } + + let mut result = InsertResult { + prefix_len: matched_length, + total_len, + last_device_node_id: None, + inserted_host_node: None, + host_insert_dropped: false, + mamba_exist: false, + adopted_ranges: None, + cache_actions, + }; + if matched_length == total_len { + let node = self.arena.node(node_id); + if !node.is_root() && node.has_host_value(FULL) { + result.inserted_host_node = Some(self.arena.node(node_id).id); + } + return Ok(result); + } + + // Under write-through, a host-only suffix below a device-only parent + // would violate the invariant that a backed-up child has a backed-up + // parent. Keep any split actions produced by the walk, but do not + // materialize the suffix. + let parent = self.arena.node(node_id); + if !self.is_write_back && !parent.is_root() && !parent.backuped() { + result.host_insert_dropped = true; + return Ok(result); + } + + let priority = self.arena.node(node_id).priority; + let new_node_id = self.new_node_in_namespace_( + key.suffix(matched_length), + node_id, + priority, + /* hit_count = */ 0, + /* creation_counter = */ None, + namespace, + ); + { + // The suffix moves into a right-sized list; the matched head drops. + let mut hash_value = hash_value; + let suffix = hash_value.split_off(matched_length / self.page_size); + self.arena.node_mut(new_node_id).hash_value = Some(suffix); + } + self.arena.set_host_value( + new_node_id, + FULL, + host_value + .narrow( + 0, + matched_length as i64, + (total_len - matched_length) as i64, + ) + .copy(), + ); + let child_map_key = self.arena.node(new_node_id).key.child_key(self.page_size); + let displaced = self + .arena + .insert_child_edge(node_id, child_map_key, new_node_id); + assert!( + displaced.is_none(), + "insert_host: parent {node_id} already has a child on the new node's page" + ); + self.update_evictable_leaf_sets_(new_node_id); + self.update_evictable_leaf_sets_(node_id); + result.inserted_host_node = Some(self.arena.node(new_node_id).id); + Ok(result) + } + + /// Read a node's device->host backup spec (device value + component transfers) now. + pub fn build_backup_spec( + &self, + node_id: NodeId, + ) -> (Tensor, HashMap>) { + self.build_backup_spec_(self.arena.node(self.arena.resolve(node_id))) + } + + /// Gather device value backup spec. + pub fn build_backup_spec_( + &self, + node: &Node, + ) -> (Tensor, HashMap>) { + // Overlapping backup chains may revisit a node whose Full KV already + // has a host copy. Keep building transfers for auxiliary components, + // but do not allocate and overwrite Full host KV a second time. + let device_value = if node.backuped() { + self.empty_device_indices.shallow_clone() + } else { + node.device_value(FULL).shallow_clone() + }; + let mut comp_xfers: HashMap> = HashMap::new(); + for i in 0..self.components.len() { + let component_type = self.components[i].component_type(); + if component_type == BASE_COMPONENT_TYPE { + continue; + } + let transfers = self.components[i] + .build_hicache_transfers( + self, + node.idx, + CacheTransferPhase::BackupHost, + /* mamba_pool_idx = */ None, + /* host_indices = */ None, + /* token_ids = */ None, + /* prefetch_tokens = */ 0, + /* last_hash = */ None, + ) + .unwrap(); + if let Some(transfers) = transfers + && !transfers.is_empty() + { + comp_xfers.insert(component_type, transfers); + } + } + (device_value, comp_xfers) + } + + /// Gather a node's device->storage backup spec; None if the node is not backuped. + pub fn build_storage_backup_spec( + &self, + node_id: NodeId, + pass_prefix_keys: bool, + ) -> Option { + let node_id = self.arena.resolve(node_id); + let node = self.arena.node(node_id); + if !node.backuped() { + return None; + } + let prefix_keys = pass_prefix_keys.then(|| self.arena.prefix_hash_values(node.parent)); + let mut comp_xfers: HashMap> = HashMap::new(); + for i in 0..self.components.len() { + let component_type = self.components[i].component_type(); + if component_type == BASE_COMPONENT_TYPE { + continue; + } + let transfers = self.components[i] + .build_hicache_transfers( + self, + node_id, + CacheTransferPhase::BackupStorage, + /* mamba_pool_idx = */ None, + /* host_indices = */ None, + /* token_ids = */ None, + /* prefetch_tokens = */ 0, + /* last_hash = */ None, + ) + .unwrap(); + if let Some(transfers) = transfers + && !transfers.is_empty() + { + comp_xfers.insert(component_type, transfers); + } + } + Some(StorageBackupSpec { + host_value: node.host_value(FULL).shallow_clone(), + token_ids: K::raw_token_ids(node.key.as_ref()).into_owned(), + hash_value: node.hash_value.clone(), + prefix_keys, + comp_xfers, + }) + } + + /// Route a build_hicache_transfers call to the component for the given type. + pub fn build_hicache_transfers( + &self, + component_type: ComponentType, + node_id: NodeId, + phase: CacheTransferPhase, + host_indices: Option, + token_ids: Option<&[i64]>, + prefetch_tokens: usize, + last_hash: Option<&str>, + ) -> Option> { + self.try_build_hicache_transfers( + component_type, + node_id, + phase, + host_indices, + token_ids, + prefetch_tokens, + last_hash, + ) + .unwrap() + } + + pub fn try_build_hicache_transfers( + &self, + component_type: ComponentType, + node_id: NodeId, + phase: CacheTransferPhase, + host_indices: Option, + token_ids: Option<&[i64]>, + prefetch_tokens: usize, + last_hash: Option<&str>, + ) -> Result>, TreeCoreRuntimeError> { + let node_id = self.try_resolve_node_handle_(node_id)?; + self.component_by_type_(component_type) + .build_hicache_transfers( + self, + node_id, + phase, + /* mamba_pool_idx = */ None, + host_indices, + token_ids, + prefetch_tokens, + last_hash, + ) + } + + /// Build the H->D load-back KV transfer plus per-component aux transfers. + pub fn build_load_back_spec( + &self, + node_id: NodeId, + req: Option<&Req>, + ) -> (PoolTransfer, HashMap>) { + self.try_build_load_back_spec(node_id, req).unwrap() + } + + pub fn try_build_load_back_spec( + &self, + node_id: NodeId, + req: Option<&Req>, + ) -> Result<(PoolTransfer, HashMap>), TreeCoreRuntimeError> + { + let anchor_id = node_id; + let node_id = self.try_resolve_node_handle_(node_id)?; + // Component hooks take primitives, not Req: extract its fields here. + let mamba_pool_idx = req.and_then(|r| r.mamba_pool_idx.as_ref()); + let mut kv_transfers = self + .component_by_type_(BASE_COMPONENT_TYPE) + .build_hicache_transfers( + self, + node_id, + CacheTransferPhase::LoadBack, + /* mamba_pool_idx = */ None, + /* host_indices = */ None, + /* token_ids = */ None, + /* prefetch_tokens = */ 0, + /* last_hash = */ None, + )? + .unwrap(); + let kv_xfer = kv_transfers.remove(0); + let mut comp_xfers: HashMap> = HashMap::new(); + for i in 0..self.components.len() { + let component_type = self.components[i].component_type(); + if component_type == BASE_COMPONENT_TYPE { + continue; + } + let transfers = self.components[i].build_hicache_transfers( + self, + node_id, + CacheTransferPhase::LoadBack, + mamba_pool_idx.map(Tensor::shallow_clone), + /* host_indices = */ None, + /* token_ids = */ None, + /* prefetch_tokens = */ 0, + /* last_hash = */ None, + )?; + if let Some(transfers) = transfers + && !transfers.is_empty() + { + comp_xfers.insert(component_type, transfers); + } + } + // Reject transfers that would claim a node pinned by another load-back anchor. + let any_foreign_pin = kv_xfer + .nodes_to_load + .iter() + .chain( + comp_xfers + .values() + .flatten() + .filter_map(|xfer| xfer.nodes_to_load.as_ref()), + ) + .flatten() + .any(|&pinned_id| { + let pinned_idx = self.arena.resolve(pinned_id); + self.arena + .node(pinned_idx) + .load_back_pending_id + .is_some_and(|id| id != anchor_id) + }); + if any_foreign_pin { + let empty_kv = PoolTransfer { + name: PoolName::Kv, + host_indices: Some(Tensor::empty([0], (Kind::Int64, tch::Device::Cpu))), + nodes_to_load: Some(Vec::new()), + ..Default::default() + }; + return Ok((empty_kv, HashMap::new())); + } + Ok((kv_xfer, comp_xfers)) + } + + fn try_resolve_node_handle_(&self, node_id: NodeId) -> Result { + self.arena + .try_resolve(node_id) + .ok_or(TreeCoreRuntimeError::NodeNotAllocated { node_id }) + } + + /// The anchor node's caller-defined key and cache salt. + pub fn prefetch_anchor_info(&self, node_id: NodeId) -> (Option, Option) { + self.try_prefetch_anchor_info(node_id) + .unwrap_or_else(|error| panic!("{error}")) + } + + pub fn try_prefetch_anchor_info( + &self, + node_id: NodeId, + ) -> Result<(Option, Option), TreeCoreRuntimeError> { + let node_id = self.try_resolve_node_handle_(node_id)?; + Ok(( + self.arena.node_extra_key(node_id).map(str::to_string), + self.arena.node_cache_salt(node_id).map(str::to_string), + )) + } + + /// Whether the node's Full KV is present on host. + pub fn node_backuped(&self, node_id: NodeId) -> bool { + self.try_node_backuped(node_id) + .unwrap_or_else(|error| panic!("{error}")) + } + + pub fn try_node_backuped(&self, node_id: NodeId) -> Result { + let node_id = self.try_resolve_node_handle_(node_id)?; + Ok(self.arena.node(node_id).backuped()) + } + + /// Whether the node is a (default or named) root. + pub fn is_root(&self, node_id: NodeId) -> bool { + self.try_is_root(node_id) + .unwrap_or_else(|error| panic!("{error}")) + } + + pub fn try_is_root(&self, node_id: NodeId) -> Result { + let node_id = self.try_resolve_node_handle_(node_id)?; + Ok(self.arena.node(node_id).is_root()) + } + + /// The node's last page hash, or None when it was never hashed. + pub fn get_last_hash_value(&self, node_id: NodeId) -> Option { + self.try_get_last_hash_value(node_id) + .unwrap_or_else(|error| panic!("{error}")) + } + + pub fn try_get_last_hash_value( + &self, + node_id: NodeId, + ) -> Result, TreeCoreRuntimeError> { + let node_id = self.try_resolve_node_handle_(node_id)?; + Ok(self + .arena + .node(node_id) + .get_last_hash_value() + .map(str::to_string)) + } + + /// The hash chain of the node's ancestors, in root-to-parent order. + pub fn get_prefix_hash_values(&self, node_id: NodeId) -> Vec { + self.try_get_prefix_hash_values(node_id) + .unwrap_or_else(|error| panic!("{error}")) + } + + pub fn try_get_prefix_hash_values( + &self, + node_id: NodeId, + ) -> Result, TreeCoreRuntimeError> { + let node_id = self.try_resolve_node_handle_(node_id)?; + Ok(self + .arena + .prefix_hash_values(self.arena.node(node_id).parent)) + } + + /// The hash values owned by this node, excluding its ancestors. + pub fn get_hash_values(&self, node_id: NodeId) -> Vec { + self.try_get_hash_values(node_id) + .unwrap_or_else(|error| panic!("{error}")) + } + + pub fn try_get_hash_values( + &self, + node_id: NodeId, + ) -> Result, TreeCoreRuntimeError> { + let node_id = self.try_resolve_node_handle_(node_id)?; + Ok(self + .arena + .node(node_id) + .hash_value + .clone() + .unwrap_or_default()) + } + + pub fn snapshot_buffer_backup( + &self, + node_id: NodeId, + pass_prefix_keys: bool, + ) -> Option { + let node_id = self.arena.try_resolve(node_id)?; + let node = self.arena.node(node_id); + if node.is_root() || !node.has_device_value(FULL) { + return None; + } + let hash_values = node.hash_value.as_ref()?.clone(); + if hash_values.is_empty() { + return None; + } + let parent_node_id = node.try_parent()?; + let parent = self.arena.node(parent_node_id); + Some(BufferBackupSnapshot { + node_id: node.id, + parent_node_id: parent.id, + parent_is_root: parent.is_root(), + parent_last_hash: parent.get_last_hash_value().map(str::to_string), + token_ids: K::raw_token_ids(node.key.as_ref()).into_owned(), + extra_key: node.namespace.extra_key().map(str::to_string), + cache_salt: node.namespace.cache_salt().map(str::to_string), + is_bigram: K::IS_BIGRAM, + hash_values, + prefix_keys: pass_prefix_keys.then(|| self.arena.prefix_hash_values(node.parent)), + }) + } + + pub fn validate_buffer_backup( + &self, + node_id: NodeId, + expected_key_length: usize, + ) -> Option { + let node_id = self.arena.try_resolve(node_id)?; + let node = self.arena.node(node_id); + if !node.has_device_value(FULL) || node.key.atom_len() != expected_key_length { + return None; + } + let parent_node_id = node.try_parent()?; + let parent = self.arena.node(parent_node_id); + Some(BufferBackupState { + parent_node_id: parent.id, + parent_is_root: parent.is_root(), + parent_last_hash: parent.get_last_hash_value().map(str::to_string), + }) + } + + /// Hash every node built while storage was disabled. + pub fn backfill_missing_hash_values(&mut self) -> usize { + let root_id = self.arena.root(); + let mut filled = 0; + for node_id in self.collect_all_nodes_() { + if node_id == root_id || self.arena.node(node_id).hash_value.is_some() { + continue; + } + let hash_values = self.arena.compute_node_hash_values(node_id, self.page_size); + self.arena.node_mut(node_id).hash_value = Some(hash_values); + filled += 1; + } + filled + } + + /// The NodeId anchoring matches; the single root serves every namespace. + pub fn root_node_handle(&self, _extra_key: Option<&str>) -> NodeId { + self.arena.node(self.arena.root()).id + } + + /// Return input indices in depth-first, subtree-weight order. + pub fn dfs_weight_order(&self, node_ids: &[NodeId]) -> Vec { + self.try_dfs_weight_order(node_ids) + .unwrap_or_else(|error| panic!("{error}")) + } + + pub fn try_dfs_weight_order( + &self, + node_ids: &[NodeId], + ) -> Result, TreeCoreRuntimeError> { + let mut node_to_indices: HashMap> = HashMap::new(); + for (index, &node_id) in node_ids.iter().enumerate() { + let node_id = self.try_resolve_node_handle_(node_id)?; + node_to_indices.entry(node_id).or_default().push(index); + } + + let mut node_to_weight: HashMap = HashMap::new(); + for (&node_id, indices) in &node_to_indices { + let mut cursor = node_id; + loop { + *node_to_weight.entry(cursor).or_default() += indices.len(); + let Some(parent) = self.arena.node(cursor).try_parent() else { + break; + }; + cursor = parent; + } + } + + let mut order = Vec::with_capacity(node_ids.len()); + let mut stack = vec![(self.arena.root(), false)]; + while let Some((node_id, emit)) = stack.pop() { + if emit { + if let Some(indices) = node_to_indices.get(&node_id) { + order.extend(indices); + } + continue; + } + + stack.push((node_id, true)); + let mut children: Vec = self + .arena + .node(node_id) + .children + .values() + .copied() + .filter(|child| node_to_weight.contains_key(child)) + .collect(); + children.sort_by(|left, right| { + node_to_weight[right] + .cmp(&node_to_weight[left]) + .then_with(|| { + self.arena + .node(*left) + .creation_counter + .cmp(&self.arena.node(*right).creation_counter) + }) + }); + stack.extend(children.into_iter().rev().map(|child| (child, false))); + } + Ok(order) + } + + /// Build the backup action for a node and its unbacked ancestors. + pub fn build_backup_kv_action_(&self, node: &Node, write_back: bool) -> BackupKV { + let mut chain = vec![node.id]; + if !write_back { + let mut ancestor = node.try_parent(); + while let Some(ancestor_idx) = ancestor { + let ancestor_node = self.arena.node(ancestor_idx); + if ancestor_node.is_root() || ancestor_node.backuped() { + break; + } + chain.push(ancestor_node.id); + ancestor = ancestor_node.try_parent(); + } + // write_through: Ancestors first to preserve backup invariant + chain.reverse(); + } + BackupKV { node_ids: chain } + } + + /// Commit each component's HiCache transfers onto the node. + pub fn commit_hicache_transfers( + &mut self, + node_id: NodeId, + phase: CacheTransferPhase, + comp_xfers: HashMap>, + cache_actions: &mut Vec, + mut insert_result: Option<&mut InsertResult>, + pool_storage_result: Option<&PoolTransferResult>, + ) { + let node_id = self.arena.resolve(node_id); + for (component_type, transfers) in comp_xfers { + self.component_by_type_(component_type) + .commit_hicache_transfer( + self, + node_id, + phase, + transfers, + cache_actions, + insert_result.as_deref_mut(), + pool_storage_result, + ); + } + } + + /// Commit a successful backup to the node. + pub fn commit_backup( + &mut self, + node_id: NodeId, + host_indices: Tensor, + comp_xfers: HashMap>, + ) { + let node_id = self.arena.resolve(node_id); + let mut cache_actions: Vec = Vec::new(); + if host_indices.numel() > 0 { + let kv_xfer = PoolTransfer { + name: PoolName::Kv, + host_indices: Some(host_indices), + ..Default::default() + }; + self.component_by_type_(BASE_COMPONENT_TYPE) + .commit_hicache_transfer( + self, + node_id, + CacheTransferPhase::BackupHost, + vec![kv_xfer], + &mut cache_actions, + /* insert_result = */ None, + /* pool_storage_result = */ None, + ); + } + for (component_type, transfers) in comp_xfers { + self.component_by_type_(component_type) + .commit_hicache_transfer( + self, + node_id, + CacheTransferPhase::BackupHost, + transfers, + &mut cache_actions, + /* insert_result = */ None, + /* pool_storage_result = */ None, + ); + } + assert!(cache_actions.is_empty()); // BACKUP_HOST emits no actions + self.update_full_coexisting_host_tracking_(node_id); + } + + /// Commit a successful H->D load-back onto the node; the SWA full->swa mapping + /// rebuild is deferred to the orchestration layer. + pub fn commit_load_back( + &mut self, + node_id: NodeId, + device_indices: Tensor, + mut kv_xfer: PoolTransfer, + comp_xfers: HashMap>, + ) -> Vec { + let anchor_id = node_id; + let node_id = self.arena.resolve(node_id); + let mut cache_actions: Vec = Vec::new(); + kv_xfer.device_indices = Some(device_indices); + let nodes_to_load = kv_xfer.nodes_to_load.clone(); + if self.is_write_back { + // Pin Full KV host slots against duplicate reclaim until the ack. + // Auxiliary pools have independent host locks and may legitimately + // load the same radix node under a different anchor. + for &pinned_id in nodes_to_load.iter().flatten() { + let pinned_idx = self.arena.resolve(pinned_id); + let pinned = self.arena.node_mut(pinned_idx); + assert!( + pinned.load_back_pending_id.is_none_or(|id| id == anchor_id), + "node {pinned_id} pinned by load-back {:?}, new anchor {anchor_id}", + pinned.load_back_pending_id + ); + pinned.load_back_pending_id = Some(anchor_id); + self.update_evictable_leaf_sets_(pinned_idx); + } + } + self.component_by_type_(BASE_COMPONENT_TYPE) + .commit_hicache_transfer( + self, + node_id, + CacheTransferPhase::LoadBack, + vec![kv_xfer], + &mut cache_actions, + /* insert_result = */ None, + /* pool_storage_result = */ None, + ); + for loaded_id in nodes_to_load.unwrap_or_default() { + let loaded_idx = self.arena.resolve(loaded_id); + self.record_store_event_(loaded_idx, StorageMedium::Gpu); + } + for (component_type, transfers) in comp_xfers { + self.component_by_type_(component_type) + .commit_hicache_transfer( + self, + node_id, + CacheTransferPhase::LoadBack, + transfers, + &mut cache_actions, + /* insert_result = */ None, + /* pool_storage_result = */ None, + ); + } + self.update_evictable_leaf_sets_(node_id); + cache_actions + } + + /// Finalize load-back state along the anchor's root path. + /// + /// Write-back clears matching Full KV source pins. Write-through has no + /// pins, but both policies refresh Full host/device duplicate tracking once + /// the device copies are visible. + pub fn finish_load_back(&mut self, anchor_node_id: NodeId) { + let mut node_id = Some(self.arena.resolve(anchor_node_id)); + while let Some(idx) = node_id { + if self.arena.node(idx).is_root() { + break; + } + if self.is_write_back { + if self.arena.node(idx).load_back_pending_id != Some(anchor_node_id) { + node_id = self.arena.node(idx).try_parent(); + continue; + } + self.arena.node_mut(idx).load_back_pending_id = None; + // The pin blocked leaf-set membership; re-evaluate it. + self.update_evictable_leaf_sets_(idx); + } + self.update_full_coexisting_host_tracking_(idx); + node_id = self.arena.node(idx).try_parent(); + } + } + + /// Mark a node as having an in-flight write-through backup. + pub fn mark_write_through_pending(&mut self, node_id: NodeId) { + let node_idx = self.arena.resolve(node_id); + self.arena.node_mut(node_idx).write_through_pending_id = Some(node_id); + } + + /// Clear the write-through-pending mark (when it matches ack_id) and record the + /// host store event for each acked node. + pub fn finish_write_through(&mut self, node_ids: Vec, ack_id: usize) { + for node_id in node_ids { + let node_idx = self.arena.resolve(node_id); + let node = self.arena.node_mut(node_idx); + if node.write_through_pending_id == Some(ack_id) { + node.write_through_pending_id = None; + self.update_full_coexisting_host_tracking_(node_idx); + } + self.record_store_event_(node_idx, StorageMedium::Cpu); + } + } + + /// Store an auxiliary component's device value onto a node and restamp + /// its LRU. + pub fn set_component_device_value( + &mut self, + node_id: NodeId, + component_type: ComponentType, + value: Tensor, + ) { + self.assert_component_enabled_(component_type); + let node_idx = self.arena.resolve(node_id); + self.set_component_device_value_(node_idx, component_type, value); + } + + /// Slot-keyed aux store (internal): set the device value and restamp the LRU. + pub(crate) fn set_component_device_value_( + &mut self, + node_id: NodeIdx_, + component_type: ComponentType, + value: Tensor, + ) { + assert!( + component_type != BASE_COMPONENT_TYPE, + "set_component_device_value: auxiliary components only" + ); + let tokens = value.size()[0] as usize; + self.arena.set_device_value(node_id, component_type, value); + let host_lru = self.host_lru_list_mut(component_type); + if host_lru.in_list(Some(node_id)) { + host_lru.remove_node(node_id); + } + self.device_lru_list_mut(component_type).insert_mru(node_id); + self.inc_evictable_size(component_type, tokens); + } + + /// The component's device value on the node, or None if evicted. + pub fn get_component_device_value( + &self, + node_id: NodeId, + component_type: ComponentType, + ) -> Option<&Tensor> { + self.assert_component_enabled_(component_type); + self.arena + .try_device_value(self.arena.resolve(node_id), component_type) + } + + /// Whether the component's data is device-evicted but host-backed. + pub fn component_has_host_value_only( + &self, + node_id: NodeId, + component_type: ComponentType, + ) -> bool { + self.assert_component_enabled_(component_type); + let node_idx = self.arena.resolve(node_id); + !self.arena.has_device_value(node_idx, component_type) + && self.arena.has_host_value(node_idx, component_type) + } + + /// Verify tree-structure, leaf-set, LRU, size, and ongoing-op invariants; raise + /// AssertionError on any violation. ongoing_* args are (id, node_id) pairs. + pub fn sanity_check( + &self, + ongoing_write_through: &[(i64, NodeId)], + ongoing_load_back: &[(i64, NodeId)], + ) { + if let Err(message) = self.try_sanity_check(ongoing_write_through, ongoing_load_back) { + self.pretty_print(); + panic!("{message}"); + } + } + + /// Fallible variant of [`Self::sanity_check`] for language bindings. + pub fn try_sanity_check( + &self, + ongoing_write_through: &[(i64, NodeId)], + ongoing_load_back: &[(i64, NodeId)], + ) -> Result<(), String> { + let mut errors: Vec = Vec::new(); + let all_nodes = self.collect_all_nodes_(); + let all_node_set: HashSet = all_nodes.iter().copied().collect(); + + // ── PART 1: Tree Structure ── + // The single root: value-less, protected, parent-less, no node-level edges. + let root_idx = self.arena.root(); + let root = self.arena.node(root_idx); + for i in 0..self.components.len() { + let ct = self.components[i].component_type(); + if root.values[ct.idx()].value.is_some() { + errors.push(format!( + "[Root] root {root_idx} holds a {ct:?} device value" + )); + } + if root.has_host_value(ct) { + errors.push(format!("[Root] root {root_idx} holds a {ct:?} host value")); + } + if root.values[ct.idx()].lock_ref == 0 { + errors.push(format!("[Root] root {root_idx} {ct:?} lock_ref=0")); + } + } + if root.try_parent().is_some() { + errors.push(format!("[Root] root {root_idx} has a parent pointer")); + } + // Leaf sets aside, every live arena slot must be tree-reachable. + let orphans: Vec = self + .arena + .live_ids() + .filter(|id| !all_node_set.contains(id)) + .collect(); + if !orphans.is_empty() { + errors.push(format!( + "[Tree] {} orphaned live nodes: {:?}", + orphans.len(), + &orphans[..orphans.len().min(5)] + )); + } + for (&node_handle, hashes) in &self.salted_event_hashes { + let Some(node_id) = self.arena.try_resolve(node_handle) else { + errors.push(format!( + "[Events] salted hashes reference freed node {node_handle}" + )); + continue; + }; + let node = self.arena.node(node_id); + if node.namespace.cache_salt().is_none() { + errors.push(format!( + "[Events] unsalted node {node_handle} carries salted hashes" + )); + } + let expected_pages = node.key.atom_len().div_ceil(self.page_size); + if hashes.len() != expected_pages { + errors.push(format!( + "[Events] node {node_handle} has {} salted hashes for {expected_pages} pages", + hashes.len() + )); + } + } + // Parent ↔ child bidirectional consistency + for &node_id in &all_nodes { + for ((edge_namespace, edge_key), &child_id) in &self.arena.node(node_id).children { + let child = self.arena.node(child_id); + let child_parent = child.try_parent(); + if child_parent != Some(node_id) { + errors.push(format!( + "[Tree] child {child_id} parent={child_parent:?}, expected {node_id}" + )); + } + if child.key.atom_len() == 0 { + errors.push(format!("[Tree] node {child_id} has an empty key")); + continue; + } + if !child.key.atom_len().is_multiple_of(self.page_size) { + errors.push(format!("[Tree] node {child_id} key is not page-aligned")); + continue; + } + // The edge key must be the child's own namespaced child key. + if *edge_key != child.key.child_key(self.page_size) { + errors.push(format!( + "[Tree] child {child_id} not mapped under its own child key" + )); + } + if *edge_namespace != child.namespace { + errors.push(format!( + "[Tree] child {child_id} namespace {:?} filed under {edge_namespace:?}", + child.namespace + )); + } + // Namespaces partition at the root; below it children inherit. + if !self.arena.node(node_id).is_root() + && child.namespace != self.arena.node(node_id).namespace + { + errors.push(format!( + "[Tree] child {child_id} namespace differs from its parent's" + )); + } + if let Some(value) = child.try_device_value(FULL) + && value.size()[0] as usize != child.key.atom_len() + { + errors.push(format!( + "[Tree] node {child_id} Full value length {} != key length {}", + value.size()[0], + child.key.atom_len() + )); + } + if let Some(value) = child.try_host_value(FULL) + && value.size()[0] as usize != child.key.atom_len() + { + errors.push(format!( + "[Tree] node {child_id} Full host value length {} != key length {}", + value.size()[0], + child.key.atom_len() + )); + } + } + } + + // ── PART 2: Per-node state machine and leaf qualification ── + let mut expected_dev_leaves: HashSet = HashSet::new(); + let mut expected_hst_leaves: HashSet = HashSet::new(); + let mut expected_full_coexisting_host_nodes: HashSet = HashSet::new(); + + for &node_id in &all_nodes { + if self.arena.node(node_id).is_root() { + continue; + } + let node = self.arena.node(node_id); + let full_dev = node.has_device_value(FULL); + let full_hst = node.has_host_value(FULL); + + // Full is the tree backbone, so aux data requires Full data. + for i in 0..self.components.len() { + let ct = self.components[i].component_type(); + if ct == BASE_COMPONENT_TYPE { + continue; + } + if node.values[ct.idx()].value.is_some() && !full_dev { + errors.push(format!( + "node {node_id} {ct:?} device present but Full.value=None" + )); + } + // Auxiliary host data may outlive Full host data under write-back. + if node.has_host_value(ct) && !full_hst && !(self.is_write_back && full_dev) { + errors.push(format!( + "node {node_id} {ct:?} host present but Full.host_value=None" + )); + } + } + + // Every node must keep Full data on at least one layer. + if !full_dev && !full_hst { + errors.push(format!( + "node {node_id} dead: no Full device and no Full host" + )); + } + + // Parent prefixes must keep data whenever the child does. + let parent_id = node.parent(); + if !self.arena.node(parent_id).is_root() { + let parent = self.arena.node(parent_id); + if full_dev && !parent.has_device_value(FULL) { + errors.push(format!( + "node {node_id} device present but parent {parent_id} evicted" + )); + } + if full_hst && !parent.has_host_value(FULL) && !self.is_write_back { + errors.push(format!( + "node {node_id} backed up but parent {parent_id} not backed up" + )); + } + } + + // Lock hierarchy must stay sane (the u32 counters cannot go negative). + let full_lock = node.device_lock_ref(FULL); + for i in 0..self.components.len() { + let ct = self.components[i].component_type(); + let device_state = &node.values[ct.idx()]; + if ct != BASE_COMPONENT_TYPE && full_lock < device_state.lock_ref { + errors.push(format!( + "node {node_id} full_lock={full_lock} < {ct:?}_lock={}", + device_state.lock_ref + )); + } + if device_state.value.is_none() && device_state.lock_ref > 0 { + errors.push(format!( + "node {node_id} {ct:?} evicted but lock_ref={}", + device_state.lock_ref + )); + } + } + + // Collect expected leaf qualification (single pass) + if self.is_evictable_device_leaf_(node) { + expected_dev_leaves.insert(node_id); + } + if self.is_evictable_host_leaf_(node) { + expected_hst_leaves.insert(node_id); + } + if self.is_settled_full_coexisting_host_node_(node) { + expected_full_coexisting_host_nodes.insert(node_id); + } + } + + // ── PART 3: Tracking structures ── + + // Device leaf set must match the expected leaves. + let device_leaves: HashSet = self.evictable_device_leaves.iter().collect(); + if device_leaves != expected_dev_leaves { + let extra: Vec = device_leaves + .difference(&expected_dev_leaves) + .copied() + .take(5) + .collect(); + let missing: Vec = expected_dev_leaves + .difference(&device_leaves) + .copied() + .take(5) + .collect(); + if !extra.is_empty() { + errors.push(format!("D-leaf extra: {extra:?}")); + } + if !missing.is_empty() { + errors.push(format!("D-leaf missing: {missing:?}")); + } + } + + // Host leaf set must match the expected leaves. + let host_leaves: HashSet = self.evictable_host_leaves.iter().collect(); + if host_leaves != expected_hst_leaves { + let extra: Vec = host_leaves + .difference(&expected_hst_leaves) + .copied() + .take(5) + .collect(); + let missing: Vec = expected_hst_leaves + .difference(&host_leaves) + .copied() + .take(5) + .collect(); + if !extra.is_empty() { + errors.push(format!("H-leaf extra: {extra:?}")); + } + if !missing.is_empty() { + errors.push(format!("H-leaf missing: {missing:?}")); + } + } + + // Lazy tracking permits stale entries, but not missing or recycled ones. + let full_coexisting_host_nodes: HashSet = + self.full_coexisting_host_nodes.iter().collect(); + let missing: Vec = expected_full_coexisting_host_nodes + .difference(&full_coexisting_host_nodes) + .copied() + .take(5) + .collect(); + if !missing.is_empty() { + errors.push(format!("Full host coexistence missing: {missing:?}")); + } + let ghosts: Vec = full_coexisting_host_nodes + .difference(&all_node_set) + .copied() + .take(5) + .collect(); + if !ghosts.is_empty() { + errors.push(format!("Full host coexistence ghosts: {ghosts:?}")); + } + + // D-leaf ∩ H-leaf = ∅ + let overlap: Vec = device_leaves.intersection(&host_leaves).copied().collect(); + if !overlap.is_empty() { + errors.push(format!( + "[Leaf] {} in both sets: {:?}", + overlap.len(), + &overlap[..overlap.len().min(5)] + )); + } + + // Stale nodes: leaf sets must only contain tree-reachable nodes + let stale: Vec = device_leaves.difference(&all_node_set).copied().collect(); + if !stale.is_empty() { + errors.push(format!( + "{} stale nodes in device_leaves: {:?}", + stale.len(), + &stale[..stale.len().min(5)] + )); + } + let stale: Vec = host_leaves.difference(&all_node_set).copied().collect(); + if !stale.is_empty() { + errors.push(format!( + "{} stale nodes in host_leaves: {:?}", + stale.len(), + &stale[..stale.len().min(5)] + )); + } + + // Per-component LRU tracking + for i in 0..self.components.len() { + let ct = self.components[i].component_type(); + let lru = self.device_lru_list(ct); + let host_lru = self.host_lru_list(ct); + if ct == BASE_COMPONENT_TYPE { + // Full uses leaf sets, not LRU + if lru.len() > 0 { + errors.push(format!("Full device LRU not empty: {}", lru.len())); + } + if host_lru.len() > 0 { + errors.push(format!("Full host LRU not empty: {}", host_lru.len())); + } + // Linked-list integrity + lru.check_linked_list_(&format!("[device][{ct:?}]"), &mut errors); + host_lru.check_linked_list_(&format!("[host][{ct:?}]"), &mut errors); + } else { + // Aux device values must match the device LRU; aux host-only + // states must match the host LRU; never both at once. + let mut device_count = 0; + let mut host_only_count = 0; + for &node_id in &all_nodes { + if self.arena.node(node_id).is_root() { + continue; + } + let node = self.arena.node(node_id); + let has_device = node.values[ct.idx()].value.is_some(); + if has_device != lru.in_list(Some(node_id)) { + errors.push(format!( + "{ct:?} device LRU mismatch at node {node_id}: value={has_device} in_lru={}", + lru.in_list(Some(node_id)) + )); + } + let host_only = !has_device && node.has_host_value(ct); + if host_only != host_lru.in_list(Some(node_id)) { + errors.push(format!( + "{ct:?} host LRU mismatch at node {node_id}: host_only={host_only} in_lru={}", + host_lru.in_list(Some(node_id)) + )); + } + if lru.in_list(Some(node_id)) && host_lru.in_list(Some(node_id)) { + errors.push(format!("{ct:?} node {node_id} in both device and host LRU")); + } + device_count += has_device as usize; + host_only_count += host_only as usize; + } + if device_count != lru.len() { + errors.push(format!( + "{ct:?} device LRU: tree={device_count} != lru={}", + lru.len() + )); + } + if host_only_count != host_lru.len() { + errors.push(format!( + "{ct:?} host LRU: tree={host_only_count} != lru={}", + host_lru.len() + )); + } + // Linked-list integrity + lru.check_linked_list_(&format!("[device][{ct:?}]"), &mut errors); + host_lru.check_linked_list_(&format!("[host][{ct:?}]"), &mut errors); + } + } + + // ── PART 4: Size Accounting ── + for i in 0..self.components.len() { + let ct = self.components[i].component_type(); + let mut evictable = 0usize; + let mut protected = 0usize; + for &node_id in &all_nodes { + if self.arena.node(node_id).is_root() { + continue; + } + let state = &self.arena.node(node_id).values[ct.idx()]; + if let Some(value) = &state.value { + let tokens = value.size()[0] as usize; + if state.lock_ref > 0 { + protected += tokens; + } else { + evictable += tokens; + } + } + } + let recorded = self.component_state(ct); + if recorded.evictable_size != evictable { + errors.push(format!( + "[Size] {ct:?} evictable={} != recomputed={evictable}", + recorded.evictable_size + )); + } + if recorded.protected_size != protected { + errors.push(format!( + "[Size] {ct:?} protected={} != recomputed={protected}", + recorded.protected_size + )); + } + } + + // ── PART 5: Ongoing Operations ── + for &(op_id, node_id) in ongoing_write_through { + match self.arena.try_resolve(node_id) { + None => { + errors.push(format!("[Ongoing] write_through node {op_id} not in tree")); + } + Some(idx) if self.arena.device_lock_ref(idx, FULL) == 0 => { + errors.push(format!("[Ongoing] write_through node {op_id} lock_ref=0")); + } + Some(_) => {} + } + } + for &(op_id, node_id) in ongoing_load_back { + match self.arena.try_resolve(node_id) { + None => { + errors.push(format!("[Ongoing] load_back node {op_id} not in tree")); + } + Some(idx) if self.arena.device_lock_ref(idx, FULL) == 0 => { + errors.push(format!("[Ongoing] load_back node {op_id} lock_ref=0")); + } + Some(_) => {} + } + } + // Reject load-back pins that would survive their operation. + let ongoing_load_ids: HashSet = + ongoing_load_back.iter().map(|&(_, id)| id).collect(); + for &node_id in &all_nodes { + let pending = self.arena.node(node_id).load_back_pending_id; + if let Some(anchor) = pending + && !ongoing_load_ids.contains(&anchor) + { + errors.push(format!( + "[Ongoing] node {node_id} load_back_pending_id={anchor} \ + has no live load-back" + )); + } + } + + if !errors.is_empty() { + return Err(format!( + "Sanity check FAILED ({} violations across {} nodes):\n{}", + errors.len(), + all_nodes.len(), + errors + .iter() + .map(|e| format!(" {e}")) + .collect::>() + .join("\n") + )); + } + Ok(()) + } + + /// Every live node in the tree. + pub fn collect_all_nodes_(&self) -> Vec { + let mut nodes: Vec = Vec::new(); + // The visited guard keeps a corrupted (cyclic) tree from hanging the walk. + let mut visited: HashSet = HashSet::new(); + let mut stack: Vec = vec![self.arena.root()]; + while let Some(node_id) = stack.pop() { + if !visited.insert(node_id) { + continue; + } + nodes.push(node_id); + stack.extend(self.arena.node(node_id).children.values().copied()); + } + nodes + } +} + +#[cfg(any(test, feature = "inspection"))] +impl UnifiedTreeCore { + // Test-only inspection support for the backend-neutral Python suite. + + /// Whether the external node handle is currently live. + pub fn inspect_contains_node(&self, node_id: NodeId) -> bool { + self.arena.try_resolve(node_id).is_some() + } + + /// The parent node's external handle, or None for the root. + pub fn inspect_get_parent_node_id(&self, node_id: NodeId) -> Option { + let node_id = self.arena.resolve(node_id); + self.arena + .node(node_id) + .try_parent() + .map(|parent_id| self.arena.node(parent_id).id) + } + + /// A materialized snapshot of the node's child handles. + pub fn inspect_get_child_node_ids(&self, node_id: NodeId) -> Vec { + let node_id = self.arena.resolve(node_id); + self.arena + .node(node_id) + .children + .values() + .map(|&child_id| self.arena.node(child_id).id) + .collect() + } + + /// Logical radix-key length in key atoms. + pub fn inspect_get_node_key_length(&self, node_id: NodeId) -> usize { + self.arena.node(self.arena.resolve(node_id)).key.atom_len() + } + + /// Materialized raw token ids spanned by the node key. + pub fn inspect_get_node_token_ids(&self, node_id: NodeId) -> Vec { + K::raw_token_ids(self.arena.node(self.arena.resolve(node_id)).key.as_ref()).into_owned() + } + + /// Whether this core's key representation uses overlapping bigrams. + pub fn inspect_is_node_key_bigram(&self, node_id: NodeId) -> bool { + let node = self.arena.node(self.arena.resolve(node_id)); + !node.is_root() && K::IS_BIGRAM + } + + /// A shallow tensor snapshot of a component's host value. + pub fn inspect_get_component_host_value( + &self, + node_id: NodeId, + component_type: ComponentType, + ) -> Option { + self.assert_component_enabled_(component_type); + self.arena + .node(self.arena.resolve(node_id)) + .try_host_value(component_type) + .map(Tensor::shallow_clone) + } + + /// A component's device lock count on a node. + pub fn inspect_get_component_device_lock_ref( + &self, + node_id: NodeId, + component_type: ComponentType, + ) -> u32 { + self.assert_component_enabled_(component_type); + self.arena + .node(self.arena.resolve(node_id)) + .device_lock_ref(component_type) + } + + /// A node's accumulated match count. + pub fn inspect_get_node_hit_count(&self, node_id: NodeId) -> i64 { + self.arena.node(self.arena.resolve(node_id)).hit_count + } + + /// A node's in-flight write-through acknowledgement id. + pub fn inspect_get_write_through_pending_id(&self, node_id: NodeId) -> Option { + self.arena + .node(self.arena.resolve(node_id)) + .write_through_pending_id + } + + /// Whether a node is in a component's device LRU. + pub fn inspect_is_node_in_device_lru( + &self, + node_id: NodeId, + component_type: ComponentType, + ) -> bool { + if self.try_component_by_type_(component_type).is_none() { + return false; + } + self.device_lru_list(component_type) + .in_list(Some(self.arena.resolve(node_id))) + } + + /// Whether a node is in a component's host LRU. + pub fn inspect_is_node_in_host_lru( + &self, + node_id: NodeId, + component_type: ComponentType, + ) -> bool { + if self.try_component_by_type_(component_type).is_none() { + return false; + } + self.host_lru_list(component_type) + .in_list(Some(self.arena.resolve(node_id))) + } + + /// Materialize a component's device LRU from most to least recent. + pub fn inspect_get_component_device_lru_node_ids( + &self, + component_type: ComponentType, + ) -> Vec { + if self.try_component_by_type_(component_type).is_none() { + return Vec::new(); + } + self.device_lru_list(component_type) + .snapshot_node_ids() + .into_iter() + .map(|node_id| self.arena.node(node_id).id) + .collect() + } + + /// Whether a live node belongs to the device-evictable leaf set. + pub fn inspect_is_device_evictable_leaf(&self, node_id: NodeId) -> bool { + self.arena + .try_resolve(node_id) + .is_some_and(|node_id| self.evictable_device_leaves.contains(node_id)) + } + + /// Whether a live node belongs to the host-evictable leaf set. + pub fn inspect_is_host_evictable_leaf(&self, node_id: NodeId) -> bool { + self.arena + .try_resolve(node_id) + .is_some_and(|node_id| self.evictable_host_leaves.contains(node_id)) + } + + /// Whether the node is currently eligible as a Full device leaf. + pub fn inspect_is_device_leaf(&self, node_id: NodeId) -> bool { + let node_id = self.arena.resolve(node_id); + let node = self.arena.node(node_id); + if node.is_root() || node.evicted() || node.is_device_locked() { + return false; + } + !node + .children + .values() + .any(|&child_id| self.arena.has_device_value(child_id, FULL)) + } + + /// Materialize every live tree node handle. + pub fn inspect_get_all_node_ids(&self) -> Vec { + self.collect_all_nodes_() + .into_iter() + .map(|node_id| self.arena.node(node_id).id) + .collect() + } + + /// Protected token count for one component. + pub fn inspect_component_protected_size(&self, component_type: ComponentType) -> usize { + self.component_protected_size(component_type) + } + + /// Replace a node's hash chain without updating other bookkeeping. + pub fn inspect_set_node_hash_values( + &mut self, + node_id: NodeId, + hash_values: Option>, + ) { + let node_id = self.arena.resolve(node_id); + self.arena.node_mut(node_id).hash_value = hash_values; + } + + /// Replace a component's device value without updating bookkeeping. + pub fn inspect_set_component_device_value_raw( + &mut self, + node_id: NodeId, + component_type: ComponentType, + value: Option, + ) { + self.assert_component_enabled_(component_type); + let node_id = self.arena.resolve(node_id); + self.arena + .node_mut(node_id) + .state_mut_(ValueSlotIdx::device(component_type)) + .value = value; + } + + /// Replace a component's host value without updating bookkeeping. + pub fn inspect_set_component_host_value_raw( + &mut self, + node_id: NodeId, + component_type: ComponentType, + value: Option, + ) { + self.assert_component_enabled_(component_type); + let node_id = self.arena.resolve(node_id); + self.arena + .node_mut(node_id) + .state_mut_(ValueSlotIdx::host(component_type)) + .value = value; + } + + /// Replace a component's device lock count without updating bookkeeping. + pub fn inspect_set_component_device_lock_ref( + &mut self, + node_id: NodeId, + component_type: ComponentType, + lock_ref: u32, + ) { + self.assert_component_enabled_(component_type); + let node_id = self.arena.resolve(node_id); + self.arena + .node_mut(node_id) + .set_lock_ref_(ValueSlotIdx::device(component_type), lock_ref); + } + + /// Remove a node from a component's device LRU. + pub fn inspect_remove_node_from_device_lru( + &mut self, + node_id: NodeId, + component_type: ComponentType, + ) { + self.assert_component_enabled_(component_type); + let node_id = self.arena.resolve(node_id); + self.device_lru_list_mut(component_type) + .remove_node(node_id); + } + + /// Insert a node as a component's most-recent host-LRU entry. + pub fn inspect_insert_node_into_host_lru( + &mut self, + node_id: NodeId, + component_type: ComponentType, + ) { + self.assert_component_enabled_(component_type); + let node_id = self.arena.resolve(node_id); + self.host_lru_list_mut(component_type).insert_mru(node_id); + } + + /// Replace a component's evictable-device token count. + pub fn inspect_set_component_evictable_size( + &mut self, + component_type: ComponentType, + value: usize, + ) { + self.assert_component_enabled_(component_type); + self.component_state_mut(component_type).evictable_size = value; + } + + /// Replace a component's protected-device token count. + pub fn inspect_set_component_protected_size( + &mut self, + component_type: ComponentType, + value: usize, + ) { + self.assert_component_enabled_(component_type); + self.component_state_mut(component_type).protected_size = value; + } + + /// Refresh Full device/host duplicate tracking for a node. + pub fn inspect_update_duplicate_tracking(&mut self, node_id: NodeId) { + let node_id = self.arena.resolve(node_id); + self.update_full_coexisting_host_tracking_(node_id); + } + + /// Advance one suspended insert walk step without flushing its pending actions. + pub fn inspect_advance_insert_walk_once(&mut self) -> Result<(), &'static str> { + let Some(mut state) = self.ongoing_insert_walk_state.take() else { + return Err("no in-flight insert"); + }; + if !matches!(state.phase, InsertPhase::Walk) { + self.ongoing_insert_walk_state = Some(state); + return Err("in-flight insert is not in walk phase"); + } + self.insert_walk_step_(&mut state); + self.ongoing_insert_walk_state = Some(state); + Ok(()) + } + + /// Evict one component layer and detach the corresponding LRU entry. + pub fn inspect_evict_component( + &mut self, + node_id: NodeId, + component_type: ComponentType, + target: EvictLayer, + ) -> EvictionStepResult { + self.assert_component_enabled_(component_type); + let node_id = self.arena.resolve(node_id); + let mut result = EvictionStepResult::default(); + self.evict_component_and_detach_lru_( + node_id, + component_type, + &mut result.device_frees, + &mut result.host_frees, + target, + Some(&mut result.tracker), + ); + result + } + + /// Validate component locks for a cascade without mutating the tree. + pub fn inspect_validate_cascade_evict( + &self, + node_id: NodeId, + trigger_component_type: ComponentType, + target: EvictLayer, + ) -> Result<(), String> { + self.assert_component_enabled_(trigger_component_type); + let node_id = self.arena.resolve(node_id); + let is_leaf = match target { + EvictLayer::Device => self.evictable_device_leaves.contains(node_id), + EvictLayer::Host => self.evictable_host_leaves.contains(node_id), + EvictLayer::All => false, + }; + let trigger_component = self.component_by_type_(trigger_component_type); + let trigger_priority = trigger_component.eviction_priority(is_leaf); + let trigger_internal_priority = + trigger_component.eviction_priority(/* is_leaf = */ false); + for component in &self.components { + self.should_cascade_evict_component_( + node_id, + trigger_component_type, + component.as_ref(), + target, + is_leaf, + trigger_priority, + trigger_internal_priority, + )?; + } + Ok(()) + } + + /// Delete childless tombstone ancestors starting at `node_id`. + pub fn inspect_cleanup_tombstone_ancestors(&mut self, node_id: NodeId) -> EvictionStepResult { + let node_id = self.arena.resolve(node_id); + let mut result = EvictionStepResult::default(); + self.iteratively_delete_tombstone_leaf_( + node_id, + &mut result.tracker, + &mut result.device_frees, + &mut result.host_frees, + ); + result + } + + /// Run one component's real match-result finalizer. + pub fn inspect_finalize_component_match_result( + &self, + component_type: ComponentType, + result: MatchResult, + params: &MatchPrefixParams<'_, K>, + value_chunks: &[Tensor], + best_value_len: usize, + ) -> MatchResult { + self.component_by_type_(component_type) + .finalize_match_result_in_tree_core(self, result, params, value_chunks, best_value_len) + } + + /// Build the ordered device-to-host backup node list. + pub fn inspect_build_backup_node_ids(&self, node_id: NodeId, write_back: bool) -> Vec { + let node_id = self.arena.resolve(node_id); + self.build_backup_kv_action_(self.arena.node(node_id), write_back) + .node_ids + } +} + +impl UnifiedTreeCore { + /// Print the tree structure for debugging. + pub fn pretty_print(&self) { + println!("{}", self.pretty_format_()); + } + + /// The pretty_print rendering: one indented + /// `[id] key_len full_lock component=yes/no` line per node. + fn pretty_format_(&self) -> String { + let mut lines: Vec = Vec::new(); + let mut visited: HashSet = HashSet::new(); + let mut stack: Vec<(NodeIdx_, usize)> = vec![(self.arena.root(), 0)]; + while let Some((node_id, indent)) = stack.pop() { + if !visited.insert(node_id) { + continue; + } + let node = self.arena.node(node_id); + let component_str = self + .components + .iter() + .map(|component| { + let ct = component.component_type(); + let state = if node.values[ct.idx()].value.is_some() { + "yes" + } else { + "no" + }; + format!("{ct:?}={state}") + }) + .collect::>() + .join(" "); + lines.push(format!( + "{} [{}] {} full_lock={} {}", + " ".repeat(indent), + node.id, + node.key.atom_len(), + node.device_lock_ref(FULL), + component_str + )); + stack.extend(node.children.values().map(|&child| (child, indent + 2))); + } + lines.join("\n") + } + + /// Evictable token count of the FULL (base) component. + pub fn evictable_size(&self) -> usize { + self.evictable_size_(FULL) + } + + /// Protected (locked) token count of the FULL (base) component. + pub fn protected_size(&self) -> usize { + self.protected_size_(FULL) + } + + /// Evictable token count for one component (0 if the component is absent). + pub fn component_evictable_size(&self, component_type: ComponentType) -> usize { + self.try_component_by_type_(component_type) + .map_or(0, |_| self.evictable_size_(component_type)) + } + + /// Protected token count for one component (0 if the component is absent). + pub fn component_protected_size(&self, component_type: ComponentType) -> usize { + self.try_component_by_type_(component_type) + .map_or(0, |_| self.protected_size_(component_type)) + } + + /// FULL component evictable token count. + pub fn full_evictable_size(&self) -> usize { + self.evictable_size() + } + + /// FULL component protected token count. + pub fn full_protected_size(&self) -> usize { + self.protected_size() + } + + /// SWA component evictable token count. + pub fn swa_evictable_size(&self) -> usize { + self.evictable_size_(SWA) + } + + /// Mamba component evictable token count. + pub fn mamba_evictable_size(&self) -> usize { + self.evictable_size_(MAMBA) + } + + /// SWA component protected token count. + pub fn swa_protected_size(&self) -> usize { + self.protected_size_(SWA) + } + + /// Mamba component protected token count. + pub fn mamba_protected_size(&self) -> usize { + self.protected_size_(MAMBA) + } + + /// (full_tokens, aux_tokens) summed across the whole tree. + pub fn total_size(&self) -> (usize, usize) { + let mut total_size = 0; + let mut total_aux_size = 0; + let mut stack: Vec = vec![self.arena.root()]; + while let Some(node_id) = stack.pop() { + let node = self.arena.node(node_id); + total_size += node.device_value_len(FULL); + for i in 0..self.components.len() { + let ct = self.components[i].component_type(); + if ct == BASE_COMPONENT_TYPE { + continue; + } + if let Some(value) = &node.values[ct.idx()].value { + total_aux_size += value.size()[0] as usize; + } + } + stack.extend(self.arena.node(node_id).children.values().copied()); + } + (total_size, total_aux_size) + } + + /// Every FULL device value in the tree, concatenated. + pub fn all_values_flatten(&self) -> Tensor { + components::all_values_flatten(self, FULL) + } + + /// Flatten every FULL device slot into (slot, position, prev-slot) rows for the KV-canary sweep. + pub fn walk_for_kv_canary( + &self, + unlocked_only: bool, + swa_resident_only: bool, + ) -> KvCanaryWalkResult { + let swa_filter = swa_resident_only && self.components_by_type[SWA.idx()].is_some(); + let mut slot_indices: Vec = Vec::new(); + let mut positions: Vec = Vec::new(); + let mut prev_slot_indices: Vec = Vec::new(); + // (node, is_root, atom depth from root, last device slot on the path above) + let mut stack: Vec<(NodeIdx_, bool, i64, i64)> = vec![(self.arena.root(), true, 0, -1)]; + while let Some((node_id, is_root, depth, parent_last_slot)) = stack.pop() { + let node = self.arena.node(node_id); + let node_slots: Vec = node + .try_device_value(FULL) + .map(|value| { + Vec::::try_from(&value.to(Device::Cpu)) + .expect("device values are 1-D int64 tensors") + }) + .unwrap_or_default(); + + let mut emit = !is_root; + if unlocked_only { + // Unified SWA owns an independent component lock. A node can still + // hold Full KV for a running request while its SWA slots are unused. + emit = emit + && if swa_filter { + node.device_lock_ref(SWA) == 0 + } else { + node.device_lock_ref(FULL) == 0 + }; + } + if swa_filter { + emit = emit && node.has_device_value(SWA); + } + + // Skipped nodes still advance the chain/depth so descendants stay consistent. + let mut chain_last_slot = parent_last_slot; + for (j, &slot) in node_slots.iter().enumerate() { + if emit { + slot_indices.push(slot); + positions.push(depth + j as i64); + prev_slot_indices.push(if j == 0 { + parent_last_slot + } else { + node_slots[j - 1] + }); + } + chain_last_slot = slot; + } + + // Device-evicted nodes hold no slots but still span their key length. + let child_depth = depth + node.key.atom_len() as i64; + for &child_id in node.children.values() { + stack.push((child_id, false, child_depth, chain_last_slot)); + } + } + KvCanaryWalkResult { + slot_indices, + positions, + prev_slot_indices, + } + } + + /// Every Mamba device value in the tree, concatenated. + pub fn all_mamba_values_flatten(&self) -> Tensor { + components::all_values_flatten(self, MAMBA) + } +} + +// KV cache placement events. + +/// Storage tier of a stored/removed block. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub enum StorageMedium { + Gpu, + Cpu, +} + +impl StorageMedium { + /// The python StorageMedium enum value. + pub fn as_str(self) -> &'static str { + match self { + StorageMedium::Gpu => "GPU", + StorageMedium::Cpu => "CPU_PINNED", + } + } +} + +/// A KV placement event; one stored event may carry multiple same-sized pages. +#[derive(Clone, PartialEq, Eq, Debug)] +pub enum KvCacheEvent { + BlockStored { + block_hashes: Vec, + parent_block_hash: Option, + token_ids: Vec, + block_size: usize, + medium: StorageMedium, + cache_salt: Option>, + }, + BlockRemoved { + block_hashes: Vec, + medium: StorageMedium, + }, + AllBlocksCleared, +} +#[cfg(test)] +#[path = "tests/unified_tree_core.rs"] +mod tests; diff --git a/rust/mem-cache/torch_2_13_compat.h b/rust/mem-cache/torch_2_13_compat.h new file mode 100644 index 000000000..f70196ec8 --- /dev/null +++ b/rust/mem-cache/torch_2_13_compat.h @@ -0,0 +1,15 @@ +#pragma once + +#include +#include + +#if TORCH_VERSION_MAJOR > 2 || \ + (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR >= 13) +// Keep tch 0.24's removed alignment wrappers as explicit runtime errors. +#define align_as(...) \ + alias(); \ + throw std::runtime_error("align_as is unavailable in PyTorch 2.13+") +#define align_tensors(...) \ + autograd::variable_list{}; \ + throw std::runtime_error("align_tensors is unavailable in PyTorch 2.13+") +#endif diff --git a/scripts/ci/cuda/ci_install_dependency.sh b/scripts/ci/cuda/ci_install_dependency.sh index fa5ffafa2..c14275e0d 100755 --- a/scripts/ci/cuda/ci_install_dependency.sh +++ b/scripts/ci/cuda/ci_install_dependency.sh @@ -277,11 +277,24 @@ clean_site_packages() { } setup_cargo_cache() { + if [ "${SGLANG_BUILD_RUST_EXTS:-}" = "none" ]; then + echo "Using prebuilt Rust extensions; skipping Cargo target setup" + mark_step_done "${FUNCNAME[0]}" + return + fi + # actions/checkout's `git clean -ffdx` deletes the gitignored in-repo # rust/target, so every job recompiles the whole dependency graph. Move the # target dir out of the tree: setuptools-rust has no target-dir option of its # own and defers to CARGO_TARGET_DIR, which uv passes to the build backend. export CARGO_TARGET_DIR="${HOME}/.cache/sglang-cargo-target" + local cargo_target_lock="${HOME}/.cache/sglang-cargo-target.lock" + mkdir -p "${HOME}/.cache" + exec 9>"${cargo_target_lock}" + echo "Waiting for exclusive cargo target lock: ${cargo_target_lock}" + flock --exclusive 9 + CARGO_TARGET_LOCK_HELD=1 + echo "Acquired cargo target lock" mkdir -p "${CARGO_TARGET_DIR}" # Same disk-pressure guard as the uv cache in ci_cleanup_venv.sh (which @@ -298,6 +311,15 @@ setup_cargo_cache() { mark_step_done "${FUNCNAME[0]}" } +release_cargo_cache_lock() { + if [ "${CARGO_TARGET_LOCK_HELD:-0}" = "1" ]; then + flock --unlock 9 + exec 9>&- + CARGO_TARGET_LOCK_HELD=0 + echo "Released cargo target lock" + fi +} + setup_pip_toolchain() { if [ "$USE_VENV" = "1" ]; then # The bootstrap upgrade hit system pip; this upgrades the venv's own. @@ -473,10 +495,19 @@ require_prebuilt_rust_exts() { for module in server grpc multimodal; do [ -f "python/sglang/srt/rust_extensions/_${module}${suffix}" ] || missing+=("${module}") done + [ -f "python/sglang/srt/mem_cache/rust_tree_core/mem_cache${suffix}" ] \ + || missing+=("mem_cache") + [ -f "python/sglang/srt/mem_cache/rust_tree_core/mem_cache_inspection${suffix}" ] \ + || missing+=("mem_cache_inspection") if [ ${#missing[@]} -gt 0 ]; then echo "::warning::no prebuilt Rust extension ${suffix} for: ${missing[*]}; building from source" ls -l python/sglang/srt/rust_extensions/_*.so 2>/dev/null || echo "(no extension modules at all)" + ls -l python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so 2>/dev/null || true export SGLANG_BUILD_RUST_EXTS= + export SGLANG_RUST_BUILD_MODE=auto + if [ -n "${GITHUB_ENV:-}" ]; then + echo "SGLANG_RUST_BUILD_MODE=auto" >> "${GITHUB_ENV}" + fi mark_step_done "${FUNCNAME[0]}" return fi @@ -846,14 +877,15 @@ main() { install_apt_packages install_gdrcopy clean_site_packages - setup_cargo_cache require_prebuilt_rust_exts setup_pip_toolchain remove_stale_cuda12_nvidia_wheels uninstall_stale_flashinfer install_pytorch_stack install_cuda12_deepep_wheel + setup_cargo_cache install_sglang + release_cargo_cache_lock install_nccl # Diffusion B200 CI imports torch inside install_sglang_kernel after removing # stale CUDA 12 NVIDIA wheels, so opt into one early LD_LIBRARY_PATH refresh. diff --git a/scripts/ci/utils/stage_rust_ext_modules.sh b/scripts/ci/utils/stage_rust_ext_modules.sh index 63110075f..9bb1b04b8 100755 --- a/scripts/ci/utils/stage_rust_ext_modules.sh +++ b/scripts/ci/utils/stage_rust_ext_modules.sh @@ -1,7 +1,7 @@ #!/bin/bash -# Copy the built PyO3 extension modules into rust-ext-staging/rust_extensions/ for -# upload-artifact. Shared by both jobs of _pr-test-rust-ext-build.yml, so the -# archive layout and the module-count check cannot drift between them. +# Copy the built PyO3 extension modules into their package-relative paths under +# rust-ext-staging/. Shared by both jobs of _pr-test-rust-ext-build.yml, so the +# archive layout and module-count checks cannot drift between them. # # MAX_GLIBC (optional): also reject a module requiring a newer GLIBC symbol # version than the test runners have. Only set where the modules were just @@ -32,6 +32,23 @@ for module in server grpc multimodal; do cp "${found[@]}" rust-ext-staging/rust_extensions/ built+=("${found[@]}") done + +mkdir -p rust-ext-staging/mem_cache/rust_tree_core +for module in mem_cache mem_cache_inspection; do + tree_core=(python/sglang/srt/mem_cache/rust_tree_core/"${module}".*.so) + if [ ${#tree_core[@]} -eq 0 ]; then + echo "::error::no Rust TreeCore ${module} extension module found" + exit 1 + fi + tree_core_suffixes=$(printf '%s\n' "${tree_core[@]##*/${module}}" | sort) + if [ "${tree_core_suffixes}" != "${expected_suffixes}" ]; then + echo "::error::Rust TreeCore ${module} extension does not match the interpreter set" + printf 'have:\n%s\nwant:\n%s\n' "${tree_core_suffixes}" "${expected_suffixes}" + exit 1 + fi + cp "${tree_core[@]}" rust-ext-staging/mem_cache/rust_tree_core/ + built+=("${tree_core[@]}") +done max_allowed="${MAX_GLIBC:-}" [ -n "${max_allowed}" ] || exit 0 diff --git a/scripts/release/prepare_sglang_wheel.py b/scripts/release/prepare_sglang_wheel.py new file mode 100755 index 000000000..140f3a92f --- /dev/null +++ b/scripts/release/prepare_sglang_wheel.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python3 +"""Repair an SGLang wheel and smoke-test its production Rust TreeCore.""" + +from __future__ import annotations + +import argparse +import os +import shutil +import subprocess +import sys +import tempfile +import textwrap +import zipfile +from email.parser import BytesParser +from pathlib import Path, PurePosixPath + +_LIBTORCH_EXCLUDES = ( + "libc10.so", + "libc10_cuda.so", + "libtorch.so", + "libtorch_cpu.so", + "libtorch_cuda.so", + "libtorch_python.so", +) +_TREE_CORE_DIR = PurePosixPath("sglang/srt/mem_cache/rust_tree_core") +_BINDING_CLASSES = ( + "RustUnifiedTreeCoreBinding", + "RustBigramUnifiedTreeCoreBinding", + "TreeCoreInitParamsBinding", +) + + +def _single_wheel(directory: Path) -> Path: + wheels = sorted(directory.glob("*.whl")) + if len(wheels) != 1: + raise RuntimeError(f"expected one wheel in {directory}, found {wheels}") + return wheels[0] + + +def _metadata(wheel: Path) -> tuple[str, str]: + with zipfile.ZipFile(wheel) as archive: + metadata_files = [ + name for name in archive.namelist() if name.endswith(".dist-info/METADATA") + ] + if len(metadata_files) != 1: + raise RuntimeError( + f"expected one METADATA file in {wheel}, found {metadata_files}" + ) + metadata = BytesParser().parsebytes(archive.read(metadata_files[0])) + return str(metadata["Name"]), str(metadata["Version"]) + + +def _smoke_test_tree_core(wheel: Path) -> None: + with tempfile.TemporaryDirectory(prefix="sglang-wheel-smoke-") as temp_dir: + root = Path(temp_dir) + with zipfile.ZipFile(wheel) as archive: + names = archive.namelist() + inspection_modules = [ + name + for name in names + if PurePosixPath(name).parent == _TREE_CORE_DIR + and PurePosixPath(name).name.startswith("mem_cache_inspection") + and name.endswith(".so") + ] + if inspection_modules: + raise RuntimeError( + f"production wheel contains inspection modules: {inspection_modules}" + ) + production_modules = [ + name + for name in names + if PurePosixPath(name).parent == _TREE_CORE_DIR + and PurePosixPath(name).name.startswith("mem_cache.") + and name.endswith(".so") + ] + if len(production_modules) != 1: + raise RuntimeError( + "expected one production Rust TreeCore module, found " + f"{production_modules}" + ) + + install_dir = root / "installed" + subprocess.run( + [ + sys.executable, + "-m", + "pip", + "install", + "--disable-pip-version-check", + "--no-compile", + "--no-deps", + "--no-index", + "--target", + os.fspath(install_dir), + os.fspath(wheel), + ], + check=True, + ) + + smoke_script = textwrap.dedent(f""" + import sys + import types + from pathlib import Path + + site_packages = Path({os.fspath(install_dir)!r}).resolve() + sys.path.insert(0, str(site_packages)) + package = types.ModuleType("sglang") + package.__package__ = "sglang" + package.__path__ = [str(site_packages / "sglang")] + sys.modules["sglang"] = package + + from sglang.srt.mem_cache.rust_tree_core.extension import bindings + + module_path = Path(bindings.__file__).resolve() + if site_packages not in module_path.parents: + raise RuntimeError( + f"loaded TreeCore outside installed wheel: {{module_path}}" + ) + if bindings.__name__ != "sglang.srt.mem_cache.rust_tree_core.mem_cache": + raise RuntimeError( + f"loaded unexpected TreeCore module: {{bindings.__name__}}" + ) + for class_name in {_BINDING_CLASSES!r}: + binding = getattr(bindings, class_name, None) + if binding is None: + raise RuntimeError( + f"production TreeCore is missing {{class_name}}" + ) + inspection_methods = [ + name for name in dir(binding) if name.startswith("inspect_") + ] + if inspection_methods: + raise RuntimeError( + f"production {{class_name}} exposes inspection methods: " + f"{{inspection_methods}}" + ) + + from array import array + + hashes = bindings.get_hash_str(array("q", [1, 2]), None, 1) + if len(hashes) != 2 or any(len(value) != 64 for value in hashes): + raise RuntimeError(f"unexpected TreeCore hash result: {{hashes}}") + """) + environment = os.environ.copy() + environment["SGLANG_RUST_BUILD_MODE"] = "never" + environment.pop("PYTHONPATH", None) + subprocess.run( + [sys.executable, "-I", "-c", smoke_script], + cwd=root, + env=environment, + check=True, + ) + + +def _write_github_outputs(path: Path, *, wheel: Path, version: str) -> None: + with path.open("a", encoding="utf-8") as output: + output.write(f"wheel_filename={wheel.name}\n") + output.write(f"wheel_version={version}\n") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("wheel_dir", type=Path) + parser.add_argument("--github-output", type=Path) + args = parser.parse_args() + + wheel_dir = args.wheel_dir.resolve() + source_wheel = _single_wheel(wheel_dir) + with tempfile.TemporaryDirectory( + prefix="sglang-wheel-repair-", dir=wheel_dir.parent + ) as repair_dir: + command = [ + sys.executable, + "-m", + "auditwheel", + "repair", + os.fspath(source_wheel), + "--wheel-dir", + repair_dir, + ] + for library in _LIBTORCH_EXCLUDES: + command.extend(("--exclude", library)) + subprocess.run(command, check=True) + + repaired_wheel = _single_wheel(Path(repair_dir)) + name, version = _metadata(repaired_wheel) + if name.casefold() != "sglang": + raise RuntimeError(f"expected sglang wheel, found {name!r}") + _smoke_test_tree_core(repaired_wheel) + + destination = wheel_dir / repaired_wheel.name + source_wheel.unlink() + shutil.move(repaired_wheel, destination) + + if args.github_output is not None: + _write_github_outputs( + args.github_output.resolve(), wheel=destination, version=version + ) + print(f"Prepared {destination.name} (sglang {version})") + + +if __name__ == "__main__": + main() diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py index 874b4396c..e553b433a 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py @@ -18,13 +18,14 @@ from sglang.test.test_utils import ( is_in_ci, popen_launch_server, terminate_and_kill_process_tree, + unified_radix_tree_server_env, ) DSV4_FLASH_MODEL = "sgl-project/DeepSeek-V4-Flash-FP8" DSV4_DSPARK_MODEL = "deepseek-ai/DeepSeek-V4-Flash-DSpark" DSV4_FLASH_LAUNCH_TIMEOUT = 3600 -register_cuda_ci(est_time=2400, stage="extra-b", runner_config="4-gpu-h100") +register_cuda_ci(est_time=4800, stage="extra-b", runner_config="4-gpu-h100") def _assert_dsv4_decode_cached_tokens(result, history_len, output_len, label): @@ -37,6 +38,7 @@ def _assert_dsv4_decode_cached_tokens(result, history_len, output_len, label): class TestUnifiedDeepSeekV4FlashHiCache(UnifiedRadixTreeTestMixin, CustomTestCase): """DeepSeek V4 Flash FP8 + HiCache + UnifiedRadixCache.""" + tree_core_backend = "python" tp_size = 4 pp_size = 1 hicache_io_backend = "direct" @@ -100,10 +102,10 @@ class TestUnifiedDeepSeekV4FlashHiCache(UnifiedRadixTreeTestMixin, CustomTestCas cls.base_url, timeout=DSV4_FLASH_LAUNCH_TIMEOUT, other_args=cls._server_args(), - env={ - "SGLANG_DSV4_FP4_EXPERTS": "0", - "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1", - }, + env=unified_radix_tree_server_env( + cls.tree_core_backend, + SGLANG_DSV4_FP4_EXPERTS="0", + ), ) cls.input_ids = get_input_ids(cls.model, num_samples=18) @@ -127,6 +129,7 @@ class TestUnifiedDeepSeekV4FlashHiCachePageFirstDirect( class TestUnifiedDeepSeekV4FlashHiCacheL3(AccuracyTwoPassMixin, CustomTestCase): """DeepSeek V4 Flash FP8 + HiCache L3 (file backend) + UnifiedRadixCache.""" + tree_core_backend = "python" l3_prefetch_page_size = 256 l3_prefetch_prompt_pages = 4 max_running_requests = 4 @@ -171,11 +174,11 @@ class TestUnifiedDeepSeekV4FlashHiCacheL3(AccuracyTwoPassMixin, CustomTestCase): "--max-running-requests", str(cls.max_running_requests), ], - env={ - "SGLANG_DSV4_FP4_EXPERTS": "0", - "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1", - "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.hicache_dir, - }, + env=unified_radix_tree_server_env( + cls.tree_core_backend, + SGLANG_DSV4_FP4_EXPERTS="0", + SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR=cls.hicache_dir, + ), ) @classmethod @@ -188,6 +191,7 @@ class TestUnifiedDeepSeekV4FlashHiCacheL3(AccuracyTwoPassMixin, CustomTestCase): class TestUnifiedDeepSeekV4FlashEagleHiCacheL3(AccuracyTwoPassMixin, CustomTestCase): """DeepSeek V4 Flash EAGLE + HiCache L3 should load from storage.""" + tree_core_backend = "python" page_size = 256 l3_prefetch_page_size = 256 l3_prefetch_prompt_pages = 4 @@ -248,11 +252,11 @@ class TestUnifiedDeepSeekV4FlashEagleHiCacheL3(AccuracyTwoPassMixin, CustomTestC "--speculative-num-draft-tokens", "4", ], - env={ - "SGLANG_DSV4_FP4_EXPERTS": "0", - "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1", - "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.hicache_dir, - }, + env=unified_radix_tree_server_env( + cls.tree_core_backend, + SGLANG_DSV4_FP4_EXPERTS="0", + SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR=cls.hicache_dir, + ), ) @classmethod @@ -385,12 +389,38 @@ class TestUnifiedDeepSeekV4FlashDSparkHiCacheL3( "--speculative-algorithm", "DSPARK", ], - env={ - "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1", - "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.hicache_dir, - }, + env=unified_radix_tree_server_env( + cls.tree_core_backend, + SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR=cls.hicache_dir, + ), ) +class TestRustUnifiedDeepSeekV4FlashHiCache(TestUnifiedDeepSeekV4FlashHiCache): + tree_core_backend = "rust" + + +class TestRustUnifiedDeepSeekV4FlashHiCachePageFirstDirect( + TestUnifiedDeepSeekV4FlashHiCachePageFirstDirect +): + tree_core_backend = "rust" + + +class TestRustUnifiedDeepSeekV4FlashHiCacheL3(TestUnifiedDeepSeekV4FlashHiCacheL3): + tree_core_backend = "rust" + + +class TestRustUnifiedDeepSeekV4FlashEagleHiCacheL3( + TestUnifiedDeepSeekV4FlashEagleHiCacheL3 +): + tree_core_backend = "rust" + + +class TestRustUnifiedDeepSeekV4FlashDSparkHiCacheL3( + TestUnifiedDeepSeekV4FlashDSparkHiCacheL3 +): + tree_core_backend = "rust" + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py index 96f357fc3..f8e60fcbd 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py @@ -9,10 +9,11 @@ from sglang.test.test_utils import ( CustomTestCase, popen_launch_server, terminate_and_kill_process_tree, + unified_radix_tree_server_env, ) -register_cuda_ci(est_time=250, stage="base-b", runner_config="2-gpu-large") -register_amd_ci(est_time=400, suite="stage-b-test-2-gpu-large-amd") +register_cuda_ci(est_time=500, stage="base-b", runner_config="2-gpu-large") +register_amd_ci(est_time=800, suite="stage-b-test-2-gpu-large-amd") FULL_MODEL = "Qwen/Qwen3-32B" @@ -20,6 +21,7 @@ FULL_MODEL = "Qwen/Qwen3-32B" class TestUnifiedFullRadixCache(UnifiedRadixTreeTestMixin, CustomTestCase): """Full attention.""" + tree_core_backend = "python" kl_threshold = 0.0025 @classmethod @@ -38,7 +40,7 @@ class TestUnifiedFullRadixCache(UnifiedRadixTreeTestMixin, CustomTestCase): "--page-size", "64", ], - env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}, + env=unified_radix_tree_server_env(cls.tree_core_backend), ) cls.input_ids = get_input_ids(cls.model, num_samples=18) @@ -47,5 +49,9 @@ class TestUnifiedFullRadixCache(UnifiedRadixTreeTestMixin, CustomTestCase): terminate_and_kill_process_tree(cls.process, wait_timeout=60) +class TestRustUnifiedFullRadixCache(TestUnifiedFullRadixCache): + tree_core_backend = "rust" + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py index 7f4ef8b40..aa7583bee 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py @@ -70,9 +70,10 @@ from sglang.test.test_utils import ( CustomTestCase, popen_launch_server, terminate_and_kill_process_tree, + unified_radix_tree_server_env, ) -register_cuda_ci(est_time=1150, stage="base-b", runner_config="1-gpu-large") +register_cuda_ci(est_time=2300, stage="base-b", runner_config="1-gpu-large") _MODEL_PATH = os.environ.get("INKLING_TEST_MODEL_PATH", "thinkingmachines/Inkling") _MODEL_REVISION = os.environ.get("INKLING_TEST_MODEL_REVISION", "test") @@ -149,6 +150,8 @@ class TestUnifiedHybridBitExact(CustomTestCase): decode-region state reuse in general rather than that regression. """ + tree_core_backend = "python" + @classmethod def setUpClass(cls): cls.model = _MODEL_PATH @@ -168,7 +171,7 @@ class TestUnifiedHybridBitExact(CustomTestCase): cls.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, other_args=other_args, - env={**os.environ, "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}, + env=unified_radix_tree_server_env(cls.tree_core_backend), ) @classmethod @@ -226,7 +229,7 @@ class TestUnifiedHybridLazyBitExact(TestUnifiedHybridBitExact): cls.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, other_args=other_args, - env={**os.environ, "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}, + env=unified_radix_tree_server_env(cls.tree_core_backend), ) @@ -243,6 +246,8 @@ class TestUnifiedHybridHiCacheBitExact(CustomTestCase): cannot produce a non-aligned hit length, which this regression needs. """ + tree_core_backend = "python" + @classmethod def setUpClass(cls): cls.model = _MODEL_PATH @@ -277,7 +282,7 @@ class TestUnifiedHybridHiCacheBitExact(CustomTestCase): cls.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, other_args=other_args, - env={**os.environ, "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}, + env=unified_radix_tree_server_env(cls.tree_core_backend), ) cls.input_ids = get_input_ids( tokenizer_path=cls.model, num_samples=9, trust_remote_code=True @@ -333,6 +338,8 @@ class TestUnifiedHybridMTPBitExact(CustomTestCase): environment override; a regression there surfaces here as a nonzero KL. """ + tree_core_backend = "python" + @classmethod def setUpClass(cls): cls.model = _MODEL_PATH @@ -360,10 +367,7 @@ class TestUnifiedHybridMTPBitExact(CustomTestCase): cls.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, other_args=other_args, - env={ - **os.environ, - "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1", - }, + env=unified_radix_tree_server_env(cls.tree_core_backend), ) @classmethod @@ -396,5 +400,21 @@ class TestUnifiedHybridMTPBitExact(CustomTestCase): self._run(assert_decode_cache_hit) +class TestRustUnifiedHybridBitExact(TestUnifiedHybridBitExact): + tree_core_backend = "rust" + + +class TestRustUnifiedHybridLazyBitExact(TestUnifiedHybridLazyBitExact): + tree_core_backend = "rust" + + +class TestRustUnifiedHybridHiCacheBitExact(TestUnifiedHybridHiCacheBitExact): + tree_core_backend = "rust" + + +class TestRustUnifiedHybridMTPBitExact(TestUnifiedHybridMTPBitExact): + tree_core_backend = "rust" + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/rust/test_run_mem_cache_rust_tests.py b/test/registered/rust/test_run_mem_cache_rust_tests.py new file mode 100644 index 000000000..55758516c --- /dev/null +++ b/test/registered/rust/test_run_mem_cache_rust_tests.py @@ -0,0 +1,64 @@ +"""Run the standalone mem-cache crate's native Rust unit tests.""" + +import shutil +import subprocess +import unittest +from pathlib import Path + +from sglang.srt.environ import envs +from sglang.srt.rust_extensions.torch_build import torch_build_configuration +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +BUILD_AND_RUN_TIMEOUT_S = 900 +RUST_WORKSPACE = Path(__file__).resolve().parents[3] / "rust" +MEM_CACHE_MANIFEST = RUST_WORKSPACE / "mem-cache" / "Cargo.toml" + +register_cpu_ci(est_time=900, suite="base-a-test-cpu") + + +@unittest.skipIf( + envs.SGLANG_SKIP_RUST_TESTS.get(), + "SGLANG_SKIP_RUST_TESTS is set (no rust/ workspace changes per CI check-changes)", +) +class TestMemCacheCargo(CustomTestCase): + def test_mem_cache_native_tests(self): + self.assertIsNotNone( + shutil.which("cargo"), + "cargo not found on PATH; install a Rust toolchain " + "(scripts/ci/utils/install_rust_protoc.sh)", + ) + self.assertTrue( + MEM_CACHE_MANIFEST.is_file(), + f"mem-cache manifest not found at {MEM_CACHE_MANIFEST}", + ) + build = torch_build_configuration( + compat_header=MEM_CACHE_MANIFEST.parent / "torch_2_13_compat.h", + python_module="sglang.srt.mem_cache.rust_tree_core.mem_cache", + ) + proc = subprocess.run( + [ + "cargo", + "test", + "--manifest-path", + str(MEM_CACHE_MANIFEST), + "--locked", + "--no-default-features", + ], + cwd=RUST_WORKSPACE, + env=build.environment, + capture_output=True, + text=True, + timeout=BUILD_AND_RUN_TIMEOUT_S, + ) + print(proc.stdout) + self.assertEqual( + proc.returncode, + 0, + f"mem-cache native tests failed\n" + f"--- stdout ---\n{proc.stdout}\n--- stderr ---\n{proc.stderr}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/rust/test_run_rust_tests.py b/test/registered/rust/test_run_rust_tests.py index 085dcf1b7..67679d19f 100644 --- a/test/registered/rust/test_run_rust_tests.py +++ b/test/registered/rust/test_run_rust_tests.py @@ -1,4 +1,4 @@ -"""Run the `rust/` Cargo workspace's unit tests from the CPU CI suite.""" +"""Run the repository's native Rust unit tests from the CPU CI suite.""" import shutil import subprocess @@ -11,7 +11,6 @@ from sglang.test.test_utils import CustomTestCase BUILD_AND_RUN_TIMEOUT_S = 900 RUST_WORKSPACE = Path(__file__).resolve().parents[3] / "rust" - register_cpu_ci(est_time=900, suite="base-a-test-cpu") @@ -23,6 +22,24 @@ register_cpu_ci(est_time=900, suite="base-a-test-cpu") "SGLANG_SKIP_RUST_TESTS is set (no rust/ workspace changes per CI check-changes)", ) class TestCargoWorkspace(CustomTestCase): + def _run_cargo(self, args: list[str], *, cwd: Path, env: dict | None = None): + proc = subprocess.run( + ["cargo", *args], + cwd=cwd, + env=env, + capture_output=True, + text=True, + timeout=BUILD_AND_RUN_TIMEOUT_S, + ) + # Print unconditionally so a green run still shows which tests ran. + print(proc.stdout) + self.assertEqual( + proc.returncode, + 0, + f"`cargo {' '.join(args)}` failed in {cwd}\n" + f"--- stdout ---\n{proc.stdout}\n--- stderr ---\n{proc.stderr}", + ) + def test_cargo_test_workspace(self): # Not skipUnless: cargo is a hard dependency of the editable install # (setuptools-rust builds sglang-grpc), so a missing toolchain is a @@ -37,21 +54,7 @@ class TestCargoWorkspace(CustomTestCase): f"rust workspace manifest not found at {RUST_WORKSPACE}", ) - proc = subprocess.run( - ["cargo", "test", "--workspace"], - cwd=RUST_WORKSPACE, - capture_output=True, - text=True, - timeout=BUILD_AND_RUN_TIMEOUT_S, - ) - # Print unconditionally so a green run still shows which tests ran. - print(proc.stdout) - self.assertEqual( - proc.returncode, - 0, - f"`cargo test --workspace` failed in {RUST_WORKSPACE}\n" - f"--- stdout ---\n{proc.stdout}\n--- stderr ---\n{proc.stderr}", - ) + self._run_cargo(["test", "--workspace"], cwd=RUST_WORKSPACE) if __name__ == "__main__": diff --git a/test/registered/rust/test_rust_extension.py b/test/registered/rust/test_rust_extension.py index dc3fe94dd..85ea9c2d1 100644 --- a/test/registered/rust/test_rust_extension.py +++ b/test/registered/rust/test_rust_extension.py @@ -9,11 +9,12 @@ import time import unittest from pathlib import Path from tempfile import TemporaryDirectory -from types import ModuleType +from types import ModuleType, SimpleNamespace from unittest import mock from sglang.srt.rust_extensions import load_rust_extension from sglang.srt.rust_extensions import loader as rust_extension +from sglang.srt.rust_extensions.torch_build import torch_build_configuration from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -84,6 +85,74 @@ crate-type = ["cdylib"] fingerprint.assert_not_called() cargo_build.assert_not_called() + def test_bundled_named_variant_never_touches_source_or_cargo(self): + bundled = ModuleType("demo._inspection") + with ( + mock.patch.object( + rust_extension.importlib, "import_module", return_value=bundled + ) as import_module, + mock.patch.object(rust_extension, "_discover_crate") as discover, + mock.patch.object(rust_extension, "_build_context") as fingerprint, + mock.patch.object(rust_extension, "_cargo_build") as cargo_build, + ): + self.assertIs( + load_rust_extension( + "demo._core", + mode="never", + workspace=Path("/workspace/not-present"), + additional_features=("inspection",), + extension_module="demo._inspection", + ), + bundled, + ) + import_module.assert_called_once_with("demo._inspection") + discover.assert_not_called() + fingerprint.assert_not_called() + cargo_build.assert_not_called() + + def test_auto_ignores_a_stale_bundled_extension_in_a_source_tree(self): + with TemporaryDirectory() as directory: + root = Path(directory) + workspace = self._workspace(root) + (workspace / "demo/lib.rs").write_text( + "fn source_changed() {}\n", encoding="utf-8" + ) + stale = ModuleType("demo._core") + built = ModuleType("demo._core") + artifact = root / "libdemo_extension.so" + artifact.write_bytes(b"fresh extension") + context = rust_extension._BuildContext( + "changed-source", "fingerprint", "target" + ) + with ( + mock.patch.object( + rust_extension, "_import_bundled_extension", return_value=stale + ) as bundled_import, + mock.patch.object( + rust_extension, "_build_context", return_value=context + ), + mock.patch.object( + rust_extension, "_source_digest", return_value="changed-source" + ), + mock.patch.object( + rust_extension, "_cargo_build", return_value=artifact + ) as cargo_build, + mock.patch.object( + rust_extension, "_load_extension_from_path", return_value=built + ), + ): + self.assertIs( + load_rust_extension( + "demo._core", + mode="auto", + workspace=workspace, + cache_dir=root / "cache", + ), + built, + ) + bundled_import.assert_not_called() + cargo_build.assert_called_once() + def test_discovery_reads_crate_manifest_metadata(self): with TemporaryDirectory() as directory: workspace = self._workspace(Path(directory)) @@ -126,6 +195,18 @@ crate-type = ["cdylib"] changed_flags.target_fingerprint, ) + inspection = rust_extension._build_context( + crate, + features=(*crate.features, "inspection"), + extension_module="demo._inspection", + build_fingerprint={"torch": "2.13"}, + ) + self.assertNotEqual(changed_source.fingerprint, inspection.fingerprint) + self.assertNotEqual( + changed_source.target_fingerprint, + inspection.target_fingerprint, + ) + def test_auto_builds_once_then_uses_cache(self): with TemporaryDirectory() as directory: root = Path(directory) @@ -155,13 +236,19 @@ crate-type = ["cdylib"] ): self.assertIs( rust_extension.load_rust_extension( - "demo._core", workspace=workspace, cache_dir=root / "cache" + "demo._core", + mode="auto", + workspace=workspace, + cache_dir=root / "cache", ), loaded, ) self.assertIs( rust_extension.load_rust_extension( - "demo._core", workspace=workspace, cache_dir=root / "cache" + "demo._core", + mode="auto", + workspace=workspace, + cache_dir=root / "cache", ), loaded, ) @@ -270,6 +357,122 @@ crate-type = ["cdylib"] ], ) + def test_variant_uses_its_own_module_name_features_and_environment(self): + with TemporaryDirectory() as directory: + root = Path(directory) + workspace = self._workspace(root) + artifact = root / "libdemo_extension.so" + artifact.write_bytes(b"extension") + context = rust_extension._BuildContext("source", "fingerprint", "target") + loaded = ModuleType("demo._inspection") + environment = {"CUSTOM_BUILD_INPUT": "value"} + with ( + mock.patch.object( + rust_extension, "_import_bundled_extension", return_value=None + ) as bundled_import, + mock.patch.object( + rust_extension, "_build_context", return_value=context + ) as build_context, + mock.patch.object( + rust_extension, "_source_digest", return_value="source" + ), + mock.patch.object( + rust_extension, "_cargo_build", return_value=artifact + ) as cargo_build, + mock.patch.object( + rust_extension, + "_load_extension_from_path", + return_value=loaded, + ) as load_from_path, + ): + self.assertIs( + load_rust_extension( + "demo._core", + mode="auto", + workspace=workspace, + cache_dir=root / "cache", + additional_features=("inspection",), + extension_module="demo._inspection", + build_environment=environment, + build_fingerprint={"native": "abi"}, + ), + loaded, + ) + bundled_import.assert_not_called() + self.assertEqual( + build_context.call_args.kwargs, + { + "features": ("python", "inspection"), + "build_fingerprint": {"native": "abi"}, + "extension_module": "demo._inspection", + }, + ) + self.assertEqual( + cargo_build.call_args.kwargs, + { + "features": ("python", "inspection"), + "build_environment": environment, + }, + ) + self.assertEqual(load_from_path.call_args.args[0], "demo._inspection") + + def test_torch_build_configuration_is_versioned_and_relocatable(self): + with TemporaryDirectory() as directory: + root = Path(directory) + torch_root = root / "torch" + (torch_root / "lib").mkdir(parents=True) + torch_init = torch_root / "__init__.py" + torch_init.write_text("", encoding="utf-8") + compat_header = root / "compat.h" + compat_header.write_text("// compatibility\n", encoding="utf-8") + fake_torch = SimpleNamespace( + __version__="2.13.0+cu130", + __file__=str(torch_init), + compiled_with_cxx11_abi=lambda: True, + version=SimpleNamespace(cuda="13.0", hip=None), + ) + + build = torch_build_configuration( + compat_header=compat_header, + python_module="sglang.srt.mem_cache.rust_tree_core.mem_cache", + torch_module=fake_torch, + base_environment={ + "PATH": "/usr/bin", + "CXXFLAGS": "-O2", + "RUSTFLAGS": "-Ctarget-cpu=x86-64", + }, + ) + + self.assertEqual(build.environment["LIBTORCH_USE_PYTORCH"], "1") + self.assertEqual(build.environment["LIBTORCH_BYPASS_VERSION_CHECK"], "1") + self.assertIn(str(compat_header), build.environment["CXXFLAGS"]) + self.assertIn( + "$ORIGIN/../../../../torch/lib", build.environment["RUSTFLAGS"] + ) + self.assertIn(str(torch_root / "lib"), build.environment["RUSTFLAGS"]) + self.assertEqual(build.fingerprint["torch_version"], "2.13.0+cu130") + self.assertTrue(build.fingerprint["torch_cxx11_abi"]) + + wheel_build = torch_build_configuration( + compat_header=compat_header, + python_module="sglang.srt.mem_cache.rust_tree_core.mem_cache", + torch_module=fake_torch, + base_environment={}, + include_absolute_rpath=False, + ) + self.assertNotIn( + str(torch_root / "lib"), wheel_build.environment["RUSTFLAGS"] + ) + self.assertFalse(wheel_build.fingerprint["include_absolute_rpath"]) + + fake_torch.__version__ = "2.14.0" + with self.assertRaisesRegex(RuntimeError, "PyTorch 2.11 through 2.13"): + torch_build_configuration( + compat_header=compat_header, + python_module="sglang.srt.mem_cache.rust_tree_core.mem_cache", + torch_module=fake_torch, + ) + def test_filesystem_lock_serializes_processes(self): with TemporaryDirectory() as directory: lock_path = Path(directory) / "build.lock" @@ -337,6 +540,12 @@ crate-type = ["cdylib"] "sglang_mm_core", ("python", "parallel"), ), + ( + "sglang.srt.mem_cache.rust_tree_core.mem_cache", + "mem_cache", + "mem_cache", + ("python-extension",), + ), ): crate = rust_extension._discover_crate( rust_extension._RUST_WORKSPACE, python_module diff --git a/test/registered/unit/disaggregation/test_decode_hicache_tree_core.py b/test/registered/unit/disaggregation/test_decode_hicache_tree_core.py new file mode 100644 index 000000000..d5902cb14 --- /dev/null +++ b/test/registered/unit/disaggregation/test_decode_hicache_tree_core.py @@ -0,0 +1,107 @@ +"""Unit tests for decode HiCache TreeCore interactions.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import Mock + +import torch + +from sglang.srt.disaggregation.decode_hicache_mixin import ( + DecodeHiCachePreallocMixin, + DecodePrefixMatch, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestDecodeHiCacheTreeCore(CustomTestCase): + def test_storage_probe_and_prefetch_use_node_handles(self): + ongoing_prefetch = {} + + def register_prefetch(req_id, *_args, **_kwargs): + ongoing_prefetch[req_id] = object() + + tree_cache = SimpleNamespace( + hicache_storage_pass_prefix_keys=True, + ongoing_prefetch=ongoing_prefetch, + is_backuped=Mock(return_value=True), + is_root=Mock(return_value=False), + get_last_hash_value=Mock(return_value="h2"), + get_prefix_hash_values=Mock(return_value=["h0", "h1"]), + query_storage_hit_length=Mock(return_value=2), + prefetch_from_storage=Mock(side_effect=register_prefetch), + ) + harness = SimpleNamespace( + scheduler=SimpleNamespace(enable_decode_hicache=True), + tree_cache=tree_cache, + ) + req = SimpleNamespace( + rid="req-0", + origin_input_ids=[0, 1, 2, 3, 4, 5, 6, 7], + extra_key="model", + cache_salt=None, + ) + result = SimpleNamespace( + device_indices=torch.tensor([10, 11]), + host_hit_length=2, + last_device_node=11, + last_host_node=22, + ) + + prefix_match = DecodeHiCachePreallocMixin._build_decode_prefix_match( + harness, req, result + ) + + self.assertEqual(prefix_match.l3_storage_hit_length, 2) + tree_cache.query_storage_hit_length.assert_called_once_with( + 22, [4, 5, 6, 7], "h2", ["h0", "h1"] + ) + + DecodeHiCachePreallocMixin._start_hicache_prefetch(harness, req, prefix_match) + + self.assertTrue(prefix_match.prefetch_registered) + tree_cache.prefetch_from_storage.assert_called_once_with( + "req-0", + 22, + [4, 5], + "h2", + ["h0", "h1"], + extra_key="model", + cache_salt=None, + ) + + def test_stale_prefetch_anchor_degrades_to_l2(self): + tree_cache = SimpleNamespace( + hicache_storage_pass_prefix_keys=True, + ongoing_prefetch={}, + get_last_hash_value=Mock(side_effect=KeyError(22)), + get_prefix_hash_values=Mock(), + prefetch_from_storage=Mock(), + ) + harness = SimpleNamespace(tree_cache=tree_cache) + req = SimpleNamespace( + rid="req-0", + origin_input_ids=[0, 1, 2, 3, 4, 5], + extra_key=None, + cache_salt=None, + ) + prefix_match = DecodePrefixMatch( + prefix_indices=torch.tensor([10, 11]), + l2_host_hit_length=2, + l3_storage_hit_length=2, + last_device_node=11, + last_host_node=22, + ) + + DecodeHiCachePreallocMixin._start_hicache_prefetch(harness, req, prefix_match) + + self.assertEqual(prefix_match.l3_storage_hit_length, 0) + self.assertFalse(prefix_match.prefetch_registered) + tree_cache.get_prefix_hash_values.assert_not_called() + tree_cache.prefetch_from_storage.assert_not_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/managers/test_schedule_policy_dfs_weight.py b/test/registered/unit/managers/test_schedule_policy_dfs_weight.py new file mode 100644 index 000000000..8afef7d55 --- /dev/null +++ b/test/registered/unit/managers/test_schedule_policy_dfs_weight.py @@ -0,0 +1,55 @@ +"""Unit tests for DFS-weight schedule-policy delegation.""" + +import unittest +from types import SimpleNamespace + +from sglang.srt.managers.schedule_policy import SchedulePolicy +from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestSchedulePolicyDfsWeight(CustomTestCase): + def test_orders_requests_by_subtree_weight(self): + class Node: + def __init__(self): + self.children = {} + + root = Node() + branch_a = Node() + branch_b = Node() + leaf_a1 = Node() + leaf_a2 = Node() + root.children = {"a": branch_a, "b": branch_b} + branch_a.children = {"a1": leaf_a1, "a2": leaf_a2} + + class TreeCache: + dfs_weight_order = BasePrefixCache.dfs_weight_order + + def __init__(self): + self.root_node = root + + @staticmethod + def resolve_node_handle(node): + return node + + waiting_queue = [ + SimpleNamespace(last_node=branch_b, name="b"), + SimpleNamespace(last_node=leaf_a2, name="a2"), + SimpleNamespace(last_node=leaf_a1, name="a1-first"), + SimpleNamespace(last_node=leaf_a1, name="a1-second"), + SimpleNamespace(last_node=branch_a, name="a-parent"), + ] + + SchedulePolicy._sort_by_dfs_weight(waiting_queue, TreeCache()) + + self.assertEqual( + [req.name for req in waiting_queue], + ["a1-first", "a1-second", "a2", "a-parent", "b"], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/rust_unified_tree_core_inspector.py b/test/registered/unit/mem_cache/rust_unified_tree_core_inspector.py new file mode 100644 index 000000000..097e4ffa9 --- /dev/null +++ b/test/registered/unit/mem_cache/rust_unified_tree_core_inspector.py @@ -0,0 +1,228 @@ +"""Test-only inspection adapter for the Rust Unified TreeCore.""" + +from __future__ import annotations + +from typing import Optional + +import torch +from unified_tree_core_inspection_interface import ( + UnifiedTreeCoreInspectionInterface, +) + +from sglang.srt.mem_cache.base_prefix_cache import MatchPrefixParams, MatchResult +from sglang.srt.mem_cache.rust_tree_core.adapter import ( + RustUnifiedTreeCore, + _fill_evict_result, + _match_result_from_binding, + _radix_key_buffer, +) +from sglang.srt.mem_cache.rust_tree_core.extension import load_tree_core_extension +from sglang.srt.mem_cache.unified_cache.components import ComponentType, EvictLayer +from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import ( + BaseEvictionResult, + NodeId, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci( + est_time=0, suite="base-a-test-cpu", disabled="Rust TreeCore test inspector" +) + +_inspection_bindings = load_tree_core_extension(inspection=True) + + +class RustUnifiedTreeCoreInspector( + RustUnifiedTreeCore, UnifiedTreeCoreInspectionInterface +): + """Rust TreeCore variant used by the shared backend-conformance tests. + + The production adapter deliberately implements only + ``UnifiedTreeCoreInterface``. These forwarding methods keep white-box state + controls in test code while the binding returns snapshots rather than Rust + iterators across the Python boundary. + """ + + _bindings = _inspection_bindings + + def contains_node(self, node_id: NodeId) -> bool: + return self._binding.inspect_contains_node(node_id) + + def get_parent_node_id(self, node_id: NodeId) -> Optional[NodeId]: + return self._binding.inspect_get_parent_node_id(node_id) + + def get_child_node_ids(self, node_id: NodeId) -> list[NodeId]: + return self._binding.inspect_get_child_node_ids(node_id) + + def get_node_key_length(self, node_id: NodeId) -> int: + return self._binding.inspect_get_node_key_length(node_id) + + def get_node_token_ids(self, node_id: NodeId) -> list[int]: + return self._binding.inspect_get_node_token_ids(node_id) + + def is_node_key_bigram(self, node_id: NodeId) -> bool: + return self._binding.inspect_is_node_key_bigram(node_id) + + def get_component_host_value( + self, node_id: NodeId, component_type: ComponentType + ) -> Optional[torch.Tensor]: + return self._binding.inspect_get_component_host_value( + node_id, int(component_type) + ) + + def get_component_device_lock_ref( + self, node_id: NodeId, component_type: ComponentType + ) -> int: + return self._binding.inspect_get_component_device_lock_ref( + node_id, int(component_type) + ) + + def get_node_hit_count(self, node_id: NodeId) -> int: + return self._binding.inspect_get_node_hit_count(node_id) + + def get_write_through_pending_id(self, node_id: NodeId) -> Optional[int]: + return self._binding.inspect_get_write_through_pending_id(node_id) + + def is_node_in_device_lru( + self, node_id: NodeId, component_type: ComponentType + ) -> bool: + return self._binding.inspect_is_node_in_device_lru(node_id, int(component_type)) + + def is_node_in_host_lru( + self, node_id: NodeId, component_type: ComponentType + ) -> bool: + return self._binding.inspect_is_node_in_host_lru(node_id, int(component_type)) + + def get_component_device_lru_node_ids( + self, component_type: ComponentType + ) -> list[NodeId]: + return self._binding.inspect_get_component_device_lru_node_ids( + int(component_type) + ) + + def is_device_evictable_leaf(self, node_id: NodeId) -> bool: + return self._binding.inspect_is_device_evictable_leaf(node_id) + + def is_host_evictable_leaf(self, node_id: NodeId) -> bool: + return self._binding.inspect_is_host_evictable_leaf(node_id) + + def is_device_leaf(self, node_id: NodeId) -> bool: + return self._binding.inspect_is_device_leaf(node_id) + + def get_all_node_ids(self) -> list[NodeId]: + return self._binding.inspect_get_all_node_ids() + + def component_protected_size(self, component_type: ComponentType) -> int: + return self._binding.inspect_component_protected_size(int(component_type)) + + def set_node_hash_values( + self, node_id: NodeId, hash_values: Optional[list[str]] + ) -> None: + self._binding.inspect_set_node_hash_values(node_id, hash_values) + + def set_component_device_value_raw( + self, + node_id: NodeId, + component_type: ComponentType, + value: Optional[torch.Tensor], + ) -> None: + self._binding.inspect_set_component_device_value_raw( + node_id, int(component_type), value + ) + + def set_component_host_value_raw( + self, + node_id: NodeId, + component_type: ComponentType, + value: Optional[torch.Tensor], + ) -> None: + self._binding.inspect_set_component_host_value_raw( + node_id, int(component_type), value + ) + + def set_component_device_lock_ref( + self, node_id: NodeId, component_type: ComponentType, lock_ref: int + ) -> None: + assert lock_ref >= 0 + self._binding.inspect_set_component_device_lock_ref( + node_id, int(component_type), lock_ref + ) + + def remove_node_from_device_lru( + self, node_id: NodeId, component_type: ComponentType + ) -> None: + self._binding.inspect_remove_node_from_device_lru(node_id, int(component_type)) + + def insert_node_into_host_lru( + self, node_id: NodeId, component_type: ComponentType + ) -> None: + self._binding.inspect_insert_node_into_host_lru(node_id, int(component_type)) + + def set_component_evictable_size( + self, component_type: ComponentType, value: int + ) -> None: + assert value >= 0 + self._binding.inspect_set_component_evictable_size(int(component_type), value) + + def set_component_protected_size( + self, component_type: ComponentType, value: int + ) -> None: + assert value >= 0 + self._binding.inspect_set_component_protected_size(int(component_type), value) + + def update_duplicate_tracking(self, node_id: NodeId) -> None: + self._binding.inspect_update_duplicate_tracking(node_id) + + def advance_insert_walk_once(self) -> None: + self._binding.inspect_advance_insert_walk_once() + + def evict_component( + self, + node_id: NodeId, + component_type: ComponentType, + target: EvictLayer, + ) -> BaseEvictionResult: + binding_result = self._binding.inspect_evict_component( + node_id, int(component_type), int(target) + ) + return _fill_evict_result(binding_result, BaseEvictionResult()) + + def validate_cascade_evict( + self, + node_id: NodeId, + component_type: ComponentType, + target: EvictLayer, + ) -> None: + self._binding.inspect_validate_cascade_evict( + node_id, int(component_type), int(target) + ) + + def cleanup_tombstone_ancestors(self, node_id: NodeId) -> BaseEvictionResult: + binding_result = self._binding.inspect_cleanup_tombstone_ancestors(node_id) + return _fill_evict_result(binding_result, BaseEvictionResult()) + + def finalize_component_match_result( + self, + component_type: ComponentType, + result: MatchResult, + params: MatchPrefixParams, + value_chunks: list[torch.Tensor], + best_value_len: int, + ) -> MatchResult: + binding_result = self._binding.inspect_finalize_component_match_result( + int(component_type), + result, + _radix_key_buffer(params.key), + params.key.extra_key, + params.key.cache_salt, + value_chunks, + best_value_len, + ) + return _match_result_from_binding(binding_result)._replace( + cache_protected_len=result.cache_protected_len, + cache_actions=result.cache_actions, + ) + + def build_backup_node_ids( + self, node_id: NodeId, write_back: bool = False + ) -> list[NodeId]: + return self._binding.inspect_build_backup_node_ids(node_id, write_back) diff --git a/test/registered/unit/mem_cache/test_rust_tree_core.py b/test/registered/unit/mem_cache/test_rust_tree_core.py new file mode 100644 index 000000000..c74cee1d2 --- /dev/null +++ b/test/registered/unit/mem_cache/test_rust_tree_core.py @@ -0,0 +1,149 @@ +"""Smoke tests for the in-tree Rust TreeCore backend (``rust``). + +Requires a Rust toolchain: the extension builds with cargo on first use. +""" + +import shutil +from array import array + +import pytest +import torch +from unified_tree_core_inspection_interface import UnifiedTreeCoreInspectionInterface + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=90, suite="base-a-test-cpu") + +if shutil.which("cargo") is None: + pytest.skip("the rust backend builds with cargo", allow_module_level=True) + +from sglang.srt.mem_cache.base_prefix_cache import InsertParams, MatchPrefixParams +from sglang.srt.mem_cache.cache_init_params import CacheInitParams +from sglang.srt.mem_cache.radix_cache import RadixKey +from sglang.srt.mem_cache.unified_cache.component_type import ComponentType +from sglang.srt.mem_cache.unified_cache.tree_core_registry import create_tree_core + + +def _tree_core(): + return create_tree_core( + "rust", + CacheInitParams( + disable=False, + req_to_token_pool=None, + token_to_kv_pool_allocator=None, + page_size=1, + tree_components=(ComponentType.FULL,), + ), + components={}, + ) + + +def _key(token_ids, extra_key=None): + return RadixKey(array("q", token_ids), extra_key=extra_key) + + +def _pump_insert(core, params): + step = core.begin_insert(params) + while step.result is None: + step = core.resume_insert() + core.end_insert() + return step.result + + +def test_registry_resolves_the_rust_backend_lazily(): + core = _tree_core() + assert type(core).__name__ == "RustUnifiedTreeCore" + assert not isinstance(core, UnifiedTreeCoreInspectionInterface) + assert not any(name.startswith("inspect_") for name in dir(core._binding)) + + +def test_insert_then_match_round_trips(): + core = _tree_core() + _pump_insert( + core, + InsertParams( + key=_key([1, 2, 3]), value=torch.tensor([10, 11, 12], dtype=torch.int64) + ), + ) + matched = core.match_prefix(MatchPrefixParams(key=_key([1, 2, 3]))) + assert matched.device_indices.tolist() == [10, 11, 12] + + +def test_lock_moves_tokens_between_evictable_and_protected(): + core = _tree_core() + _pump_insert( + core, + InsertParams(key=_key([1, 2]), value=torch.tensor([10, 11], dtype=torch.int64)), + ) + matched = core.match_prefix(MatchPrefixParams(key=_key([1, 2]))) + core.inc_lock_ref(matched.best_match_node) + assert core.protected_size() == 2 + assert core.evictable_size() == 0 + core.dec_lock_ref(matched.best_match_node) + assert core.evictable_size() == 2 + + +def test_namespaces_isolate_the_same_tokens(): + core = _tree_core() + _pump_insert( + core, + InsertParams( + key=_key([1, 2], extra_key="chat"), + value=torch.tensor([20, 21], dtype=torch.int64), + ), + ) + salted = core.match_prefix(MatchPrefixParams(key=_key([1, 2], extra_key="chat"))) + assert salted.device_indices.tolist() == [20, 21] + unsalted = core.match_prefix(MatchPrefixParams(key=_key([1, 2]))) + assert unsalted.device_indices.numel() == 0 + + +def test_backfill_hashes_existing_nodes_in_parent_order(): + expected = _tree_core() + expected.enable_storage = True + _pump_insert( + expected, + InsertParams(key=_key([1, 2]), value=torch.tensor([10, 11], dtype=torch.int64)), + ) + _pump_insert( + expected, + InsertParams( + key=_key([1, 2, 3, 4]), + value=torch.tensor([10, 11, 12, 13], dtype=torch.int64), + ), + ) + + late = _tree_core() + _pump_insert( + late, + InsertParams(key=_key([1, 2]), value=torch.tensor([10, 11], dtype=torch.int64)), + ) + _pump_insert( + late, + InsertParams( + key=_key([1, 2, 3, 4]), + value=torch.tensor([10, 11, 12, 13], dtype=torch.int64), + ), + ) + + parent = late.match_prefix(MatchPrefixParams(key=_key([1, 2]))).best_match_node + child = late.match_prefix(MatchPrefixParams(key=_key([1, 2, 3, 4]))).best_match_node + expected_parent = expected.match_prefix( + MatchPrefixParams(key=_key([1, 2])) + ).best_match_node + expected_child = expected.match_prefix( + MatchPrefixParams(key=_key([1, 2, 3, 4])) + ).best_match_node + + assert late.get_hash_values(parent) == [] + assert late.get_hash_values(child) == [] + assert late.backfill_missing_hash_values() == 2 + assert late.get_hash_values(parent) == expected.get_hash_values(expected_parent) + assert late.get_hash_values(child) == expected.get_hash_values(expected_child) + assert late.backfill_missing_hash_values() == 0 + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/mem_cache/test_rust_tree_core_integration.py b/test/registered/unit/mem_cache/test_rust_tree_core_integration.py new file mode 100644 index 000000000..45056b2ab --- /dev/null +++ b/test/registered/unit/mem_cache/test_rust_tree_core_integration.py @@ -0,0 +1,1973 @@ +"""Integration tests driving the real compiled Rust mem_cache extension.""" + +import hashlib +import shutil +import sys +from array import array +from types import SimpleNamespace + +import pytest +import torch + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=90, suite="base-a-test-cpu") + +if shutil.which("cargo") is None: + pytest.skip("the rust backend builds with cargo", allow_module_level=True) + +from sglang.srt.disaggregation.kv_events import ( + AllBlocksCleared, + BlockRemoved, + BlockStored, + BlockStoredMetadata, + BlockStoredWithMetadata, + StorageMedium, +) +from sglang.srt.environ import envs +from sglang.srt.mem_cache.base_prefix_cache import ( + InsertParams, + InsertResult, + MatchPrefixParams, +) +from sglang.srt.mem_cache.cache_init_params import CacheInitParams +from sglang.srt.mem_cache.hicache_storage import ( + PoolHitPolicy, + PoolName, + PoolTransfer, + PoolTransferResult, +) +from sglang.srt.mem_cache.radix_cache import RadixKey +from sglang.srt.mem_cache.rust_tree_core.adapter import RustUnifiedTreeCore +from sglang.srt.mem_cache.rust_tree_core.extension import bindings as mem_cache +from sglang.srt.mem_cache.unified_cache.cache_action import ( + BackupKV, + FreeComponentHostSlot, + FreeDeviceKV, + FreeDeviceKVFullOnly, + RecoverSWAWithLockedFull, + ReplaceWriteThroughOnNodeSplit, + SWARebuild, +) +from sglang.srt.mem_cache.unified_cache.component_type import ComponentType +from sglang.srt.mem_cache.utils import hash_str_to_int64 +from sglang.srt.runtime_context import get_context + + +def _tree_core(**params_overrides) -> RustUnifiedTreeCore: + params = dict( + disable=False, + req_to_token_pool=None, + token_to_kv_pool_allocator=None, + page_size=1, + tree_components=(ComponentType.FULL,), + ) + params.update(params_overrides) + return RustUnifiedTreeCore(CacheInitParams(**params)) + + +def _key(token_ids: list[int]) -> RadixKey: + return RadixKey(array("q", token_ids)) + + +def _insert(core: RustUnifiedTreeCore, token_ids: list[int], indices: list[int]): + return _pump_insert( + core, + InsertParams( + key=_key(token_ids), + value=torch.tensor(indices, dtype=torch.int64), + ), + ) + + +def _binding(**init_overrides): + return mem_cache.RustUnifiedTreeCoreBinding( + mem_cache.TreeCoreInitParamsBinding(**init_overrides), + [int(ComponentType.FULL)], + ) + + +def _pump_insert(core: RustUnifiedTreeCore, params: InsertParams) -> InsertResult: + """Drive the resumable-insert protocol, folding step actions into the result.""" + step = core.begin_insert(params) + actions = list(step.actions) + while step.result is None: + step = core.resume_insert() + actions.extend(step.actions) + core.end_insert() + return InsertResult( + prefix_len=step.result.prefix_len, + last_device_node=step.result.last_device_node, + mamba_exist=step.result.mamba_exist, + cache_actions=actions, + ) + + +def _accumulate_step(result, tracker, device_frees, host_frees): + """Fold an eviction step into running accumulators (the Controller + consumption contract: deltas add, freed tensors append), draining it.""" + for component, delta in result.tracker.items(): + tracker[component] = tracker.get(component, 0) + delta + for component, tensors in result.device_frees.items(): + device_frees.setdefault(component, []).extend(tensors) + for component, tensors in result.host_frees.items(): + host_frees.setdefault(component, []).extend(tensors) + result.device_frees.clear() + result.host_frees.clear() + + +def test_match_on_the_empty_tree_returns_no_indices(): + core = _tree_core() + result = core.match_prefix(MatchPrefixParams(key=_key([1, 2, 3]))) + assert result.device_indices.numel() == 0 + + +def test_insert_then_match_back_returns_the_exact_indices(): + core = _tree_core() + result = _insert(core, [1, 2, 3], [10, 11, 12]) + assert result.prefix_len == 0 + assert result.cache_actions == [] + matched = core.match_prefix(MatchPrefixParams(key=_key([1, 2, 3]))) + assert result.last_device_node == matched.last_device_node + assert matched.device_indices.tolist() == [10, 11, 12] + + +def test_root_node_handle_is_namespace_independent(): + core = _tree_core() + root = core.root_node_handle() + # The single root serves every namespace, seen or not. + assert core.root_node_handle("ghost") == root + _pump_insert( + core, + InsertParams( + key=RadixKey(array("q", [1, 2]), extra_key="chat"), + value=torch.tensor([10, 11], dtype=torch.int64), + ), + ) + assert core.root_node_handle("chat") == root + # A full miss in the namespace anchors its match at the root. + missed = core.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", [9]), extra_key="chat")) + ) + assert missed.best_match_node == root + + +def test_stale_handle_reads_raise_key_error_without_poisoning_the_core(): + core = _tree_core() + stale_root = core.root_node_handle() + core.reset() + live_root = core.root_node_handle() + + accessors = ( + core.is_backuped, + core.is_root, + core.get_last_hash_value, + core.get_prefix_hash_values, + core.prefetch_anchor_info, + ) + for accessor in accessors: + with pytest.raises(Exception) as exc_info: + accessor(stale_root) + assert isinstance(exc_info.value, KeyError) + assert exc_info.value.args == (stale_root,) + assert core.is_root(live_root) + + +def test_stale_handle_operations_raise_key_error_without_poisoning_the_core(): + from sglang.srt.mem_cache.unified_cache.components import CacheTransferPhase + + core = _tree_core() + stale_root = core.root_node_handle() + core.reset() + live_root = core.root_node_handle() + + operations = ( + lambda: core.demote(stale_root), + lambda: core.build_hicache_transfers( + ComponentType.FULL, stale_root, CacheTransferPhase.BACKUP_STORAGE + ), + lambda: core.build_load_back_spec(stale_root), + lambda: core.get_hash_values(stale_root), + lambda: core.dfs_weight_order([stale_root]), + ) + for operation in operations: + with pytest.raises(KeyError) as exc_info: + operation() + assert exc_info.value.args == (stale_root,) + assert core.is_root(live_root) + + +def test_dfs_weight_order_groups_the_heaviest_subtree_first(): + core = _tree_core() + _insert(core, [1, 10], [10, 11]) + _insert(core, [1, 11], [10, 12]) + _insert(core, [2, 20], [20, 21]) + + branch_a = core.match_prefix(MatchPrefixParams(key=_key([1, 99]))).last_device_node + leaf_a1 = core.match_prefix(MatchPrefixParams(key=_key([1, 10]))).last_device_node + leaf_a2 = core.match_prefix(MatchPrefixParams(key=_key([1, 11]))).last_device_node + leaf_b = core.match_prefix(MatchPrefixParams(key=_key([2, 20]))).last_device_node + + assert core.dfs_weight_order([leaf_b, leaf_a2, leaf_a1, leaf_a1, branch_a]) == [ + 2, + 3, + 1, + 4, + 0, + ] + assert core.dfs_weight_order([leaf_b, leaf_a2]) == [1, 0] + + +def test_get_hash_values_round_trips_through_insert_host(): + core = _tree_core() + root = core.match_prefix(MatchPrefixParams(key=_key([99]))).best_match_node + result = core.insert_host( + root, _key([1]), torch.tensor([100], dtype=torch.int64), ["h0"] + ) + assert core.get_hash_values(result.inserted_host_node) == ["h0"] + # A never-hashed device node reads back empty. + _insert(core, [5], [50]) + device_node = core.match_prefix(MatchPrefixParams(key=_key([5]))).best_match_node + assert core.get_hash_values(device_node) == [] + + +def test_insert_coerces_a_none_priority(): + core = _tree_core() + result = _pump_insert( + core, + InsertParams( + key=_key([1, 2]), + value=torch.tensor([10, 11], dtype=torch.int64), + priority=None, + ), + ) + assert result.prefix_len == 0 + + +def test_extension_insert_frees_the_duplicate_overlap(): + core = _tree_core() + _insert(core, [1, 2, 3], [10, 11, 12]) + result = _insert(core, [1, 2, 3, 4, 5], [20, 21, 22, 13, 14]) + assert result.prefix_len == 3 + # The overlap's fresh indices are duplicates: freed, not stored. + assert len(result.cache_actions) == 1 + action = result.cache_actions[0] + assert isinstance(action, FreeDeviceKV) + assert torch.cat(action.indices).tolist() == [20, 21, 22] + matched = core.match_prefix(MatchPrefixParams(key=_key([1, 2, 3, 4, 5]))) + assert matched.device_indices.tolist() == [10, 11, 12, 13, 14] + + +def test_lock_and_unlock_move_tokens_between_protected_and_evictable(): + core = _tree_core() + _insert(core, [1, 2, 3], [10, 11, 12]) + _insert(core, [1, 2, 3, 4, 5], [20, 21, 22, 13, 14]) + matched = core.match_prefix(MatchPrefixParams(key=_key([1, 2, 3, 4, 5]))) + core.inc_lock_ref(matched.best_match_node) + assert core.protected_size() == 5 + assert core.evictable_size() == 0 + core.dec_lock_ref(matched.best_match_node) + assert core.protected_size() == 0 + assert core.evictable_size() == 5 + + +def test_full_eviction_walk_drains_the_tree(): + core = _tree_core() + _insert(core, [1, 2, 3], [10, 11, 12]) + _insert(core, [1, 2, 3, 4, 5], [20, 21, 22, 13, 14]) + tracker = {ComponentType.FULL: 0} + device_frees: dict = {} + host_frees: dict = {} + core.evict_device_start(ComponentType.FULL, 100) + evicted = 0 + while True: + step = core.evict_device_next_node(ComponentType.FULL, tracker) + node = step.node_id + _accumulate_step(step, tracker, device_frees, host_frees) + if node is None: + break + leaf_step = core.evict_device_leaf(node, is_write_back=False) + _accumulate_step(leaf_step, tracker, device_frees, host_frees) + evicted += 1 + core.evict_device_end(ComponentType.FULL) + assert evicted == 2 + assert tracker == {ComponentType.FULL: 5} + assert core.evictable_size() == 0 + freed = torch.cat(device_frees[ComponentType.FULL]) + assert sorted(freed.tolist()) == [10, 11, 12, 13, 14] + assert host_frees == {} + + +def test_insert_suspends_at_a_backup_barrier_through_the_binding(): + core = _tree_core() + core.write_through_threshold = 2 + core.set_hicache_enabled() + _insert(core, [1, 2, 3], [10, 11, 12]) + core.match_prefix(MatchPrefixParams(key=_key([1, 2, 3]))) + + step = core.begin_insert( + InsertParams( + key=_key([1, 2, 3, 4, 5]), + value=torch.tensor([20, 21, 22, 13, 14], dtype=torch.int64), + ) + ) + # The crossing node's backup is a barrier: the walk stays suspended in Rust. + assert step.result is None + assert core.has_ongoing_insert() + assert [type(a).__name__ for a in step.actions] == ["FreeDeviceKV", "BackupKV"] + + done = core.resume_insert() + assert done.actions == [] + assert done.result is not None + assert done.result.prefix_len == 3 + assert not core.has_ongoing_insert() + assert core.end_insert() == [] + core.sanity_check([], []) + + +def test_configuration_reads_the_locked_rust_state(): + core = _tree_core() + core._binding.set_hicache_enabled() + core._binding.set_is_write_back(True) + core._binding.set_write_through_threshold(7) + core._binding.set_enable_storage(True) + assert core.enable_hicache is True + assert core.is_write_back is True + assert core.write_through_threshold == 7 + assert core.enable_storage is True + + swa_core = _swa_tree_core() + swa_core._binding.set_has_swa_host_pool() + assert swa_core.has_swa_host_pool is True + + +def test_external_cache_linker_is_rejected(): + core = _tree_core() + assert core.enable_external_cache_linker is False + with pytest.raises(ValueError, match="External cache linker"): + core.enable_external_cache_linker = True + assert core.enable_external_cache_linker is False + + +def test_sanity_check_passes_after_the_full_flow(): + core = _tree_core() + _insert(core, [1, 2, 3], [10, 11, 12]) + _insert(core, [1, 2, 3, 4, 5], [20, 21, 22, 13, 14]) + core.match_prefix(MatchPrefixParams(key=_key([1, 2, 3, 4, 5]))) + core.sanity_check([], []) + + +def test_sanity_check_maps_invariant_failures_to_assertion_error(): + core = _tree_core() + _insert(core, [1], [10]) + leaf = core.match_prefix(MatchPrefixParams(key=_key([1]))).best_match_node + + with pytest.raises( + AssertionError, match=r"(?s)Sanity check FAILED.*load_back node 8 lock_ref=0" + ): + core.sanity_check([], [(8, leaf)]) + + # A reported invariant failure does not poison the binding mutex. + core.sanity_check([], []) + + +def test_short_value_tensor_raises_value_error(): + binding = _binding() + params = mem_cache.InsertParamsBinding( + key=array("q", [1, 2, 3]), + value=torch.tensor([10, 11], dtype=torch.int64), + ) + with pytest.raises(ValueError, match="shorter than the aligned key length"): + binding.insert(params) + + +def test_binding_stays_usable_after_a_failed_insert(): + binding = _binding() + with pytest.raises(ValueError): + binding.insert( + mem_cache.InsertParamsBinding( + key=array("q", [1, 2, 3]), + value=torch.tensor([10], dtype=torch.int64), + ) + ) + result = binding.insert( + mem_cache.InsertParamsBinding( + key=array("q", [1, 2, 3]), + value=torch.tensor([10, 11, 12], dtype=torch.int64), + ) + ) + assert result.prefix_len == 0 + matched = binding.match_prefix(mem_cache.MatchParamsBinding(array("q", [1, 2, 3]))) + assert matched.device_indices.tolist() == [10, 11, 12] + + +@pytest.mark.parametrize("prior_hash", ["abcd", "z" * 64]) +def test_hash_boundary_rejects_malformed_prior_hash(prior_hash): + with pytest.raises(ValueError, match="64-character hexadecimal digest"): + mem_cache.get_hash_str(array("q", [1, 2]), prior_hash, 2) + + +@pytest.mark.parametrize("token_id", [-1, 1 << 32]) +def test_hash_boundary_rejects_token_ids_outside_uint32(token_id): + with pytest.raises(ValueError, match="does not fit in uint32"): + mem_cache.get_hash_str(array("q", [token_id]), None, 1) + + +def test_hash_boundary_rejects_zero_page_size(): + with pytest.raises(ValueError, match="page_size must be positive"): + mem_cache.get_hash_str(array("q", [1, 2]), None, 0) + + +def test_binding_rejects_zero_page_size_before_core_construction(): + with pytest.raises(ValueError, match="page_size must be at least 1"): + _binding(page_size=0) + + +def test_binding_rejects_unknown_eviction_policy_before_core_construction(): + with pytest.raises(ValueError, match="Unknown eviction policy: clock"): + _binding(eviction_policy="clock") + + +def test_poisoned_binding_refuses_to_reuse_the_core(): + binding = _binding() + root = binding.root_node_handle() + + # Reading a backup spec from the value-less root deliberately trips a native + # invariant while the binding owns the mutex. + with pytest.raises(BaseException) as initial_panic: + binding.build_backup_spec(root) + assert initial_panic.type.__name__ == "PanicException" + + # The guard must fail closed instead of handing potentially partial state to + # the next operation through PoisonError::into_inner(). + with pytest.raises(BaseException) as poisoned: + binding.root_node_handle() + assert poisoned.type.__name__ == "PanicException" + assert "Rust TreeCore mutex poisoned" in str(poisoned.value) + + +def test_extra_key_isolates_namespaces(): + core = _tree_core() + result = _pump_insert( + core, + InsertParams( + key=RadixKey(array("q", [1, 2, 3]), extra_key="salt"), + value=torch.tensor([10, 11, 12], dtype=torch.int64), + ), + ) + assert result.prefix_len == 0 + salted = core.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", [1, 2, 3]), extra_key="salt")) + ) + assert salted.device_indices.tolist() == [10, 11, 12] + unsalted = core.match_prefix(MatchPrefixParams(key=_key([1, 2, 3]))) + assert unsalted.device_indices.numel() == 0 + other = core.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", [1, 2, 3]), extra_key="other")) + ) + assert other.device_indices.numel() == 0 + assert core.prefetch_anchor_info(salted.best_match_node) == ("salt", None) + assert core.prefetch_anchor_info(core.root_node_handle()) == (None, None) + + +def test_cache_salt_is_supported_by_all_key_entry_points(): + core = _tree_core() + tokens = array("q", [1, 2]) + first_key = RadixKey(tokens, extra_key="bc", cache_salt="a") + second_key = RadixKey(tokens, extra_key="c", cache_salt="ab") + _pump_insert( + core, + InsertParams(key=first_key, value=torch.tensor([10, 11], dtype=torch.int64)), + ) + _pump_insert( + core, + InsertParams(key=second_key, value=torch.tensor([20, 21], dtype=torch.int64)), + ) + + assert core.match_prefix( + MatchPrefixParams(key=first_key) + ).device_indices.tolist() == [ + 10, + 11, + ] + assert core.match_prefix( + MatchPrefixParams(key=second_key) + ).device_indices.tolist() == [ + 20, + 21, + ] + assert ( + core.match_prefix(MatchPrefixParams(key=_key([1, 2]))).device_indices.numel() + == 0 + ) + + host_core = _tree_core() + host_core.set_hicache_enabled() + result = host_core.insert_host( + host_core.root_node_handle(), + first_key, + torch.tensor([100, 101], dtype=torch.int64), + ["h0", "h1"], + ) + assert result.inserted_host_node is not None + host_match = host_core.match_prefix(MatchPrefixParams(key=first_key)) + assert host_match.host_hit_length == 2 + assert host_core.prefetch_anchor_info(host_match.best_match_node) == ("bc", "a") + with pytest.raises(RuntimeError, match="does not match non-root anchor"): + host_core.insert_host( + host_match.best_match_node, + RadixKey(array("q", [3, 4]), extra_key="bc", cache_salt="other"), + torch.tensor([102, 103], dtype=torch.int64), + ["h2", "h3"], + ) + + +def test_session_radix_cache_is_rejected(): + with pytest.raises(ValueError, match="enable-session-radix-cache"): + _tree_core(enable_session_radix_cache=True) + + +@pytest.mark.parametrize( + ("params", "message"), + [ + ( + {"tree_components": (ComponentType.FULL, ComponentType.C128)}, + "components: C128", + ), + ( + {"component_registry_override": {ComponentType.FULL: object}}, + "component_registry_override", + ), + ], +) +def test_unsupported_component_configuration_is_rejected(params, message): + with pytest.raises(ValueError, match=message): + _tree_core(**params) + + +def test_page_size_two_drops_the_ragged_tail(): + core = _tree_core(page_size=2) + result = _pump_insert( + core, + InsertParams( + key=_key([1, 2, 3, 4, 5]), + value=torch.tensor([10, 11, 12, 13, 14], dtype=torch.int64), + ), + ) + assert result.prefix_len == 0 + matched = core.match_prefix(MatchPrefixParams(key=_key([1, 2, 3, 4, 5]))) + assert matched.device_indices.tolist() == [10, 11, 12, 13] + + +def test_insert_value_none_materializes_the_token_ids(): + core = _tree_core() + result = _pump_insert(core, InsertParams(key=_key([1, 2, 3]))) + assert result.prefix_len == 0 + matched = core.match_prefix(MatchPrefixParams(key=_key([1, 2, 3]))) + assert matched.device_indices.tolist() == [1, 2, 3] + + +def test_empty_match_result_is_root_anchored(): + core = _tree_core() + empty = core.empty_match_result + assert empty.device_indices.numel() == 0 + assert empty.host_hit_length == 0 + probe = core.match_prefix(MatchPrefixParams(key=_key([9]))) + assert empty.best_match_node == probe.best_match_node + assert empty.last_device_node == probe.last_device_node + assert empty.last_host_node == probe.last_host_node + + +def test_set_hicache_enabled_marks_the_tree(): + core = _tree_core() + core.set_hicache_enabled() + assert core.enable_hicache + + +def test_hicache_write_through_and_load_back_round_trip(): + core = _tree_core() + core.set_hicache_enabled() + _insert(core, [1, 2], [10, 11]) + leaf = core.match_prefix(MatchPrefixParams(key=_key([1, 2]))).best_match_node + # Write-through: back the leaf up host-side, then demote it to host-only. + device_value, comp_xfers = core.build_backup_spec(leaf) + assert device_value.tolist() == [10, 11] + assert comp_xfers == {} + core.mark_write_through_pending(leaf) + core.commit_backup(leaf, torch.tensor([100, 101], dtype=torch.int64), comp_xfers) + core.finish_write_through([leaf], leaf) + tracker = {ComponentType.FULL: 0} + device_frees, host_frees = {}, {} + _accumulate_step(core.demote(leaf), tracker, device_frees, host_frees) + assert tracker[ComponentType.FULL] == 2 + assert [t.tolist() for t in device_frees[ComponentType.FULL]] == [[10, 11]] + assert core.component_has_host_value_only(leaf, ComponentType.FULL) + # Load back host -> device; the match then serves device indices again. + kv_xfer, comp_xfers = core.build_load_back_spec(leaf) + assert kv_xfer.name == PoolName.KV + assert kv_xfer.host_indices.tolist() == [100, 101] + assert kv_xfer.nodes_to_load == [leaf] + actions = core.commit_load_back( + leaf, torch.tensor([50, 51], dtype=torch.int64), kv_xfer, comp_xfers + ) + assert actions == [] + result = core.match_prefix(MatchPrefixParams(key=_key([1, 2]))) + assert result.device_indices.tolist() == [50, 51] + core.finish_load_back(leaf) + core.sanity_check([], []) + + +def test_invalid_demote_states_raise_assertion_error(): + core = _tree_core() + core.set_hicache_enabled() + _insert(core, [1], [10]) + leaf = core.match_prefix(MatchPrefixParams(key=_key([1]))).best_match_node + + with pytest.raises(AssertionError): + core.demote(leaf) + + core.commit_backup(leaf, torch.tensor([100], dtype=torch.int64), {}) + tracker = {ComponentType.FULL: 0} + _accumulate_step(core.demote(leaf), tracker, {}, {}) + with pytest.raises(AssertionError): + core.demote(leaf) + + +def test_write_through_load_back_is_unpinned_and_refreshes_duplicate_tracking(): + core = _tree_core() + core.set_hicache_enabled() + root = core.root_node_handle() + leaf = core.insert_host( + root, _key([1]), torch.tensor([100], dtype=torch.int64), ["h0"] + ).inserted_host_node + assert leaf is not None + + kv_xfer, comp_xfers = core.build_load_back_spec(leaf) + core.commit_load_back( + leaf, torch.tensor([50], dtype=torch.int64), kv_xfer, comp_xfers + ) + + # Write-through load-back does not pin Full KV against device eviction. + core.evict_device_start(ComponentType.FULL, 1) + candidate = core.evict_device_next_node( + ComponentType.FULL, {ComponentType.FULL: 0} + ).node_id + core.evict_device_end(ComponentType.FULL) + assert candidate == leaf + + # insert_host created no stale duplicate entry, so this checks the ack refresh. + core.finish_load_back(leaf) + core.sanity_check([], []) + + +def test_insert_host_extends_the_backuped_path(): + core = _tree_core() + core.set_hicache_enabled() + _insert(core, [1], [10]) + leaf = core.match_prefix(MatchPrefixParams(key=_key([1]))).best_match_node + core.commit_backup(leaf, torch.tensor([100], dtype=torch.int64), {}) + result = core.insert_host( + leaf, _key([2, 3]), torch.tensor([101, 102], dtype=torch.int64), [] + ) + assert result.prefix_len == 0 + assert result.total_len == 2 + assert result.inserted_host_node is not None + match = core.match_prefix(MatchPrefixParams(key=_key([1, 2, 3]))) + assert match.host_hit_length == 2 + + +def test_insert_host_reports_a_dropped_write_through_suffix(): + core = _tree_core() + _insert(core, [1], [10]) + parent = core.match_prefix(MatchPrefixParams(key=_key([1]))).best_match_node + + result = core.insert_host( + parent, _key([2]), torch.tensor([100], dtype=torch.int64), ["h0"] + ) + + assert result.prefix_len == 0 + assert result.total_len == 1 + assert result.inserted_host_node is None + assert result.host_insert_dropped + + +def test_host_lock_refs_round_trip(): + core = _tree_core() + core.set_hicache_enabled() + _insert(core, [1], [10]) + leaf = core.match_prefix(MatchPrefixParams(key=_key([1]))).best_match_node + core.commit_backup(leaf, torch.tensor([100], dtype=torch.int64), {}) + core.inc_host_lock_ref(leaf) + core.dec_host_lock_ref(leaf) + core.sanity_check([], []) + + +def test_drive_host_eviction_frees_the_demoted_leaf(): + core = _tree_core() + core.set_hicache_enabled() + _insert(core, [1], [10]) + leaf = core.match_prefix(MatchPrefixParams(key=_key([1]))).best_match_node + core.commit_backup(leaf, torch.tensor([100], dtype=torch.int64), {}) + _accumulate_step(core.demote(leaf), {}, {}, {}) + tracker = {ComponentType.FULL: 0} + device_frees, host_frees = {}, {} + _accumulate_step( + core.drive_host_eviction(ComponentType.FULL, 1), + tracker, + device_frees, + host_frees, + ) + assert tracker[ComponentType.FULL] == 1 + assert [t.tolist() for t in host_frees[ComponentType.FULL]] == [[100]] + # The host-only leaf is gone: the key no longer matches anywhere. + result = core.match_prefix(MatchPrefixParams(key=_key([1]))) + assert result.host_hit_length == 0 + core.sanity_check([], []) + + +def test_events_disabled_take_events_is_empty(): + core = _tree_core() + _insert(core, [1, 2], [10, 11]) + assert core.take_events() == [] + + +def test_insert_emits_block_stored_events(): + core = _tree_core(enable_kv_cache_events=True, page_size=2) + _insert(core, [1, 2, 7, 8], [10, 11, 12, 13]) + hashes = [ + hash_str_to_int64(h) + for h in mem_cache.get_hash_str(array("q", [1, 2, 7, 8]), None, 2) + ] + assert core.take_events() == [ + BlockStored( + block_hashes=hashes, + parent_block_hash=None, + token_ids=[1, 2, 7, 8], + block_size=2, + lora_id=None, + medium=StorageMedium.GPU, + ), + ] + assert core.take_events() == [] + + +def test_salted_events_match_python_hash_and_metadata_contract(): + core = _tree_core(enable_kv_cache_events=True, page_size=2) + key = RadixKey(array("q", [1, 2, 7, 8]), cache_salt="tenant-a") + _pump_insert( + core, + InsertParams( + key=key, + value=torch.tensor([10, 11, 12, 13], dtype=torch.int64), + ), + ) + seed = hashlib.sha256(b"sglang-cache-salt-v1\0tenant-a").hexdigest() + hashes = [ + hash_str_to_int64(value) + for value in mem_cache.get_hash_str(array("q", [1, 2, 7, 8]), seed, 2) + ] + assert core.take_events() == [ + BlockStoredWithMetadata( + block_hashes=hashes, + parent_block_hash=None, + token_ids=[1, 2, 7, 8], + block_size=2, + lora_id=None, + medium=StorageMedium.GPU, + metadata=BlockStoredMetadata(cache_salt="tenant-a"), + ) + ] + + tracker = {ComponentType.FULL: 0} + core.evict_device_start(ComponentType.FULL, 4) + candidate = core.evict_device_next_node(ComponentType.FULL, tracker).node_id + assert candidate is not None + evicted = core.evict_device_leaf(candidate, is_write_back=False) + evicted.device_frees.clear() + evicted.host_frees.clear() + core.evict_device_end(ComponentType.FULL) + assert core.take_events() == [ + BlockRemoved(block_hashes=hashes, medium=StorageMedium.GPU) + ] + + +def test_salted_eagle_events_match_the_bigram_hash_contract(): + core = _tree_core(enable_kv_cache_events=True, page_size=2, is_eagle=True) + raw_tokens = array("q", [1, 2, 3, 4, 5]) + key = RadixKey(raw_tokens, cache_salt="tenant-a", is_bigram=True) + _pump_insert( + core, + InsertParams( + key=key, + value=torch.tensor([10, 11, 12, 13], dtype=torch.int64), + ), + ) + seed = hashlib.sha256(b"sglang-cache-salt-v1\0tenant-a").hexdigest() + hashes = [ + hash_str_to_int64(value) + for value in mem_cache.get_hash_str(raw_tokens, seed, 2, is_bigram=True) + ] + assert core.take_events() == [ + BlockStoredWithMetadata( + block_hashes=hashes, + parent_block_hash=None, + token_ids=[(1, 2), (2, 3), (3, 4), (4, 5)], + block_size=2, + lora_id=None, + medium=StorageMedium.GPU, + metadata=BlockStoredMetadata(cache_salt="tenant-a"), + ) + ] + + +def test_demote_emits_block_removed(): + core = _tree_core(enable_kv_cache_events=True) + core.set_hicache_enabled() + _insert(core, [1, 2], [10, 11]) + leaf = core.match_prefix(MatchPrefixParams(key=_key([1, 2]))).best_match_node + core.commit_backup(leaf, torch.tensor([100, 101], dtype=torch.int64), {}) + core.take_events() + _accumulate_step(core.demote(leaf), {}, {}, {}) + hashes = [ + hash_str_to_int64(h) + for h in mem_cache.get_hash_str(array("q", [1, 2]), None, 1) + ] + assert core.take_events() == [ + BlockRemoved(block_hashes=hashes, medium=StorageMedium.GPU) + ] + + +def test_all_cleared_event_crosses_the_binding(): + core = _tree_core(enable_kv_cache_events=True) + core._record_all_cleared_event() + assert core.take_events() == [AllBlocksCleared()] + + +def test_match_result_mamba_fields_are_inert_without_mamba(): + core = _tree_core() + _insert(core, [1, 2], [10, 11]) + result = core.match_prefix(MatchPrefixParams(key=_key([1, 2]))) + assert result.mamba_host_hit_length == 0 + assert result.mamba_branching_seqlen is None + + +def test_storage_backup_spec_round_trips_the_backuped_node(): + core = _tree_core(page_size=2) + core.set_hicache_enabled() + core.enable_storage = True + _insert(core, [1, 2], [10, 11]) + _insert(core, [1, 2, 7, 8], [10, 11, 12, 13]) + parent = core.match_prefix(MatchPrefixParams(key=_key([1, 2]))).best_match_node + child = core.match_prefix(MatchPrefixParams(key=_key([1, 2, 7, 8]))).best_match_node + core.commit_backup(parent, torch.tensor([100, 101], dtype=torch.int64), {}) + core.commit_backup(child, torch.tensor([102, 103], dtype=torch.int64), {}) + + spec = core.build_storage_backup_spec(child, pass_prefix_keys=True) + assert spec.host_value.tolist() == [102, 103] + assert spec.token_ids == array("q", [7, 8]) + parent_hashes = mem_cache.get_hash_str(array("q", [1, 2]), None, 2) + assert spec.prefix_keys == parent_hashes + assert spec.hash_value == mem_cache.get_hash_str( + array("q", [7, 8]), parent_hashes[-1], 2 + ) + assert spec.comp_xfers == {} + + +def test_prefetch_node_accessors_round_trip(): + core = _tree_core(page_size=2) + core.set_hicache_enabled() + core.enable_storage = True + _insert(core, [1, 2], [10, 11]) + leaf = core.match_prefix(MatchPrefixParams(key=_key([1, 2]))).best_match_node + + assert not core.is_backuped(leaf) + assert not core.is_root(leaf) + assert core.get_last_hash_value(leaf) == ( + mem_cache.get_hash_str(array("q", [1, 2]), None, 2)[-1] + ) + assert core.get_prefix_hash_values(leaf) == [] + + core.commit_backup(leaf, torch.tensor([100, 101], dtype=torch.int64), {}) + assert core.is_backuped(leaf) + + +def test_storage_backup_spec_is_none_for_an_unbackuped_node(): + core = _tree_core() + _insert(core, [1, 2], [10, 11]) + leaf = core.match_prefix(MatchPrefixParams(key=_key([1, 2]))).best_match_node + assert core.build_storage_backup_spec(leaf, pass_prefix_keys=False) is None + + +def test_build_hicache_transfers_routes_the_backup_storage_phase(): + from sglang.srt.mem_cache.unified_cache.components import CacheTransferPhase + + core = _tree_core() + _insert(core, [1, 2], [10, 11]) + leaf = core.match_prefix(MatchPrefixParams(key=_key([1, 2]))).best_match_node + # The FULL component has no storage sidecar; the kv transfer is built by + # the controller from the spec instead. + assert ( + core.build_hicache_transfers( + ComponentType.FULL, leaf, CacheTransferPhase.BACKUP_STORAGE + ) + is None + ) + + +def _canary_rows(core, *, unlocked_only=False, swa_resident_only=False): + walk = core.walk_for_kv_canary( + unlocked_only=unlocked_only, swa_resident_only=swa_resident_only + ) + return sorted( + zip( + walk.slot_indices.tolist(), + walk.positions.tolist(), + walk.prev_slot_indices.tolist(), + ) + ) + + +def test_walk_for_kv_canary_emits_chained_rows(): + core = _tree_core() + _insert(core, [1, 2, 3], [10, 11, 12]) + _insert(core, [1, 2, 3, 4, 5], [10, 11, 12, 13, 14]) + assert _canary_rows(core) == [ + (10, 0, -1), + (11, 1, 10), + (12, 2, 11), + (13, 3, 12), + (14, 4, 13), + ] + + +def test_walk_for_kv_canary_unlocked_only_skips_locked_nodes_but_keeps_the_chain(): + core = _tree_core() + _insert(core, [1, 2, 3], [10, 11, 12]) + _insert(core, [1, 2, 3, 4, 5], [10, 11, 12, 13, 14]) + locked = core.match_prefix(MatchPrefixParams(key=_key([1, 2, 3]))).best_match_node + core.inc_lock_ref(locked) + assert _canary_rows(core, unlocked_only=True) == [(13, 3, 12), (14, 4, 13)] + + +def test_walk_for_kv_canary_skips_demoted_nodes(): + core = _tree_core() + core.set_hicache_enabled() + _insert(core, [1, 2], [10, 11]) + leaf = core.match_prefix(MatchPrefixParams(key=_key([1, 2]))).best_match_node + core.commit_backup(leaf, torch.tensor([100, 101], dtype=torch.int64), {}) + _accumulate_step(core.demote(leaf), {}, {}, {}) + assert _canary_rows(core) == [] + + +def test_walk_for_kv_canary_swa_filter_is_inert_without_the_swa_component(): + core = _tree_core() + _insert(core, [1, 2], [10, 11]) + assert _canary_rows(core, swa_resident_only=True) == [(10, 0, -1), (11, 1, 10)] + + +def test_empty_keys_cross_the_binding(): + assert mem_cache.MatchParamsBinding(array("q")).key == [] + assert mem_cache.MatchParamsBinding([]).key == [] + + +def test_empty_cache_salt_uses_the_default_namespace_at_the_binding(): + binding = _binding() + binding.insert( + mem_cache.InsertParamsBinding( + key=array("q", [1]), + value=torch.tensor([10], dtype=torch.int64), + cache_salt="", + ) + ) + result = binding.match_prefix(mem_cache.MatchParamsBinding(array("q", [1]))) + assert result.device_indices.tolist() == [10] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_cuda_core_resolves_the_current_device(): + core = _tree_core( + token_to_kv_pool_allocator=SimpleNamespace(device="cuda"), + ) + assert core.device == torch.device("cuda", torch.cuda.current_device()) + result = _pump_insert( + core, + InsertParams( + key=_key([1, 2, 3]), + value=torch.tensor([10, 11, 12], dtype=torch.int64, device=core.device), + ), + ) + assert result.prefix_len == 0 + matched = core.match_prefix(MatchPrefixParams(key=_key([1, 2, 3]))) + assert matched.device_indices.device == core.device + assert matched.device_indices.tolist() == [10, 11, 12] + # The value=None fallback also lands on the resolved device. + fallback = _pump_insert(core, InsertParams(key=_key([7, 8]))) + assert fallback.prefix_len == 0 + matched = core.match_prefix(MatchPrefixParams(key=_key([7, 8]))) + assert matched.device_indices.tolist() == [7, 8] + + +def test_unsupported_component_sets_are_rejected(): + with pytest.raises(ValueError, match="component sets are supported"): + mem_cache.RustUnifiedTreeCoreBinding( + mem_cache.TreeCoreInitParamsBinding(), [int(ComponentType.SWA)] + ) + with pytest.raises(ValueError, match="component sets are supported"): + mem_cache.RustUnifiedTreeCoreBinding( + mem_cache.TreeCoreInitParamsBinding(), [int(ComponentType.MAMBA)] + ) + + +def test_swa_requires_the_sliding_window_size(): + with pytest.raises(ValueError, match="requires swa_sliding_window_size"): + mem_cache.RustUnifiedTreeCoreBinding( + mem_cache.TreeCoreInitParamsBinding(), + [int(ComponentType.FULL), int(ComponentType.SWA)], + ) + + +def test_swa_without_a_window_is_rejected_through_the_adapter(): + with pytest.raises(ValueError, match="requires swa_sliding_window_size"): + _tree_core(tree_components=(ComponentType.FULL, ComponentType.SWA)) + + +def test_enable_hicache_constructs(): + mem_cache.RustUnifiedTreeCoreBinding( + mem_cache.TreeCoreInitParamsBinding(enable_hicache=True), + [int(ComponentType.FULL)], + ) + + +def test_is_write_back_constructs(): + mem_cache.RustUnifiedTreeCoreBinding( + mem_cache.TreeCoreInitParamsBinding(is_write_back=True), + [int(ComponentType.FULL)], + ) + + +def test_write_back_eviction_backs_up_then_drop_subtree_falls_back(): + core = _tree_core() + core.is_write_back = True + core.set_hicache_enabled() + _insert(core, [1, 2], [10, 11]) + leaf = core.match_prefix(MatchPrefixParams(key=_key([1, 2]))).best_match_node + tracker = {ComponentType.FULL: 0} + device_frees, host_frees = {}, {} + # The unbacked leaf earns a backup action; nothing is freed yet. + leaf_step = core.evict_device_leaf(leaf, is_write_back=True) + backup = leaf_step.backup_kv + _accumulate_step(leaf_step, tracker, device_frees, host_frees) + assert backup == BackupKV([leaf]) + assert device_frees == {} and host_frees == {} + # Host pressure: the backup failed, so the subtree drop keeps eviction moving. + drop_step = core.drop_subtree_no_host(leaf) + dropped = drop_step.is_dropped + _accumulate_step(drop_step, tracker, device_frees, host_frees) + assert dropped + assert tracker[ComponentType.FULL] == 2 + assert [t.tolist() for t in device_frees[ComponentType.FULL]] == [[10, 11]] + result = core.match_prefix(MatchPrefixParams(key=_key([1, 2]))) + assert result.device_indices.numel() == 0 + core.sanity_check([], []) + + +# ==== SWA wiring ==== + + +def _swa_tree_core(window: int = 8, **params_overrides) -> RustUnifiedTreeCore: + return _tree_core( + tree_components=(ComponentType.FULL, ComponentType.SWA), + sliding_window_size=window, + **params_overrides, + ) + + +def test_write_back_load_back_ignores_auxiliary_nodes_for_pending_ownership(): + core = _swa_tree_core(window=4) + core.set_hicache_enabled() + core.has_swa_host_pool = True + core.is_write_back = True + root = core.root_node_handle() + shared = core.insert_host( + root, _key([1]), torch.tensor([100], dtype=torch.int64), ["h0"] + ).inserted_host_node + anchor = core.insert_host( + root, + _key([1, 2]), + torch.tensor([100, 101], dtype=torch.int64), + ["h0", "h1"], + ).inserted_host_node + assert shared is not None and anchor is not None + + core.commit_backup( + shared, + torch.empty(0, dtype=torch.int64), + { + ComponentType.SWA: [ + PoolTransfer( + name=PoolName.SWA, + host_indices=torch.tensor([200], dtype=torch.int64), + ) + ] + }, + ) + core.commit_load_back( + shared, + torch.tensor([10], dtype=torch.int64), + PoolTransfer( + name=PoolName.KV, + host_indices=torch.tensor([100], dtype=torch.int64), + nodes_to_load=[shared], + ), + {}, + ) + + # The first Full load is genuinely pinned while awaiting its own ack. + core.evict_device_start(ComponentType.FULL, 1) + candidate = core.evict_device_next_node( + ComponentType.FULL, {ComponentType.FULL: 0} + ).node_id + core.evict_device_end(ComponentType.FULL) + assert candidate is None + + # Loading shared's SWA under another anchor must not claim its Full pin. + core.commit_load_back( + anchor, + torch.tensor([11], dtype=torch.int64), + PoolTransfer( + name=PoolName.KV, + host_indices=torch.tensor([101], dtype=torch.int64), + nodes_to_load=[anchor], + ), + { + ComponentType.SWA: [ + PoolTransfer( + name=PoolName.SWA, + host_indices=torch.tensor([200], dtype=torch.int64), + device_indices=torch.tensor([20], dtype=torch.int64), + nodes_to_load=[shared], + ) + ] + }, + ) + assert core.get_component_device_value(shared, ComponentType.SWA).tolist() == [20] + + core.finish_load_back(anchor) + core.finish_load_back(shared) + core.sanity_check([], []) + + +def _swa_cache(window: int = 8, page_size: int = 1): + """A real UnifiedRadixCache on the Rust tree core with a real SWA allocator.""" + from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator + from sglang.srt.mem_cache.memory_pool import ReqToTokenPool + from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool + from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache + from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler + + set_global_server_args_for_scheduler( + ServerArgs(model_path="dummy", page_size=page_size) + ) + req_to_token_pool = ReqToTokenPool( + size=4, max_context_len=64, device="cpu", enable_memory_saver=False + ) + kv_pool = SWAKVPool( + size=64, + size_swa=64, + page_size=page_size, + dtype=torch.bfloat16, + head_num=1, + head_dim=8, + swa_attention_layer_ids=[0], + full_attention_layer_ids=[1], + device="cpu", + ) + allocator = SWATokenToKVPoolAllocator( + size=64, + size_swa=64, + page_size=page_size, + dtype=torch.bfloat16, + device="cpu", + kvcache=kv_pool, + need_sort=False, + ) + with envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.override("rust"): + cache = UnifiedRadixCache( + params=CacheInitParams( + req_to_token_pool=req_to_token_pool, + token_to_kv_pool_allocator=allocator, + page_size=page_size, + disable=False, + sliding_window_size=window, + tree_components=(ComponentType.FULL, ComponentType.SWA), + ) + ) + return cache, allocator + + +def test_buffer_backup_snapshot_round_trips_and_detects_a_split(): + core = _tree_core() + core.enable_storage = True + key = RadixKey(array("q", [1, 2]), extra_key="adapter-a", cache_salt="tenant-a") + inserted = _pump_insert( + core, + InsertParams(key=key, value=torch.tensor([10, 11], dtype=torch.int64)), + ) + leaf = inserted.last_device_node + + snapshot = core.snapshot_buffer_backup(leaf, pass_prefix_keys=True) + assert snapshot.node_id == leaf + assert snapshot.parent_is_root + assert snapshot.key.token_ids == array("q", [1, 2]) + assert snapshot.key.extra_key == "adapter-a" + assert snapshot.key.cache_salt == "tenant-a" + assert not snapshot.key.is_bigram + assert snapshot.prefix_keys == [] + assert core.validate_buffer_backup(leaf, len(snapshot.key)) is not None + + _pump_insert( + core, + InsertParams( + key=RadixKey( + array("q", [1, 9]), extra_key="adapter-a", cache_salt="tenant-a" + ), + value=torch.tensor([12, 13], dtype=torch.int64), + ), + ) + assert core.validate_buffer_backup(leaf, len(snapshot.key)) is None + + +def test_buffer_backup_snapshot_preserves_bigram_keys(): + core = _tree_core(is_eagle=True) + core.enable_storage = True + inserted = _insert(core, [1, 2, 3], [10, 11]) + + snapshot = core.snapshot_buffer_backup( + inserted.last_device_node, pass_prefix_keys=False + ) + assert snapshot.key.token_ids == array("q", [1, 2, 3]) + assert snapshot.key.is_bigram + + +def test_swa_core_builds_with_a_window(): + core = _swa_tree_core(window=8) + result = _insert(core, [1, 2, 3], [10, 11, 12]) + assert result.prefix_len == 0 + # The in-window new leaf asks for one SWA rebuild. + (action,) = result.cache_actions + assert isinstance(action, SWARebuild) + assert action.source_value.tolist() == [10, 11, 12] + + +def test_swa_load_back_missing_value_raises_assertion_error(): + from sglang.srt.mem_cache.unified_cache.components import CacheTransferPhase + + core = _swa_tree_core(window=4) + core.set_hicache_enabled() + core.has_swa_host_pool = True + inserted = _insert(core, [1], [10]) + node = inserted.cache_actions[0].node_id + + with pytest.raises(AssertionError): + core.build_hicache_transfers( + ComponentType.SWA, node, CacheTransferPhase.LOAD_BACK + ) + with pytest.raises(AssertionError): + core.build_load_back_spec(node) + + +def test_swa_straddling_insert_crosses_the_boundary_actions(): + core = _swa_tree_core(window=8) + _insert(core, [1, 2, 3, 4], [10, 11, 12, 13]) + result = _pump_insert( + core, + InsertParams( + key=_key([1, 2, 3, 4]), + value=torch.tensor([20, 21, 22, 23], dtype=torch.int64), + swa_evicted_seqlen=2, + ), + ) + free_tail, rebuild, free_duplicates = result.cache_actions + assert isinstance(free_tail, FreeDeviceKVFullOnly) + assert free_tail.indices[0].tolist() == [12, 13] + assert isinstance(rebuild, SWARebuild) + assert rebuild.source_value.tolist() == [22, 23] + assert isinstance(free_duplicates, FreeDeviceKV) + + +def test_every_pool_name_crosses_the_prefetch_commit_boundary(): + from sglang.srt.mem_cache.unified_cache.components import CacheTransferPhase + + core = _tree_core() + core.set_hicache_enabled() + anchor = core.match_prefix(MatchPrefixParams(key=_key([99]))).best_match_node + # Sidecar pools (e.g. the EAGLE draft KV) report hit pages through the + # commit's pool_storage_result; every python pool name must parse. + for name in PoolName: + core.commit_hicache_transfers( + anchor, + CacheTransferPhase.PREFETCH, + {}, + cache_actions=[], + pool_storage_result=PoolTransferResult( + kv_hit_pages=0, extra_pool_hit_pages={name: 1} + ), + ) + + +def _mamba_tree_core( + page_size: int = 1, + mamba_max_states_per_path: int = -1, + **params_overrides, +) -> RustUnifiedTreeCore: + with get_context().override_server_args( + _mamba_cache_chunk_size=256, + mamba_max_states_per_path=mamba_max_states_per_path, + ): + return _tree_core( + tree_components=(ComponentType.FULL, ComponentType.MAMBA), + page_size=page_size, + **params_overrides, + ) + + +def _mamba_tree_core_with_cap(cap: int) -> RustUnifiedTreeCore: + return _mamba_tree_core(mamba_max_states_per_path=cap) + + +def _hybrid_swa_mamba_tree_core(window: int) -> RustUnifiedTreeCore: + with get_context().override_server_args( + _mamba_cache_chunk_size=256, + mamba_max_states_per_path=-1, + ): + return _tree_core( + tree_components=( + ComponentType.FULL, + ComponentType.SWA, + ComponentType.MAMBA, + ), + sliding_window_size=window, + ) + + +def _mamba_insert(core, token_ids, indices, mamba_slot): + return _pump_insert( + core, + InsertParams( + key=_key(token_ids), + value=torch.tensor(indices, dtype=torch.int64), + mamba_value=torch.tensor([mamba_slot], dtype=torch.int64), + ), + ) + + +def test_kv_canary_rows_exclude_mamba_slots(): + core = _mamba_tree_core() + _mamba_insert(core, [1, 2], [10, 11], 7) + # The canary walk emits FULL slots only; the mamba state slot never appears. + assert _canary_rows(core) == [(10, 0, -1), (11, 1, 10)] + + +def test_component_set_guard_accepts_the_mamba_set(): + mem_cache.RustUnifiedTreeCoreBinding( + mem_cache.TreeCoreInitParamsBinding(mamba_cache_chunk_size=256), + [int(ComponentType.FULL), int(ComponentType.MAMBA)], + ) + + +def test_component_set_guard_accepts_the_hybrid_swa_mamba_set(): + mem_cache.RustUnifiedTreeCoreBinding( + mem_cache.TreeCoreInitParamsBinding( + swa_sliding_window_size=8, mamba_cache_chunk_size=256 + ), + [ + int(ComponentType.FULL), + int(ComponentType.SWA), + int(ComponentType.MAMBA), + ], + ) + + +def test_skipped_mamba_lock_survives_swa_only_release_through_the_adapter(): + core = _hybrid_swa_mamba_tree_core(window=2) + inserted = _mamba_insert(core, [1, 2], [10, 11], 7) + for action in inserted.cache_actions: + if isinstance(action, SWARebuild): + core.set_component_device_value( + action.node_id, ComponentType.SWA, action.source_value + ) + node = core.match_prefix(MatchPrefixParams(key=_key([1, 2]))).best_match_node + + owner = core.inc_lock_ref(node) + skipped = core.inc_lock_ref(node, skip_lock_components=(ComponentType.MAMBA,)) + assert skipped.skip_lock_node_ids == {ComponentType.MAMBA: {node}} + assert core.mamba_protected_size() == 1 + + released = core.dec_swa_lock_only( + node, + skipped.swa_uuid_for_lock, + skip_lock_node_ids=skipped.skip_lock_node_ids, + ) + assert dict(released.device_frees) == {} + assert dict(released.host_frees) == {} + assert core.mamba_protected_size() == 1 + + core.dec_lock_ref(node, skipped.to_dec_params(), skip_swa=True) + core.dec_lock_ref(node, owner.to_dec_params()) + assert core.protected_size() == 0 + assert core.swa_protected_size() == 0 + assert core.mamba_protected_size() == 0 + + +def test_component_set_guard_still_rejects_invalid_sets(): + for components in ( + [ComponentType.MAMBA], + [ComponentType.SWA, ComponentType.MAMBA], + [ComponentType.MAMBA, ComponentType.FULL], + ): + with pytest.raises(ValueError, match="component sets"): + mem_cache.RustUnifiedTreeCoreBinding( + mem_cache.TreeCoreInitParamsBinding( + swa_sliding_window_size=8, mamba_cache_chunk_size=256 + ), + [int(component) for component in components], + ) + + +def test_mamba_requires_the_chunk_size_through_the_binding(): + with pytest.raises(ValueError, match="requires mamba_cache_chunk_size"): + mem_cache.RustUnifiedTreeCoreBinding( + mem_cache.TreeCoreInitParamsBinding(), + [int(ComponentType.FULL), int(ComponentType.MAMBA)], + ) + + +def test_mamba_tree_round_trips_through_the_adapter(): + core = _mamba_tree_core() + result = _mamba_insert(core, [1, 2], [10, 11], 7) + assert result.prefix_len == 0 + assert not result.mamba_exist + + matched = core.match_prefix(MatchPrefixParams(key=_key([1, 2]))) + assert matched.device_indices.tolist() == [10, 11] + assert matched.mamba_host_hit_length == 0 + assert core.mamba_evictable_size() == 1 + assert core.all_mamba_values_flatten().tolist() == [7] + + # A reinsert keeps the slot and flags the caller to free the donation. + result = _mamba_insert(core, [1, 2], [10, 11], 8) + assert result.mamba_exist + assert core.all_mamba_values_flatten().tolist() == [7] + + lock = core.inc_lock_ref(matched.best_match_node) + assert core.mamba_protected_size() == 1 + assert core.mamba_evictable_size() == 0 + core.dec_lock_ref(matched.best_match_node, lock.to_dec_params()) + assert core.mamba_protected_size() == 0 + core.sanity_check([], []) + + +def test_mamba_eviction_walk_frees_slots_through_the_adapter(): + core = _mamba_tree_core() + _mamba_insert(core, [1], [10], 7) + internal = core.match_prefix(MatchPrefixParams(key=_key([1]))).best_match_node + _mamba_insert(core, [1, 2], [10, 11], 8) + tracker = {ComponentType.MAMBA: 0} + device_frees: dict = {} + host_frees: dict = {} + core.evict_device_start(ComponentType.MAMBA, 2) + step = core.evict_device_next_node(ComponentType.MAMBA, tracker) + assert step.node_id is None + assert step.made_progress + _accumulate_step(step, tracker, device_frees, host_frees) + + step = core.evict_device_next_node(ComponentType.MAMBA, tracker) + leaf = step.node_id + _accumulate_step(step, tracker, device_frees, host_frees) + assert leaf is not None + core.evict_device_end(ComponentType.MAMBA) + assert tracker[ComponentType.MAMBA] == 1 + assert torch.cat(device_frees[ComponentType.MAMBA]).tolist() == [7] + assert core.mamba_evictable_size() == 1 + + # A pre-eviction node handle locked after the tombstoning lands in the + # skip map, and the replay keeps the release off it. + lock = core.inc_lock_ref(internal) + assert internal in lock.skip_lock_node_ids[ComponentType.MAMBA] + core.dec_lock_ref(internal, lock.to_dec_params()) + core.sanity_check([], []) + + +def test_mamba_path_cap_evicts_excess_states_through_the_adapter(): + from collections import defaultdict + + from sglang.srt.mem_cache.unified_cache.cache_action import ( + MambaEvictExcessPathStates, + ) + + core = _mamba_tree_core_with_cap(1) + _mamba_insert(core, [1], [10], 7) + _mamba_insert(core, [1, 2], [10, 11], 8) + result = _mamba_insert(core, [1, 2, 3], [10, 11, 12], 9) + (action,) = [ + a for a in result.cache_actions if isinstance(a, MambaEvictExcessPathStates) + ] + device_frees, host_frees = defaultdict(list), defaultdict(list) + core.evict_excess_path_states(action.tail_node_id, device_frees, host_frees) + # The two shallow states free; the tail's survives the soft cap. + assert sorted(t.item() for t in device_frees[ComponentType.MAMBA]) == [7, 8] + assert host_frees == {} + assert core.mamba_evictable_size() == 1 + + +def test_eagle_with_mamba_falls_back_to_the_unigram_binding(): + core = _mamba_tree_core(is_eagle=True) + assert core.is_eagle is False + assert type(core._binding) is mem_cache.RustUnifiedTreeCoreBinding + + +def test_mamba_prefetch_commit_round_trips_through_the_adapter(): + from sglang.srt.mem_cache.unified_cache.components import CacheTransferPhase + + core = _mamba_tree_core() + core.set_hicache_enabled() + root = core.match_prefix(MatchPrefixParams(key=_key([99]))).best_match_node + insert_result = core.insert_host( + root, _key([1]), torch.tensor([100], dtype=torch.int64), ["h0"] + ) + + def commit(host_indices, loaded_pages): + actions = [] + core.commit_hicache_transfers( + root, + CacheTransferPhase.PREFETCH, + { + ComponentType.MAMBA: [ + PoolTransfer( + name=PoolName.MAMBA, + host_indices=torch.tensor(host_indices, dtype=torch.int64), + ) + ] + }, + cache_actions=actions, + insert_result=insert_result, + pool_storage_result=PoolTransferResult( + kv_hit_pages=1, extra_pool_hit_pages={PoolName.MAMBA: loaded_pages} + ), + ) + return actions + + # The loaded buffer attaches to the inserted node. + assert commit([50], loaded_pages=1) == [] + assert not insert_result.mamba_exist + + # A second buffer cannot attach: it frees and flags the caller. + (free,) = commit([51], loaded_pages=1) + assert isinstance(free, FreeComponentHostSlot) + assert free.host_indices[0].tolist() == [51] + assert insert_result.mamba_exist + + # The hosted slot now publishes to storage keyed by the trailing hash. + (xfer,) = core.build_hicache_transfers( + ComponentType.MAMBA, root_child(core), CacheTransferPhase.BACKUP_STORAGE + ) + assert xfer.keys == ["h0"] + assert xfer.hit_policy == PoolHitPolicy.TRAILING_PAGES + + +def root_child(core): + """The single inserted node under the default root.""" + return core.match_prefix(MatchPrefixParams(key=_key([1]))).last_host_node + + +def test_split_of_a_write_through_pending_node_crosses_the_replace_action(): + core = _tree_core() + core.set_hicache_enabled() + _insert(core, [1, 2, 3, 4], [10, 11, 12, 13]) + leaf = core.match_prefix(MatchPrefixParams(key=_key([1, 2, 3, 4]))).best_match_node + core.mark_write_through_pending(leaf) + # A divergent prefix splits the pending node; the publish list must follow. + result = _insert(core, [1, 2], [10, 11]) + (replace,) = [ + action + for action in result.cache_actions + if isinstance(action, ReplaceWriteThroughOnNodeSplit) + ] + assert replace.ack_id == leaf + assert replace.old_node_id == leaf + assert replace.new_child_node_id == leaf + assert replace.new_node_id != leaf + + +def test_write_through_threshold_assignment_reaches_the_core(): + core = _tree_core() + core.set_hicache_enabled() + # HiCache init lowers the threshold after construction; the second hit on + # the same prefix must then emit the write-through backup. + core.write_through_threshold = 2 + assert _insert(core, [1, 2], [10, 11]).cache_actions == [] + result = _insert(core, [1, 2], [10, 11]) + assert any(isinstance(action, BackupKV) for action in result.cache_actions) + + +def test_swa_prefetch_commit_end_to_end(): + from sglang.srt.mem_cache.unified_cache.components import CacheTransferPhase + + core = _swa_tree_core(window=4) + core.set_hicache_enabled() + core.has_swa_host_pool = True + anchor = core.match_prefix(MatchPrefixParams(key=_key([99]))).best_match_node + + # The build wraps the host buffer with placeholder keys, trailing-pages policy. + (xfer,) = core.build_hicache_transfers( + ComponentType.SWA, + anchor, + CacheTransferPhase.PREFETCH, + host_indices=torch.tensor([30, 31], dtype=torch.int64), + ) + assert xfer.name == PoolName.SWA + assert xfer.keys == ["__placeholder__", "__placeholder__"] + assert xfer.hit_policy == PoolHitPolicy.TRAILING_PAGES + assert xfer.host_indices.tolist() == [30, 31] + + # The prefetched suffix lands as one host node; its SWA host is a tombstone. + insert_result = core.insert_host( + anchor, + _key([1, 2, 3]), + torch.tensor([100, 101, 102], dtype=torch.int64), + ["h0", "h1", "h2"], + ) + assert insert_result.total_len == 3 + assert insert_result.inserted_host_node is not None + + def commit(host_indices, loaded_pages): + actions = [] + core.commit_hicache_transfers( + anchor, + CacheTransferPhase.PREFETCH, + { + ComponentType.SWA: [ + PoolTransfer( + name=PoolName.SWA, + host_indices=torch.tensor(host_indices, dtype=torch.int64), + ) + ] + }, + cache_actions=actions, + insert_result=insert_result, + pool_storage_result=PoolTransferResult( + kv_hit_pages=3, extra_pool_hit_pages={PoolName.SWA: loaded_pages} + ), + ) + return actions + + # Underloaded window (1 of 2 pages): all-or-nothing frees the whole buffer. + (free,) = commit([30, 31], loaded_pages=1) + assert isinstance(free, FreeComponentHostSlot) + assert free.component_type == ComponentType.SWA + assert free.host_indices[0].tolist() == [30, 31] + + # A full window splits the partially covered node and attaches its tail. + assert commit([40, 41], loaded_pages=2) == [] + + # The window is hosted now: a re-prefetched buffer releases instead. + (release,) = commit([50, 51], loaded_pages=2) + assert isinstance(release, FreeComponentHostSlot) + assert release.host_indices[0].tolist() == [50, 51] + + +def test_swa_locked_overlap_defers_through_the_recover_action(): + core = _swa_tree_core(window=8) + first = _insert(core, [1, 2], [10, 11]) + node = first.cache_actions[0].node_id + core.inc_lock_ref(node) + result = _pump_insert( + core, + InsertParams( + key=_key([1, 2]), + value=torch.tensor([20, 21], dtype=torch.int64), + ), + ) + (recover,) = result.cache_actions + assert isinstance(recover, RecoverSWAWithLockedFull) + assert recover.node_id == node + assert recover.kept_full.tolist() == [10, 11] + assert recover.incoming_full.tolist() == [20, 21] + + +def test_component_device_value_round_trips(): + core = _swa_tree_core(window=8) + first = _insert(core, [1, 2], [10, 11]) + node = first.cache_actions[0].node_id + assert core.get_component_device_value(node, ComponentType.SWA) is None + core.set_component_device_value( + node, ComponentType.SWA, torch.tensor([50, 51], dtype=torch.int64) + ) + stored = core.get_component_device_value(node, ComponentType.SWA) + assert stored.tolist() == [50, 51] + + +def test_lock_uuid_round_trips_through_dec_lock_ref(): + from sglang.srt.mem_cache.base_prefix_cache import DecLockRefParams + + core = _swa_tree_core(window=2) + first = _insert(core, [1, 2, 3], [10, 11, 12]) + # The window cap split the leaf: rebuild the in-window nodes' SWA values. + for action in first.cache_actions: + core.set_component_device_value( + action.node_id, + ComponentType.SWA, + torch.arange(50, 50 + action.source_value.numel(), dtype=torch.int64), + ) + node = first.cache_actions[-1].node_id + result = core.inc_lock_ref(node) + assert result.swa_uuid_for_lock is not None + assert result.swa_uuid_for_host_lock is None + # The locked window is protected SWA accounting, visible through the binding. + assert core.swa_protected_size() == 2 + assert core.swa_evictable_size() == 1 + core.dec_lock_ref( + node, + DecLockRefParams( + swa_uuid_for_lock=result.swa_uuid_for_lock, + skip_lock_node_ids=result.skip_lock_node_ids, + ), + ) + # The uuid-bounded release returned the window to evictable. + assert core.swa_protected_size() == 0 + assert core.swa_evictable_size() == 3 + # A repeat acquire reuses the stamped uuid. + again = core.inc_lock_ref(node) + assert again.swa_uuid_for_lock == result.swa_uuid_for_lock + + +def test_swa_skip_map_crosses_the_binding_and_replays(): + from sglang.srt.mem_cache.base_prefix_cache import DecLockRefParams + + core = _swa_tree_core(window=8) + _insert(core, [1, 2], [10, 11]) + second = _insert(core, [1, 2, 3, 4], [10, 11, 12, 13]) + leaf = second.cache_actions[-1].node_id + # Only the leaf carries SWA; its ancestor is recorded as a tombstone skip. + core.set_component_device_value( + leaf, ComponentType.SWA, torch.tensor([52, 53], dtype=torch.int64) + ) + result = core.inc_lock_ref(leaf) + assert result.skip_lock_node_ids[ComponentType.SWA] + core.dec_lock_ref( + leaf, + DecLockRefParams( + swa_uuid_for_lock=result.swa_uuid_for_lock, + skip_lock_node_ids=result.skip_lock_node_ids, + ), + ) + assert core.swa_protected_size() == 0 + + +def test_dec_swa_lock_only_frees_flow_after_the_full_release(): + from sglang.srt.mem_cache.base_prefix_cache import DecLockRefParams + + core = _swa_tree_core(window=2) + first = _insert(core, [1, 2], [10, 11]) + node = first.cache_actions[0].node_id + core.set_component_device_value( + node, ComponentType.SWA, torch.tensor([50, 51], dtype=torch.int64) + ) + result = core.inc_lock_ref(node) + # The FULL lock releases first (skip_swa), then the early window release + # finds a fully unlocked device leaf and evicts it in place. + core.dec_lock_ref( + node, + DecLockRefParams(skip_lock_node_ids=result.skip_lock_node_ids), + skip_swa=True, + ) + device_frees: dict = {} + host_frees: dict = {} + _accumulate_step( + core.dec_swa_lock_only(node, result.swa_uuid_for_lock), + {}, + device_frees, + host_frees, + ) + assert [t.tolist() for t in device_frees[ComponentType.SWA]] == [[10, 11]] + assert core.get_component_device_value(node, ComponentType.SWA) is None + + +def test_dec_swa_lock_only_returns_the_window_frees(): + core = _swa_tree_core(window=2) + first = _insert(core, [1, 2, 3], [10, 11, 12]) + for action in first.cache_actions: + core.set_component_device_value( + action.node_id, + ComponentType.SWA, + torch.arange(50, 50 + action.source_value.numel(), dtype=torch.int64), + ) + node = first.cache_actions[-1].node_id + result = core.inc_lock_ref(node) + device_frees: dict = {} + host_frees: dict = {} + _accumulate_step( + core.dec_swa_lock_only(node, result.swa_uuid_for_lock), + {}, + device_frees, + host_frees, + ) + # The FULL lock still protects the path: the SWA release frees nothing and + # the rebuilt values survive; a repeat release is a no-op. + assert device_frees == {} + assert core.get_component_device_value(node, ComponentType.SWA) is not None + _accumulate_step( + core.dec_swa_lock_only(node, result.swa_uuid_for_lock), + {}, + device_frees, + host_frees, + ) + assert device_frees == {} + + +def test_swa_rebuild_applies_through_the_python_allocator(): + cache, allocator = _swa_cache(window=8) + full = allocator.alloc(4) + result = cache.insert(InsertParams(key=_key([1, 2, 3, 4]), value=full)) + assert result.prefix_len == 0 + node = cache.match_prefix(MatchPrefixParams(key=_key([1, 2, 3, 4]))).best_match_node + assert node != 0, "the SWA-covered match must reach the leaf" + # The cache executed SWARebuild through the allocator: the node holds the + # full slice's SWA translation. + stored = cache.tree_core.get_component_device_value(node, ComponentType.SWA) + expected = allocator.translate_loc_from_full_to_swa(full) + assert stored is not None + assert stored.tolist() == expected.tolist() + assert (allocator.full_to_swa_index_mapping[full.to(torch.int64)] > 0).all() + + +def test_recover_with_locked_full_applies_through_the_python_allocator(): + cache, allocator = _swa_cache(window=8) + kept = allocator.alloc(2) + cache.insert(InsertParams(key=_key([1, 2]), value=kept)) + node = cache.match_prefix(MatchPrefixParams(key=_key([1, 2]))).best_match_node + assert node != 0 + lock = cache.inc_lock_ref(node) + # The decode advanced past the window: the SWA lock releases early, then + # window eviction tombstones the SWA slot under the FULL lock (the state a + # locked-full overlap recovers from); its frees return to the allocator. + cache.dec_swa_lock_only(node, lock.swa_uuid_for_lock) + tracker = {ComponentType.FULL: 0, ComponentType.SWA: 0} + device_frees: dict = {} + host_frees: dict = {} + cache.tree_core.evict_device_start(ComponentType.SWA, 100) + step = cache.tree_core.evict_device_next_node(ComponentType.SWA, tracker) + assert step.node_id is None + _accumulate_step(step, tracker, device_frees, host_frees) + cache.tree_core.evict_device_end(ComponentType.SWA) + for freed in device_frees[ComponentType.SWA]: + allocator.free_swa(freed) + assert cache.tree_core.get_component_device_value(node, ComponentType.SWA) is None + incoming = allocator.alloc(2) + before_free = allocator.full_attn_allocator.available_size() + cache.components[ComponentType.SWA].apply_component_action( + RecoverSWAWithLockedFull(node_id=node, kept_full=kept, incoming_full=incoming) + ) + # The locked full keeps its slots, remapped onto the incoming full's SWA + # translation; the incoming full is freed back to the allocator. + stored = cache.tree_core.get_component_device_value(node, ComponentType.SWA) + assert stored.tolist() == allocator.translate_loc_from_full_to_swa(kept).tolist() + assert (allocator.full_to_swa_index_mapping[incoming.to(torch.int64)] == 0).all() + assert ( + allocator.full_attn_allocator.available_size() == before_free + incoming.numel() + ) + + +# ==== Bigram (EAGLE) wiring ==== + + +def _bigram_tree_core(**params_overrides) -> RustUnifiedTreeCore: + # Without mamba, the core honors is_eagle and selects the bigram binding. + params = dict( + disable=False, + req_to_token_pool=None, + token_to_kv_pool_allocator=None, + page_size=1, + is_eagle=True, + tree_components=(ComponentType.FULL,), + ) + params.update(params_overrides) + return RustUnifiedTreeCore(CacheInitParams(**params)) + + +def _bigram_key(token_ids: list[int]) -> RadixKey: + return RadixKey(array("q", token_ids), is_bigram=True) + + +def test_bigram_insert_then_a_longer_match_returns_the_inserted_prefix(): + core = _bigram_tree_core() + # 4 raw tokens = 3 bigram atoms, so the value carries 3 indices. + result = _pump_insert( + core, + InsertParams( + key=_bigram_key([1, 2, 3, 4]), + value=torch.tensor([10, 11, 12], dtype=torch.int64), + ), + ) + assert result.prefix_len == 0 + matched = core.match_prefix(MatchPrefixParams(key=_bigram_key([1, 2, 3, 4, 5]))) + assert matched.device_indices.tolist() == [10, 11, 12] + + +def test_bigram_match_diverges_on_the_pair_not_the_token(): + core = _bigram_tree_core() + _pump_insert( + core, + InsertParams( + key=_bigram_key([1, 2, 3, 4]), + value=torch.tensor([10, 11, 12], dtype=torch.int64), + ), + ) + # (1, 2) matches; (2, 9) diverges from (2, 3) despite the shared token 2. + matched = core.match_prefix(MatchPrefixParams(key=_bigram_key([1, 2, 9]))) + assert matched.device_indices.tolist() == [10] + + +def test_bigram_empty_and_single_token_keys_match_nothing(): + core = _bigram_tree_core() + _pump_insert( + core, + InsertParams( + key=_bigram_key([1, 2, 3]), + value=torch.tensor([10, 11], dtype=torch.int64), + ), + ) + empty = core.match_prefix(MatchPrefixParams(key=_bigram_key([]))) + assert empty.device_indices.numel() == 0 + single = core.match_prefix(MatchPrefixParams(key=_bigram_key([1]))) + assert single.device_indices.numel() == 0 + + +def test_bigram_insert_truncates_a_raw_length_value_to_the_bigram_count(): + core = _bigram_tree_core() + result = _pump_insert( + core, + InsertParams( + key=_bigram_key([1, 2, 3]), + value=torch.tensor([10, 11, 12], dtype=torch.int64), + ), + ) + assert result.prefix_len == 0 + matched = core.match_prefix(MatchPrefixParams(key=_bigram_key([1, 2, 3]))) + assert matched.device_indices.tolist() == [10, 11] + + +def test_bigram_insert_value_shorter_than_the_bigram_count_raises(): + core = _bigram_tree_core() + with pytest.raises(ValueError, match="shorter than the aligned key length"): + _pump_insert( + core, + InsertParams( + key=_bigram_key([1, 2, 3, 4]), + value=torch.tensor([10, 11], dtype=torch.int64), + ), + ) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__])) diff --git a/test/registered/unit/mem_cache/test_rust_unified_radix_cache_bench.py b/test/registered/unit/mem_cache/test_rust_unified_radix_cache_bench.py new file mode 100644 index 000000000..9fbd2dae5 --- /dev/null +++ b/test/registered/unit/mem_cache/test_rust_unified_radix_cache_bench.py @@ -0,0 +1,29 @@ +"""Run the UnifiedRadixCache benchmark/fuzz suite with the Rust TreeCore.""" + +import unittest + +import test_unified_radix_cache_bench as shared_suite + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=25, stage="base-b", runner_config="1-gpu-small") + + +class RustBackendSuite(unittest.TestSuite): + """Scope the backend override to this suite and restore it afterward.""" + + def run(self, result, debug=False): + previous = shared_suite._TREE_CORE_TEST_BACKEND + shared_suite._TREE_CORE_TEST_BACKEND = "rust" + try: + return super().run(result, debug) + finally: + shared_suite._TREE_CORE_TEST_BACKEND = previous + + +def load_tests(loader, standard_tests, pattern): + return RustBackendSuite(loader.loadTestsFromModule(shared_suite)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_rust_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_rust_unified_radix_cache_unittest.py new file mode 100644 index 000000000..77f1dbddc --- /dev/null +++ b/test/registered/unit/mem_cache/test_rust_unified_radix_cache_unittest.py @@ -0,0 +1,30 @@ +"""Run the shared UnifiedRadixCache unit suite with the Rust TreeCore.""" + +import unittest + +import test_unified_radix_cache_unittest as shared_suite + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=180, stage="base-b", runner_config="1-gpu-small") + + +class RustBackendSuite(unittest.TestSuite): + """Scope the test backend to this suite without polluting discovery.""" + + def run(self, result, debug=False): + previous = shared_suite._TREE_CORE_TEST_BACKEND + shared_suite._TREE_CORE_TEST_BACKEND = "rust" + try: + return super().run(result, debug) + finally: + shared_suite._TREE_CORE_TEST_BACKEND = previous + + +def load_tests(loader, standard_tests, pattern): + """Reuse the exact cache-level suite while swapping only its test factory.""" + return RustBackendSuite(loader.loadTestsFromModule(shared_suite)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py index 6b1d42b2e..9a07d97ad 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py @@ -14,7 +14,7 @@ import sys import time import unittest from array import array -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from dataclasses import dataclass from typing import Callable @@ -40,6 +40,7 @@ from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=25, stage="base-b", runner_config="1-gpu-small") register_amd_ci(est_time=25, suite="stage-b-test-1-gpu-small-amd") @@ -59,6 +60,7 @@ _BENCH_KV_SIZE = 500_000 _BENCH_CHUNK_LEN = 256 _DEFAULT_COMPONENTS = (ComponentType.FULL, ComponentType.MAMBA) +_TREE_CORE_TEST_BACKEND: str | None = None @contextmanager @@ -226,16 +228,22 @@ def create_bench_cache( # --- tree --- if tree_cls is None: tree_cls = UnifiedRadixCache - tree = tree_cls( - params=CacheInitParams( - req_to_token_pool=req_to_token_pool, - token_to_kv_pool_allocator=allocator, - page_size=page_size, - disable=False, - tree_components=components if tree_cls is UnifiedRadixCache else None, - sliding_window_size=sliding_window_size if has_swa else None, - ) + backend_override = ( + envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.override(_TREE_CORE_TEST_BACKEND) + if _TREE_CORE_TEST_BACKEND is not None and tree_cls is UnifiedRadixCache + else nullcontext() ) + with backend_override: + tree = tree_cls( + params=CacheInitParams( + req_to_token_pool=req_to_token_pool, + token_to_kv_pool_allocator=allocator, + page_size=page_size, + disable=False, + tree_components=components if tree_cls is UnifiedRadixCache else None, + sliding_window_size=sliding_window_size if has_swa else None, + ) + ) _rid = [0] @@ -780,6 +788,10 @@ class _BenchSuite: verify=True, page_size=cfg["page_size"], ) + backend = ( + _TREE_CORE_TEST_BACKEND or envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.get() + ) + print(f"[{backend}] {r.report()}") self.assertGreater(r.num_ops, 0) self.assertGreater(r.ops_per_sec, 0) @@ -803,7 +815,7 @@ for _cfg in _CI_BENCH_CONFIGS: _name = f"TestBench_{_cfg['label']}" globals()[_name] = type( _name, - (_BenchSuite, unittest.TestCase), + (_BenchSuite, CustomTestCase), {"bench_cfg": _cfg}, ) globals()[_name].__module__ = __name__ diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 960f45a3b..fd6951d03 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -21,6 +21,7 @@ from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape from sglang.srt.disaggregation.kv_events import ( BlockRemoved, BlockStored, + BlockStoredWithMetadata, StorageMedium, ) from sglang.srt.environ import envs @@ -103,6 +104,22 @@ from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=50, stage="base-b", runner_config="1-gpu-small") register_amd_ci(est_time=50, suite="stage-b-test-1-gpu-small-amd") +# A dedicated test entry point overrides this without changing the process-wide +# production backend selection. Direct Python-core tests in this module remain +# Python-only; every fixture-backed cache test is shared by both inspectors. +_TREE_CORE_TEST_BACKEND: Optional[str] = None + + +def _selected_tree_core_test_backend() -> str: + return _TREE_CORE_TEST_BACKEND or envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.get() + + +def _session_radix_cache_test_values() -> tuple[bool, ...]: + # TODO(Jialin): Restore the session-enabled case after porting #29173 to Rust. + if _selected_tree_core_test_backend() == "rust": + return (False,) + return False, True + @dataclass(frozen=True) class CacheConfig: @@ -558,15 +575,25 @@ def build_fixture( eviction_policy=cfg.eviction_policy, is_eagle=cfg.is_eagle, ) - selected_backend = envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.get() + selected_backend = _selected_tree_core_test_backend() if selected_backend == "python": - with mock.patch.dict( - _TREE_CORE_REGISTRY, - { - "python": lambda params, components: UnifiedTreeCoreInspector( - params, components - ) - }, + + def inspector_factory(params, components): + return UnifiedTreeCoreInspector(params, components) + + elif selected_backend == "rust": + from rust_unified_tree_core_inspector import RustUnifiedTreeCoreInspector + + def inspector_factory(params, _components): + return RustUnifiedTreeCoreInspector(params) + + else: + inspector_factory = None + + if inspector_factory is not None: + with ( + envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.override(selected_backend), + mock.patch.dict(_TREE_CORE_REGISTRY, {selected_backend: inspector_factory}), ): cache = UnifiedRadixCache(params=cache_init_params) else: @@ -580,6 +607,77 @@ def build_fixture( return cache, allocator, req_to_token_pool +@unittest.skipUnless(torch.cuda.is_available(), "cache fixtures need CUDA") +class TestUnifiedTreeCoreLoadBackOwnershipBackends(CustomTestCase): + """Run Full load-back ownership semantics through either TreeCore backend.""" + + cfg = CacheConfig( + page_size=1, + components=(ComponentType.FULL, ComponentType.SWA), + sliding_window_size=4, + ) + + def test_auxiliary_load_does_not_reuse_full_pending_pin(self): + cache, _, _ = build_fixture(self.cfg) + core = cache.tree_core + core.is_write_back = True + + root = cache.root_node_handle() + + def insert_host(tokens, indices, hashes): + return core.insert_host( + root, + RadixKey(array("q", tokens)), + torch.tensor(indices, dtype=torch.int64), + hashes, + ).inserted_host_node + + def full_transfer(node_id, host_index): + return PoolTransfer( + name=PoolName.KV, + host_indices=torch.tensor([host_index], dtype=torch.int64), + nodes_to_load=[node_id], + ) + + shared = insert_host([1], [100], ["h0"]) + anchor_b = insert_host([1, 2], [100, 101], ["h0", "h1"]) + self.assertIsNotNone(shared) + self.assertIsNotNone(anchor_b) + + core.set_component_host_value_raw( + shared, ComponentType.SWA, torch.tensor([200], dtype=torch.int64) + ) + core.commit_load_back( + shared, + torch.tensor([10], dtype=torch.int64), + full_transfer(shared, 100), + {}, + ) + + actions = core.commit_load_back( + anchor_b, + torch.tensor([11], dtype=torch.int64), + full_transfer(anchor_b, 101), + { + ComponentType.SWA: [ + PoolTransfer( + name=PoolName.SWA, + host_indices=torch.tensor([200], dtype=torch.int64), + device_indices=torch.tensor([20], dtype=torch.int64), + nodes_to_load=[shared], + ) + ] + }, + ) + + self.assertEqual(_device_value(cache, shared, ComponentType.SWA).tolist(), [20]) + self.assertTrue( + any(isinstance(action, RebuildFullToSWAMapping) for action in actions) + ) + core.finish_load_back(anchor_b) + core.finish_load_back(shared) + + @unittest.skipUnless(torch.cuda.is_available(), "cache fixtures need CUDA") class TestUnifiedRadixAllocationEvictionRealComponents(CustomTestCase): """Allocation targets are observed between real auxiliary-tree steps.""" @@ -673,7 +771,7 @@ class TestUnifiedRadixAllocationEvictionRealComponents(CustomTestCase): def test_allocation_target_stops_after_one_internal_tombstone(self): for component_type in (ComponentType.SWA, ComponentType.MAMBA): - for enable_session_radix_cache in (False, True): + for enable_session_radix_cache in _session_radix_cache_test_values(): with self.subTest( component_type=component_type, enable_session_radix_cache=enable_session_radix_cache, @@ -703,7 +801,7 @@ class TestUnifiedRadixAllocationEvictionRealComponents(CustomTestCase): def test_explicit_evict_continues_across_internal_steps(self): for component_type in (ComponentType.SWA, ComponentType.MAMBA): - for enable_session_radix_cache in (False, True): + for enable_session_radix_cache in _session_radix_cache_test_values(): with self.subTest( component_type=component_type, enable_session_radix_cache=enable_session_radix_cache, @@ -820,12 +918,93 @@ class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase): canonical_hashes.append(running_hash) self.assertNotEqual(canonical_hashes, cache.tree_core.get_hash_values(leaf)) + def test_buffer_anchor_rematch_preserves_bigram_boundary(self): + from sglang.srt.mem_cache.buffer_mode.pipeline import BufferModePipeline + + cache, allocator, _ = build_fixture(self.cfg) + extra_key = "adapter-a" + cache_salt = "tenant-a" + prefix_tokens = array("q", [1, 2, 3, 4, 5]) + value = allocator.alloc(len(prefix_tokens) - 1) + self.assertIsNotNone(value) + cache.insert( + InsertParams( + key=RadixKey( + prefix_tokens, + extra_key=extra_key, + cache_salt=cache_salt, + ), + value=value, + ) + ) + match = cache.match_prefix( + MatchPrefixParams( + key=RadixKey( + prefix_tokens, + extra_key=extra_key, + cache_salt=cache_salt, + ) + ) + ) + self.assertEqual(len(match.device_indices), len(prefix_tokens) - 1) + + req_id = "bigram-anchor" + prefetch_key = RadixKey( + array("q", [prefix_tokens[-1], 6, 7, 8, 9]), + extra_key=extra_key, + is_bigram=True, + cache_salt=cache_salt, + ) + cache.ongoing_prefetch[req_id] = _OngoingPrefetch( + anchor_node_id=match.last_device_node, + prefetch_key=prefetch_key, + host_indices=None, + operation=None, + anchor_lock_params=None, + comp_xfers={}, + ) + + pipeline = BufferModePipeline.__new__(BufferModePipeline) + pipeline.anchor_lock_enabled = True + pipeline.anchor_locks = {} + pipeline.anchor_locked_tokens_ = 0 + pipeline.anchor_lock_cap_tokens = 10_000 + pipeline._anchor_lock_cap_skips = 0 + pipeline._prefetch_prefix_ctx = { + req_id: (list(prefix_tokens[:-1]), extra_key, cache_salt) + } + pipeline._cache = cache + + lock_ref = _device_lock_ref(cache, match.last_device_node, ComponentType.FULL) + self.assertEqual(pipeline.try_lock_anchor(req_id), "locked") + self.assertEqual(pipeline.anchor_locks[req_id].node_id, match.last_device_node) + self.assertEqual( + _device_lock_ref(cache, match.last_device_node, ComponentType.FULL), + lock_ref + 1, + ) + + pipeline.release_anchor_lock(req_id) + cache.ongoing_prefetch.pop(req_id) + self.assertEqual( + _device_lock_ref(cache, match.last_device_node, ComponentType.FULL), + lock_ref, + ) + cache.sanity_check() + class TestUnifiedRadixCacheKVEvents(CustomTestCase): cfg = CacheConfig(page_size=2, kv_size=64, max_context_len=64) - def _insert(self, cache, allocator, tokens): - key = RadixKey(array("q", tokens)) + def _insert( + self, + cache, + allocator, + tokens, + *, + extra_key=None, + cache_salt=None, + ): + key = RadixKey(array("q", tokens), extra_key=extra_key, cache_salt=cache_salt) value = allocator.alloc(len(tokens)) self.assertIsNotNone(value) return cache.insert(InsertParams(key=key, value=value[: len(key)])) @@ -900,6 +1079,72 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase): self.assertEqual(len(removed), 1) self.assertEqual(removed[0].block_hashes, stored_hashes) + def test_cache_salt_is_included_in_store_and_remove_events(self): + cache, allocator, _ = build_fixture(self.cfg, enable_kv_cache_events=True) + cache.take_events() + + seq = [1, 2, 3, 4] + self._insert(cache, allocator, seq, cache_salt="tenant-a") + stored = self._stored_events(cache, StorageMedium.GPU) + self.assertEqual(len(stored), 1) + self.assertIsInstance(stored[0], BlockStoredWithMetadata) + self.assertEqual(stored[0].metadata.cache_salt, "tenant-a") + salted_hashes = self._event_hashes(stored) + + cache.evict(EvictParams(num_tokens=len(seq))) + removed = self._removed_events(cache, StorageMedium.GPU) + self.assertEqual(len(removed), 1) + self.assertEqual(removed[0].block_hashes, salted_hashes) + + unsalted, unsalted_allocator, _ = build_fixture( + self.cfg, enable_kv_cache_events=True + ) + unsalted.take_events() + self._insert(unsalted, unsalted_allocator, seq) + self.assertNotEqual( + self._event_hashes(self._stored_events(unsalted, StorageMedium.GPU)), + salted_hashes, + ) + + def test_cache_salt_event_parentage_survives_node_split(self): + cache, allocator, _ = build_fixture(self.cfg, enable_kv_cache_events=True) + cache.take_events() + + self._insert(cache, allocator, [1, 2, 3, 4], cache_salt="tenant-a") + original = self._stored_events(cache, StorageMedium.GPU) + self.assertEqual(len(original), 1) + self.assertEqual(len(original[0].block_hashes), 2) + + self._insert(cache, allocator, [1, 2, 5, 6], cache_salt="tenant-a") + branch = self._stored_events(cache, StorageMedium.GPU) + self.assertEqual(len(branch), 1) + self.assertIsInstance(branch[0], BlockStoredWithMetadata) + self.assertEqual(branch[0].metadata.cache_salt, "tenant-a") + self.assertEqual(branch[0].parent_block_hash, original[0].block_hashes[0]) + self.assertEqual(list(branch[0].token_ids), [5, 6]) + + def test_event_hashes_depend_on_cache_salt_but_not_extra_key(self): + def stored_hashes(*, extra_key=None, cache_salt=None): + cache, allocator, _ = build_fixture(self.cfg, enable_kv_cache_events=True) + cache.take_events() + self._insert( + cache, + allocator, + [1, 2, 3, 4], + extra_key=extra_key, + cache_salt=cache_salt, + ) + return self._event_hashes(self._stored_events(cache, StorageMedium.GPU)) + + self.assertNotEqual( + stored_hashes(cache_salt="tenant-a"), + stored_hashes(cache_salt="tenant-b"), + ) + self.assertEqual( + stored_hashes(extra_key="adapter-a", cache_salt="tenant-a"), + stored_hashes(extra_key="adapter-b", cache_salt="tenant-a"), + ) + def test_kv_events_split_preserves_block_hash_parentage(self): cache, allocator, _ = build_fixture(self.cfg, enable_kv_cache_events=True) cache.take_events() # Clear the reset event. @@ -1068,9 +1313,19 @@ class UnifiedRadixCacheSuite: allocator.full_to_swa_index_mapping[full_indices] = swa_indices return full_indices[:need_size] - def _insert(self, cache, allocator, req_to_token_pool, tokens, priority=0): + def _insert( + self, + cache, + allocator, + req_to_token_pool, + tokens, + priority=0, + *, + extra_key=None, + cache_salt=None, + ): """Insert tokens, attaching mamba data when the config has mamba.""" - key = RadixKey(array("q", tokens)) + key = RadixKey(array("q", tokens), extra_key=extra_key, cache_salt=cache_salt) value = self._alloc(allocator, len(tokens)) params = InsertParams(key=key, value=value[: len(key)], priority=priority) if self.cfg.has_mamba: @@ -1089,6 +1344,7 @@ class UnifiedRadixCacheSuite: self.assertEqual(result.prefix_len, len(seq_a)) m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_b)))) + self.assertEqual(result.last_device_node, m.last_device_node) self.assertEqual(len(m.device_indices), len(seq_b)) m = cache.match_prefix( @@ -1103,6 +1359,48 @@ class UnifiedRadixCacheSuite: cache.sanity_check() + def test_cache_salt_and_extra_key_form_independent_namespaces(self): + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + seq = self._make_seq(1, 2) + + first = self._insert( + cache, + allocator, + req_to_token_pool, + seq, + extra_key="bc", + cache_salt="a", + ) + second = self._insert( + cache, + allocator, + req_to_token_pool, + seq, + extra_key="c", + cache_salt="ab", + ) + self.assertEqual(first.prefix_len, 0) + self.assertEqual(second.prefix_len, 0) + + first_match = cache.match_prefix( + MatchPrefixParams( + key=RadixKey(array("q", seq), extra_key="bc", cache_salt="a") + ) + ) + second_match = cache.match_prefix( + MatchPrefixParams( + key=RadixKey(array("q", seq), extra_key="c", cache_salt="ab") + ) + ) + default_match = cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", seq))) + ) + self.assertEqual(len(first_match.device_indices), len(seq)) + self.assertEqual(len(second_match.device_indices), len(seq)) + self.assertEqual(len(default_match.device_indices), 0) + self.assertNotEqual(first_match.last_device_node, second_match.last_device_node) + cache.sanity_check() + def test_shared_prefix_split(self): cache, allocator, req_to_token_pool = build_fixture(self.cfg) base = self._make_seq(1, 2) @@ -1823,6 +2121,42 @@ class UnifiedRadixCacheSuite: "Full stays locked", ) + def test_swa_early_release_preserves_an_owners_skipped_mamba_lock(self): + if not self.cfg.has_swa or not self.cfg.has_mamba: + self.skipTest("requires SWA and Mamba components") + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + + seq_a = self._make_seq( + 1, (self.cfg.sliding_window_size // self.cfg.page_size) + 4 + ) + self._insert(cache, allocator, req_to_token_pool, seq_a) + node_a = cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", seq_a))) + ).last_device_node + + owner = cache.inc_lock_ref(node_a) + skipped = cache.inc_lock_ref( + node_a, skip_lock_components=(ComponentType.MAMBA,) + ) + self.assertEqual(skipped.skip_lock_node_ids, {ComponentType.MAMBA: {node_a}}) + self.assertEqual(_device_lock_ref(cache, node_a, ComponentType.MAMBA), 1) + + cache.dec_swa_lock_only( + node_a, + skipped.swa_uuid_for_lock, + skip_lock_node_ids=skipped.skip_lock_node_ids, + ) + self.assertEqual( + _device_lock_ref(cache, node_a, ComponentType.MAMBA), + 1, + "the early release must not drop another request's Mamba lock", + ) + + cache.dec_lock_ref(node_a, skipped.to_dec_params(), skip_swa=True) + cache.dec_lock_ref(node_a, owner.to_dec_params()) + self.assertEqual(_device_lock_ref(cache, node_a, ComponentType.MAMBA), 0) + cache.sanity_check() + def test_cascade_evict_asserts_on_locked_internal_mamba(self): if not self.cfg.has_swa or not self.cfg.has_mamba: self.skipTest("requires SWA and Mamba components") @@ -2707,7 +3041,16 @@ class UnifiedRadixCacheSuite: self.fail(f"prefetch {req_id} did not complete in time") def _consume_staged_prefetch( - self, cache, req_id, prefix_len=None, prefix_indices=None, timeout: float = 10.0 + self, + cache, + req_id, + prefix_len=None, + prefix_indices=None, + timeout: float = 10.0, + *, + extra_key=None, + cache_salt=None, + last_node=None, ): """Simulate the PrefillAdder consuming a staged prefetch at admission: init_load_back (buffer dispatch: device alloc + queued H2D), the batch @@ -2720,8 +3063,8 @@ class UnifiedRadixCacheSuite: prefix_len = f.matched_len req = mock.Mock() req.rid = req_id - req.extra_key = None - req.cache_salt = None + req.extra_key = extra_key + req.cache_salt = cache_salt if prefix_indices is not None: # Spliceable mid-anchor consumption publishes value=cat(prefix, # fill) — the real device prefix is required (zeros would insert @@ -2734,7 +3077,7 @@ class UnifiedRadixCacheSuite: dtype=torch.int64, device=cache.tree_core.empty_match_result.device_indices.device, ) - req.last_node = cache.root_node_handle() + req.last_node = cache.root_node_handle() if last_node is None else last_node new_indices, _last_node = cache.init_load_back( InitLoadBackParams( best_match_node=None, host_hit_length=f.num_tokens, req=req @@ -3145,6 +3488,29 @@ class UnifiedRadixCacheSuite: # Buffer-only host memory mode (host = transient staging, L3 = cache) # ================================================================ + def test_buffer_only_rejects_mamba(self): + if ( + self.cfg.components + != ( + ComponentType.FULL, + ComponentType.MAMBA, + ) + or self.cfg.page_size != 1 + ): + self.skipTest("one FULL+MAMBA page_size=1 fixture covers this guard") + + cache, _, _ = build_fixture(self.cfg) + storage_dir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True) + with self.assertRaisesRegex(ValueError, "supports only FULL/SWA"): + self._init_hicache( + cache, + storage_backend="file", + storage_dir=storage_dir, + prefetch_threshold=1, + host_memory_mode="buffer_only", + ) + def _init_buffer_hicache( self, cache, @@ -3210,17 +3576,25 @@ class UnifiedRadixCacheSuite: "buffer backup pipeline did not drain", ) - def _produce_buffer_l3(self, storage_dir, seq, marker=None): + def _produce_buffer_l3( + self, storage_dir, seq, marker=None, *, extra_key=None, cache_salt=None + ): """Producer tree in buffer mode: insert seq and push it to L3.""" prod, prod_alloc, prod_rtp = build_fixture(self.cfg) self._init_buffer_hicache(prod, storage_dir) - self._insert(prod, prod_alloc, prod_rtp, seq) - leaf = prod.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", seq))) - ).last_device_node + self._insert( + prod, + prod_alloc, + prod_rtp, + seq, + extra_key=extra_key, + cache_salt=cache_salt, + ) + key = RadixKey(array("q", seq), extra_key=extra_key, cache_salt=cache_salt) + leaf = prod.match_prefix(MatchPrefixParams(key=key)).last_device_node expected = None if marker is not None: - m = prod.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) + m = prod.match_prefix(MatchPrefixParams(key=key)) self._fill_full_kv(prod_alloc, m.device_indices, marker=marker) expected = self._snapshot_full_kv(prod_alloc, m.device_indices) self._buffer_backup_and_wait(prod, leaf) @@ -3409,6 +3783,131 @@ class UnifiedRadixCacheSuite: self.assertIn("occupancy_ratio", cons.prefetch_outcome_stats_snapshot()) cons.sanity_check() + def test_buffer_only_cache_salt_uses_the_request_namespace(self): + self._skip_unsupported_hicache_test() + if self.cfg.components != (ComponentType.FULL,) or self.cfg.page_size != 4: + self.skipTest("one FULL page_size=4 fixture covers namespace routing") + + anchor_lock = envs.SGLANG_ENABLE_HICACHE_BUFFER_ANCHOR_LOCK.override(True) + anchor_lock.__enter__() + self.addCleanup(anchor_lock.__exit__, None, None, None) + storage_dir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True) + + extra_key = "adapter-a" + cache_salt = "tenant-a" + seq = self._make_seq(1, 2) + self._produce_buffer_l3( + storage_dir, seq, extra_key=extra_key, cache_salt=cache_salt + ) + + # A root anchor has no namespace of its own. The fetched span must use + # the request namespace supplied to prefetch_from_storage. + cons, _, _ = build_fixture(self.cfg) + self._init_buffer_hicache(cons, storage_dir) + root_req = "salted-root-prefetch" + cons.prefetch_from_storage( + root_req, + cons.root_node_handle(), + array("q", seq), + None, + None, + extra_key=extra_key, + cache_salt=cache_salt, + ) + self._pump_hicache_until( + cons, + lambda: cons.check_prefetch_progress(root_req) + and cons.buffer_pipeline.has_staged(root_req), + "salted root prefetch did not stage", + ) + held = cons.buffer_pipeline.staged_prefetches[root_req] + self.assertEqual((held.extra_key, held.cache_salt), (extra_key, cache_salt)) + self.assertNotIn(root_req, cons.buffer_pipeline.anchor_locks) + loaded = self._consume_staged_prefetch( + cons, + root_req, + extra_key=extra_key, + cache_salt=cache_salt, + ) + self.assertEqual(len(loaded), len(seq)) + key = RadixKey(array("q", seq), extra_key=extra_key, cache_salt=cache_salt) + self.assertEqual( + len(cons.match_prefix(MatchPrefixParams(key=key)).device_indices), len(seq) + ) + for miss in ( + RadixKey(array("q", seq)), + RadixKey(array("q", seq), extra_key=extra_key, cache_salt="tenant-b"), + RadixKey(array("q", seq), extra_key="adapter-b", cache_salt=cache_salt), + ): + self.assertEqual( + len(cons.match_prefix(MatchPrefixParams(key=miss)).device_indices), 0 + ) + cons.sanity_check() + + # A mid-tree anchor is pinned while its suffix is in flight and + # released as soon as the staged span is consumed. + cons2, cons2_alloc, cons2_rtp = build_fixture(self.cfg) + self._init_buffer_hicache(cons2, storage_dir) + prefix = seq[: self.cfg.page_size] + suffix = seq[self.cfg.page_size :] + self._insert( + cons2, + cons2_alloc, + cons2_rtp, + prefix, + extra_key=extra_key, + cache_salt=cache_salt, + ) + prefix_match = cons2.match_prefix( + MatchPrefixParams( + key=RadixKey( + array("q", prefix), + extra_key=extra_key, + cache_salt=cache_salt, + ) + ) + ) + anchor = prefix_match.last_device_node + lock_ref = _device_lock_ref(cons2, anchor, ComponentType.FULL) + anchored_req = "salted-mid-tree-prefetch" + cons2.prefetch_from_storage( + anchored_req, + anchor, + array("q", suffix), + cons2.tree_core.get_last_hash_value(anchor), + None, + matched_prefix_tokens=prefix, + extra_key=extra_key, + cache_salt=cache_salt, + ) + self.assertIn(anchored_req, cons2.buffer_pipeline.anchor_locks) + self.assertEqual( + _device_lock_ref(cons2, anchor, ComponentType.FULL), lock_ref + 1 + ) + self._pump_hicache_until( + cons2, + lambda: cons2.check_prefetch_progress(anchored_req) + and cons2.buffer_pipeline.has_staged(anchored_req), + "salted mid-tree prefetch did not stage", + ) + self._consume_staged_prefetch( + cons2, + anchored_req, + prefix_len=len(prefix), + prefix_indices=prefix_match.device_indices, + extra_key=extra_key, + cache_salt=cache_salt, + last_node=anchor, + ) + self.assertNotIn(anchored_req, cons2.buffer_pipeline.anchor_locks) + self.assertEqual(cons2.buffer_pipeline.anchor_locked_tokens_, 0) + self.assertEqual(_device_lock_ref(cons2, anchor, ComponentType.FULL), lock_ref) + self.assertEqual( + len(cons2.match_prefix(MatchPrefixParams(key=key)).device_indices), len(seq) + ) + cons2.sanity_check() + def test_buffer_only_storage_prefetch_miss_marker_and_retry(self): """A too-early storage query misses; the miss must arm the retry marker exactly once, a re-issued check must serve once the content @@ -3430,7 +3929,7 @@ class UnifiedRadixCacheSuite: # Query BEFORE any producer wrote the span: full miss -> revoked. req_id = "early-query-miss" cons.prefetch_from_storage( - req_id, cons.root_node.id, array("q", seq), None, None + req_id, cons.root_node_handle(), array("q", seq), None, None ) self._pump_hicache_until( cons, @@ -3446,7 +3945,7 @@ class UnifiedRadixCacheSuite: # and stages what the first, too-early query could not see. self._produce_buffer_l3(storage_dir, seq) cons.prefetch_from_storage( - req_id, cons.root_node.id, array("q", seq), None, None + req_id, cons.root_node_handle(), array("q", seq), None, None ) self._pump_hicache_until( cons, @@ -3461,7 +3960,7 @@ class UnifiedRadixCacheSuite: aborted_rid = "aborted-miss" cons.prefetch_from_storage( aborted_rid, - cons.root_node.id, + cons.root_node_handle(), array("q", self._make_seq(700, 4)), None, None, @@ -3477,7 +3976,7 @@ class UnifiedRadixCacheSuite: # A fully-device-matched (empty-suffix) decline also arms the retry: # the device match can evict while the request waits in the queue. cons.prefetch_from_storage( - "fully-matched", cons.root_node.id, array("q", []), None, None + "fully-matched", cons.root_node_handle(), array("q", []), None, None ) self.assertTrue(cons.pop_storage_prefetch_miss("fully-matched")) cons.sanity_check() @@ -3763,8 +4262,10 @@ class UnifiedRadixCacheSuite: ) self._pump_hicache_until( cons, - lambda: cons.check_prefetch_progress(req_id) - and cons.buffer_pipeline.has_staged(req_id), + lambda: ( + cons.check_prefetch_progress(req_id) + and cons.buffer_pipeline.has_staged(req_id) + ), "prefetch did not stage", ) cons.pop_prefetch_loaded_tokens(req_id) @@ -3825,7 +4326,7 @@ class UnifiedRadixCacheSuite: req_id = "growth-trim" cons.prefetch_from_storage( - req_id, cons.root_node.id, array("q", seq), None, None + req_id, cons.root_node_handle(), array("q", seq), None, None ) self._pump_hicache_until( cons, @@ -3892,7 +4393,7 @@ class UnifiedRadixCacheSuite: req_id = "covered-hold" cons.prefetch_from_storage( - req_id, cons.root_node.id, array("q", seq), None, None + req_id, cons.root_node_handle(), array("q", seq), None, None ) self._pump_hicache_until( cons, @@ -3936,7 +4437,7 @@ class UnifiedRadixCacheSuite: req_id = "covered-at-commit" cons.prefetch_from_storage( - req_id, cons.root_node.id, array("q", seq), None, None + req_id, cons.root_node_handle(), array("q", seq), None, None ) # Wait for the hit verdict WITHOUT draining it (the drain is the # scheduler-thread IO commit under test). @@ -4286,20 +4787,20 @@ class UnifiedRadixCacheSuite: # each evict() attempts the drop fallback on the parent. with mock.patch.object(cache.cache_controller, "write", return_value=None): # Pinned subtree root: drop declines, chain stays intact. - cache.inc_host_lock_ref(parent) + parent_lock_params = cache.inc_host_lock_ref(parent).to_dec_params() result = cache.evict(EvictParams(num_tokens=len(parent_seq))) self.assertEqual(result.num_tokens_evicted, 0) - cache.dec_host_lock_ref(parent) + cache.dec_host_lock_ref(parent, parent_lock_params) # Pinned host-only descendant: drop declines as well. - cache.inc_host_lock_ref(child) + child_lock_params = cache.inc_host_lock_ref(child).to_dec_params() result = cache.evict(EvictParams(num_tokens=len(parent_seq))) self.assertEqual(result.num_tokens_evicted, 0) m = cache.match_prefix( MatchPrefixParams(key=RadixKey(array("q", parent_seq))) ) self.assertEqual(len(m.device_indices), len(parent_seq)) - cache.dec_host_lock_ref(child) + cache.dec_host_lock_ref(child, child_lock_params) # Unpinned: the subtree drops and the child's host slots return. result = cache.evict(EvictParams(num_tokens=len(parent_seq))) @@ -4776,7 +5277,7 @@ class UnifiedRadixCacheSuite: value = self._alloc(allocator, len(seq)) result = cache.insert( InsertParams( - key=RadixKey(seq), + key=RadixKey(array("q", seq)), value=value, swa_evicted_seqlen=ps, ) @@ -5119,15 +5620,17 @@ class UnifiedRadixCacheSuite: # device present -> CoW source is the device value (host backup irrelevant) cache.tree_core.set_component_device_value_raw(node, ComponentType.MAMBA, dev) cache.tree_core.set_component_host_value_raw(node, ComponentType.MAMBA, None) - self.assertIs( - cache.tree_core.get_component_device_value(node, ComponentType.MAMBA), - dev, + device_value = cache.tree_core.get_component_device_value( + node, ComponentType.MAMBA ) + self.assertIsNotNone(device_value) + self.assertTrue(torch.equal(device_value, dev)) cache.tree_core.set_component_host_value_raw(node, ComponentType.MAMBA, host) - self.assertIs( - cache.tree_core.get_component_device_value(node, ComponentType.MAMBA), - dev, + device_value = cache.tree_core.get_component_device_value( + node, ComponentType.MAMBA ) + self.assertIsNotNone(device_value) + self.assertTrue(torch.equal(device_value, dev)) # device evicted -> nothing to CoW from cache.tree_core.set_component_device_value_raw(node, ComponentType.MAMBA, None) self.assertIsNone( @@ -5925,7 +6428,6 @@ class UnifiedRadixCacheSuite: int(kv_xfer.host_indices.numel()), dtype=torch.int64, device=cache.device ) cache.tree_core.commit_load_back(a, device_indices, kv_xfer, {}) - self.assertEqual(cache.tree_core.node_by_id(a).load_back_pending_id, a) # Anchor `b` rejects its whole spec: its SWA window claims pinned `a`. kv_xfer, comp_xfers = cache.tree_core.build_load_back_spec(b) @@ -5935,7 +6437,6 @@ class UnifiedRadixCacheSuite: # After the ack unpins, the same spec builds fully. cache.tree_core.finish_load_back(a) - self.assertIsNone(cache.tree_core.node_by_id(a).load_back_pending_id) kv_xfer, comp_xfers = cache.tree_core.build_load_back_spec(b) self.assertEqual(kv_xfer.nodes_to_load, [b]) self.assertEqual(comp_xfers[ComponentType.SWA][0].nodes_to_load, [a, b]) @@ -6003,6 +6504,8 @@ class UnifiedRadixCacheSuite: last_host_node=leaf, best_match_node=leaf, host_hit_length=0, + cache_protected_len=7, + cache_actions=(FreeDeviceKV(indices=[]),), ) result = cache.tree_core.finalize_component_match_result( ComponentType.SWA, @@ -6015,6 +6518,8 @@ class UnifiedRadixCacheSuite: ) self.assertEqual(result.host_hit_length, 0) self.assertEqual(result.swa_host_hit_length, expected) + self.assertEqual(result.cache_protected_len, 7) + self.assertEqual(result.cache_actions, (FreeDeviceKV(indices=[]),)) def test_hicache_swa_commit_load_back_rebuilds_mapping(self): """LOAD_BACK commit must: @@ -6735,30 +7240,59 @@ class TestMambaCheckpointGrid(CustomTestCase): """ cfg = CacheConfig( - page_size=64, + page_size=32, components=(ComponentType.FULL, ComponentType.MAMBA), enable_mamba_extra_buffer=True, kv_size=1024, max_context_len=1024, ) - def _grid(self, cache): - component = next( - c - for c in cache._components_tuple - if c.component_type is ComponentType.MAMBA + def _branching_seqlen(self, *, tree_page_size: int, full_hit_length: int): + cache, allocator, req_to_token_pool = build_fixture( + self.cfg, + tree_page_size=tree_page_size, + mamba_cache_chunk_size=64, ) - return component.mamba_checkpoint_grid + prefix = list(range(1, tree_page_size + 1)) + tokens = list(range(1, full_hit_length + 1)) + for sequence in (prefix, tokens): + value = allocator.alloc(len(sequence)) + self.assertIsNotNone(value) + req = Req( + rid=f"checkpoint-grid-{len(sequence)}", + origin_input_text="", + origin_input_ids=array("q"), + sampling_params=SamplingParams(temperature=0, max_new_tokens=1), + ) + req_to_token_pool.alloc([req]) + cache.insert( + InsertParams( + key=RadixKey(array("q", sequence)), + value=value, + mamba_value=req.kv.mamba_pool_idx.unsqueeze(0), + ) + ) + + leaf = cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", tokens))) + ).last_device_node + cache.tree_core.set_component_device_value_raw(leaf, ComponentType.MAMBA, None) + result = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens)))) + self.assertEqual(result.full_kv_hit_length, full_hit_length) + self.assertEqual(len(result.device_indices), tree_page_size) + return result.mamba_branching_seqlen def test_grid_follows_the_widened_tree_page(self): - cache, _, _ = build_fixture( - self.cfg, tree_page_size=256, mamba_cache_chunk_size=64 + # lcm(chunk=64, tree page=96) is 192. Chunk-only alignment would + # incorrectly report 256, which is not a radix-node boundary. + self.assertEqual( + self._branching_seqlen(tree_page_size=96, full_hit_length=288), 192 ) - self.assertEqual(self._grid(cache), 256) def test_grid_is_the_chunk_size_without_widening(self): - cache, _, _ = build_fixture(self.cfg, mamba_cache_chunk_size=64) - self.assertEqual(self._grid(cache), 64) + self.assertEqual( + self._branching_seqlen(tree_page_size=32, full_hit_length=160), 128 + ) class TestUnifiedRadixCacheInt8MambaCheckpoint(CustomTestCase): @@ -7292,28 +7826,46 @@ class TestResumableInsertWalk(_InsertWalkSuite): ) step = cache.tree_core.begin_insert(params) self.assertIsNone(step.result) - with self.assertRaises(AssertionError): + with self.assertRaisesRegex(RuntimeError, "concurrent insert walks"): cache.tree_core.begin_insert(params) cache.tree_core.end_insert() + def test_resume_insert_rejects_missing_walk(self): + cache, _, _ = build_fixture(self.cfg) + with self.assertRaisesRegex(RuntimeError, "no in-flight insert"): + cache.tree_core.resume_insert() + def test_insert_abort_drains_pending_deferred_frees(self): """A mid-insert failure after a deferred dup-free accumulated must still return those slots to the allocator via the end_insert drain.""" - cache, allocator, req_to_token_pool = build_fixture(self.cfg) + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + cache.write_through_threshold = sys.maxsize self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4]) + self._insert(cache, allocator, req_to_token_pool, [1, 2]) - # The overlap walk defers a 4-slot dup-free; the commit hook then raises. + (parent,) = _node_children(cache, cache.root_node_handle()) + cache.write_through_threshold = cache.tree_core.get_node_hit_count(parent) + 1 + + # The parent emits a backup barrier. The injected failure advances the + # next overlap step, leaving its duplicate free pending in TreeCore. available = allocator.available_size() - full_comp = cache.components[ComponentType.FULL] + + def fail_after_one_walk_step(): + cache.tree_core.advance_insert_walk_once() + raise RuntimeError("boom") + with mock.patch.object( - full_comp, "commit_insert_component_data", side_effect=RuntimeError("boom") + cache.tree_core, "resume_insert", fail_after_one_walk_step ): with self.assertRaises(RuntimeError): - self._insert(cache, allocator, req_to_token_pool, list(range(1, 9))) + self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4]) - # 8 alloc'd for the insert, 4 dup slots drained back on abort. - self.assertEqual(allocator.available_size(), available - 4) + # All four duplicate slots are returned: two at the barrier, then two + # from end_insert after the injected abort. + self.assertEqual(allocator.available_size(), available) self.assertFalse(cache.tree_core.has_ongoing_insert()) + cache.writing_check() + cache.sanity_check() def test_deferrable_actions_ride_final_step_without_suspension(self): """A walk whose only actions are deferrable frees completes in a single @@ -7447,6 +7999,68 @@ class TestResumableInsertWalkSWA(_InsertWalkSuite): components=(ComponentType.FULL, ComponentType.SWA), sliding_window_size=8 ) + def _recover_swa_tombstone_with_tracking(self, *, lock_full: bool): + cache, allocator, _ = build_fixture(self.cfg) + seq = list(range(1, self.cfg.sliding_window_size + 1)) + key = RadixKey(array("q", seq)) + cache.insert( + InsertParams( + key=key, + value=self._alloc(allocator, len(seq)), + swa_evicted_seqlen=len(seq), + ) + ) + (leaf,) = _node_children(cache, cache.root_node_handle()) + lock_result = cache.inc_lock_ref(leaf) if lock_full else None + try: + result = cache.insert( + InsertParams( + key=key, + value=self._alloc(allocator, len(seq)), + swa_evicted_seqlen=0, + track_adopted_ranges=True, + ) + ) + finally: + if lock_result is not None: + cache.dec_lock_ref(leaf, lock_result.to_dec_params()) + return result, len(seq) + + def test_swa_tombstone_recovery_reports_unlocked_full_adoption(self): + result, seq_len = self._recover_swa_tombstone_with_tracking(lock_full=False) + self.assertEqual(result.adopted_ranges[ComponentType.FULL], [(0, seq_len)]) + self.assertEqual(result.adopted_ranges[ComponentType.SWA], [(0, seq_len)]) + + def test_swa_tombstone_recovery_keeps_locked_full_out_of_adoption(self): + result, seq_len = self._recover_swa_tombstone_with_tracking(lock_full=True) + self.assertNotIn(ComponentType.FULL, result.adopted_ranges) + self.assertEqual(result.adopted_ranges[ComponentType.SWA], [(0, seq_len)]) + + def test_swa_tombstone_recovery_frees_full_only(self): + sw = self.cfg.sliding_window_size + cache, allocator, _ = build_fixture(self.cfg) + seq = list(range(1, 2 * sw + 1)) + key = RadixKey(array("q", seq)) + cache.insert( + InsertParams( + key=key, value=self._alloc(allocator, len(seq)), swa_evicted_seqlen=sw + ) + ) + value = self._alloc(allocator, len(seq)) + full_available = allocator.full_attn_allocator.available_size() + swa_available = allocator.swa_attn_allocator.available_size() + cache.insert(InsertParams(key=key, value=value, swa_evicted_seqlen=0)) + + self.assertEqual( + allocator.full_attn_allocator.available_size(), + full_available + len(seq), + ) + self.assertEqual( + allocator.swa_attn_allocator.available_size(), + swa_available + sw, + ) + cache.sanity_check() + def test_swa_recovery_keeps_recovered_node_below_window_nodes(self): """A tombstone recovered during the walk lands below the in-window path in the SWA LRU, so eviction takes the recovered span first.""" @@ -8209,6 +8823,26 @@ class TestUnifiedRadixCacheStorageAttachBackfill(CustomTestCase): self.assertEqual(cache.tree_core.backfill_missing_hash_values(), 0) self.assertEqual(self._hashes_by_token_ids(cache), before) + def test_buffer_backup_snapshot_is_detached_and_revalidates_parent(self): + cache = self._build_two_level_tree(storage_on_from_the_start=True) + child = cache.match_prefix( + MatchPrefixParams(key=RadixKey(self.prefix_tokens + self.suffix_tokens)) + ).last_device_node + parent = cache.tree_core.get_parent_node_id(child) + snapshot = cache.tree_core.snapshot_buffer_backup(child, pass_prefix_keys=True) + + self.assertEqual(snapshot.key.token_ids, self.suffix_tokens) + snapshot.key.token_ids[0] = -1 + self.assertEqual(_node_token_ids(cache, child), list(self.suffix_tokens)) + + parent_hashes = cache.tree_core.get_hash_values(parent) + replacement_hashes = ["a" * 64] * len(parent_hashes) + cache.tree_core.set_node_hash_values(parent, replacement_hashes) + state = cache.tree_core.validate_buffer_backup(child, len(snapshot.key)) + self.assertEqual(state.parent_node_id, parent) + self.assertEqual(state.parent_last_hash, replacement_hashes[-1]) + cache.tree_core.set_node_hash_values(parent, parent_hashes) + def test_enabling_storage_backfills_the_tree(self): """The tree is hashed by the time `enable_storage` flips on.""" cache = self._build_two_level_tree(storage_on_from_the_start=False) diff --git a/test/registered/unit/mem_cache/unified_tree_core_inspection_interface.py b/test/registered/unit/mem_cache/unified_tree_core_inspection_interface.py index ea2d5ede2..1567b9a15 100644 --- a/test/registered/unit/mem_cache/unified_tree_core_inspection_interface.py +++ b/test/registered/unit/mem_cache/unified_tree_core_inspection_interface.py @@ -205,6 +205,11 @@ class UnifiedTreeCoreInspectionInterface(UnifiedTreeCoreInterface): # ==== Targeted white-box operations ==== + @abstractmethod + def advance_insert_walk_once(self) -> None: + """Advance one suspended insert walk step without flushing its actions.""" + ... + @abstractmethod def evict_component( self, diff --git a/test/registered/unit/mem_cache/unified_tree_core_inspector.py b/test/registered/unit/mem_cache/unified_tree_core_inspector.py index d66494b9a..529b1ee50 100644 --- a/test/registered/unit/mem_cache/unified_tree_core_inspector.py +++ b/test/registered/unit/mem_cache/unified_tree_core_inspector.py @@ -14,6 +14,7 @@ from sglang.srt.mem_cache.unified_cache.components import ComponentType, EvictLa from sglang.srt.mem_cache.unified_cache.unified_tree_core import ( UnifiedLRUList, UnifiedTreeCore, + _InsertPhase, ) from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import ( BaseEvictionResult, @@ -195,6 +196,15 @@ class UnifiedTreeCoreInspector(UnifiedTreeCore, UnifiedTreeCoreInspectionInterfa """Refresh duplicate-host tracking for the node.""" self._update_duplicate_tracking(self.node_by_id(node_id)) + def advance_insert_walk_once(self) -> None: + """Advance one suspended insert walk step without flushing its actions.""" + state = self._ongoing_insert_walk_state + if state is None: + raise RuntimeError("no in-flight insert") + if state.phase is not _InsertPhase.WALK: + raise RuntimeError("in-flight insert is not in walk phase") + self._insert_walk_step(state) + def evict_component( self, node_id: NodeId, diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 803a62343..6272676b7 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -61,7 +61,9 @@ from sglang.srt.arg_groups.serving_hook import ( ssl_verify_of, ) from sglang.srt.arg_groups.speculative_hook import handle_speculative_decoding -from sglang.srt.arg_groups.validation_hook import check_two_batch_overlap +from sglang.srt.arg_groups.validation_hook import ( + check_two_batch_overlap, +) from sglang.srt.entrypoints.sidecar import ( SGLANG_GRPC_ENDPOINT_ENV, Sidecar, @@ -789,6 +791,19 @@ class TestLoadBalanceMethod(unittest.TestCase): "mooncake", ) + def test_pd_decode_hicache_allows_rust_tree_core(self): + server_args = ServerArgs( + model_path="dummy", + disaggregation_mode="decode", + disaggregation_decode_enable_radix_cache=True, + disaggregation_transfer_backend="nixl", + enable_hierarchical_cache=True, + ) + with envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.override("rust"): + handle_pd_disaggregation(server_args) + + self.assertFalse(resolution_result(server_args, "disable_radix_cache")) + class TestSkipTokenizerInit(unittest.TestCase): def test_skip_tokenizer_worker_counts(self): @@ -1475,6 +1490,16 @@ class TestHiCacheArgs(unittest.TestCase): expected_decode_backend, ) + def test_buffer_only_accepts_both_tree_cores(self): + for backend in ("python", "rust"): + args = self._make_args( + enable_hierarchical_cache=True, + hicache_host_memory_mode="buffer_only", + hicache_storage_backend="file", + ) + with envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.override(backend): + handle_hicache(args) + def test_hicache_io_backend_and_mem_layout_compatibility(self): cases = [ {