From dbbd15e6bd1b8e68fc0a3343a3d2fbb7f4d67af7 Mon Sep 17 00:00:00 2001 From: scama Date: Sun, 13 Sep 2026 14:03:45 +1000 Subject: [PATCH 01/16] feat(skippy): consolidate next-generation KV cache stack --- .../manage-ci/references/current-inventory.md | 8 +- .github/actions/plan-ci/action.yml | 7 +- .../action.yml | 5 + .github/workflows/ci-linux-lane.yml | 2 +- .github/workflows/ci-macos-lane.yml | 2 +- .github/workflows/ci-windows-lane.yml | 17 +- .../ci-windows-product-smoke-slice.yml | 41 + .github/workflows/main_windows.yml | 2 + .../workflows/product-integration-smoke.yml | 10 +- .omo/specs/pr-ci-optimization.md | 9 +- Cargo.lock | 1569 ++++++++++- ci/ci.md | 18 +- .../mesh-client/src/client/control_plane.rs | 65 + .../mesh-client/tests/control_plane_client.rs | 52 +- crates/mesh-client/tests/protocol_wire.rs | 2 + crates/mesh-llm-cli/src/lib.rs | 4 +- crates/mesh-llm-cli/src/parser.rs | 74 +- crates/mesh-llm-cli/src/parser/commands.rs | 65 + .../src/gpus/tune/apply_write_tests.rs | 6 +- .../src/gpus/tune/planning.rs | 2 +- .../tune/recommendation_defaults_tests.rs | 33 +- .../src/gpus/tune/recommendation_writes.rs | 20 + crates/mesh-llm-commands/src/kv_cache.rs | 217 ++ crates/mesh-llm-commands/src/lib.rs | 1 + .../src/operational_logging.rs | 1 + .../operational_logging/command_summary.rs | 1 + .../command_summary/dispatch.rs | 3 +- .../command_summary/kv_cache.rs | 49 + .../command_summary_tests.rs | 34 + crates/mesh-llm-config/src/lib.rs | 107 +- crates/mesh-llm-config/src/model.rs | 47 + .../control_behavior/runtime_controls.rs | 25 +- .../src/model/built_in_schema/declarations.rs | 20 + .../src/model/built_in_schema/presentation.rs | 59 +- .../model/built_in_schema/setting_schema.rs | 24 + crates/mesh-llm-config/src/size.rs | 76 + crates/mesh-llm-config/src/validate.rs | 75 + crates/mesh-llm-config/src/wiring_status.rs | 5 + .../src/wiring_status/runtime.rs | 19 + .../command_summary_grammar/descriptors.rs | 3 + .../descriptors/kv_cache.rs | 25 + .../command_summary_grammar/raw_options.rs | 10 + .../src/command_summary_grammar/vocabulary.rs | 3 + crates/mesh-llm-host-runtime/Cargo.toml | 1 + crates/mesh-llm-host-runtime/src/api/mod.rs | 1 + .../src/api/routes/kv_cache.rs | 306 +++ .../src/api/routes/mod.rs | 5 + .../src/api/routes/runtime.rs | 185 +- .../src/api/tests/support.rs | 2 + .../src/inference/skippy/certification.rs | 54 +- .../src/inference/skippy/mod.rs | 1 + .../src/inference/skippy/resolver/tests.rs | 2 +- .../inference/skippy/resolver/translation.rs | 1 + .../src/inference/skippy/resolver/types.rs | 8 +- .../src/inference/skippy/stage/mod.rs | 1 + .../mesh/owner_control/commands/kv_cache.rs | 122 + .../src/mesh/owner_control/commands/mod.rs | 62 +- .../owner_control/commands/model_lifecycle.rs | 1 + .../owner_control/commands/scan_refresh.rs | 1 + .../src/mesh/owner_control/mod.rs | 9 + .../src/mesh/owner_control_response.rs | 1 + .../src/mesh/owner_lifecycle_cache/tests.rs | 1 + .../src/mesh/tests/control_plane_unique.rs | 2 + .../src/models/profile.rs | 8 + .../network/openai/moa_gateway/self_fill.rs | 173 +- .../openai/moa_gateway/self_fill/tests.rs | 48 + .../src/network/openai/routing_rank.rs | 25 + .../src/runtime/config_state.rs | 42 +- .../src/runtime/config_state_tests.rs | 2 + .../src/runtime/config_state_tests/kv_disk.rs | 48 + .../src/runtime/context_planning.rs | 551 +++- .../src/runtime/kv_disk_config.rs | 611 ++++ .../src/runtime/local.rs | 39 + .../src/runtime/local_memory_plan.rs | 259 ++ .../mesh-llm-host-runtime/src/runtime/mod.rs | 2 + .../src/runtime/options.rs | 6 + .../src/runtime/proxy/tests/mod.rs | 3 +- .../src/runtime/run_auto.rs | 16 +- crates/mesh-llm-protocol/proto/node.proto | 22 + crates/mesh-llm-protocol/src/proto/node.rs | 59 + crates/mesh-llm-protocol/src/protocol/mod.rs | 90 + crates/mesh-llm-routing/src/lib.rs | 107 +- crates/mesh-llm-system/src/hardware/tests.rs | 59 +- .../configuration-defaults-runtime.ts | 3 +- crates/mesh-llm/src/commands/mod.rs | 12 +- crates/mesh-llm/src/lib.rs | 3 + .../mesh-llm/tests/protocol_convert_matrix.rs | 2 + crates/skippy-bench/Cargo.toml | 1 + crates/skippy-bench/src/cli.rs | 31 + crates/skippy-bench/src/l2_tier.rs | 309 +++ crates/skippy-bench/src/main.rs | 2 + crates/skippy-cache/Cargo.toml | 24 + crates/skippy-cache/README.md | 10 + .../examples/cachegen_cubecl_spike.rs | 387 +++ crates/skippy-cache/src/cachegen/archive.rs | 1133 ++++++++ crates/skippy-cache/src/cachegen/container.rs | 742 +++++ .../fixtures/generate_lmcache_compat.py | 273 ++ .../fixtures/lmcache_b5d109e_bins16.bin | Bin 0 -> 1231 bytes .../fixtures/lmcache_b5d109e_bins32.bin | Bin 0 -> 1241 bytes .../src/cachegen/fixtures/ryg_rans_golden.bin | Bin 0 -> 137 bytes crates/skippy-cache/src/cachegen/lmcache.rs | 819 ++++++ crates/skippy-cache/src/cachegen/mod.rs | 12 + crates/skippy-cache/src/cachegen/rans.rs | 311 +++ crates/skippy-cache/src/cachegen/reference.rs | 305 ++ crates/skippy-cache/src/fsinfo.rs | 482 ++++ crates/skippy-cache/src/identity.rs | 419 ++- crates/skippy-cache/src/l2/mod.rs | 2445 +++++++++++++++++ crates/skippy-cache/src/l3.rs | 2215 +++++++++++++++ crates/skippy-cache/src/l3/packed.rs | 570 ++++ crates/skippy-cache/src/l3/tests.rs | 1263 +++++++++ crates/skippy-cache/src/l3_remote.rs | 479 ++++ crates/skippy-cache/src/lib.rs | 37 +- crates/skippy-cache/src/manager.rs | 717 +++++ crates/skippy-cache/src/payload/blob_store.rs | 9 +- crates/skippy-cache/src/payload/bytes.rs | 41 +- crates/skippy-cache/src/policy/accounting.rs | 121 + crates/skippy-cache/src/policy/admission.rs | 295 ++ crates/skippy-cache/src/policy/decay.rs | 25 + .../skippy-cache/src/policy/lru_baseline.rs | 60 + crates/skippy-cache/src/policy/mod.rs | 615 +++++ crates/skippy-cache/src/policy/score.rs | 50 + crates/skippy-cache/src/policy/tests.rs | 1156 ++++++++ crates/skippy-cache/src/policy/traces.rs | 137 + crates/skippy-cache/src/radix.rs | 14 +- crates/skippy-cache/src/source.rs | 82 + crates/skippy-cache/src/tier.rs | 1407 ++++++++++ crates/skippy-correctness/Cargo.toml | 2 + crates/skippy-correctness/src/cli.rs | 187 ++ crates/skippy-correctness/src/main.rs | 7 +- crates/skippy-correctness/src/report.rs | 105 +- .../src/runner/cachegen_gate.rs | 606 ++++ .../src/runner/kv_page_growth.rs | 433 +++ crates/skippy-correctness/src/runner/mod.rs | 5 + .../src/runner/remote_handoff.rs | 1833 ++++++++++++ .../src/runner/remote_handoff/identity.rs | 173 ++ .../src/runner/stage_execution.rs | 7 + .../src/runner/state_handoff.rs | 82 +- crates/skippy-ffi/src/abi.rs | 1 + crates/skippy-ffi/src/dynamic.rs | 20 +- crates/skippy-ffi/src/lib.rs | 89 +- crates/skippy-ffi/src/state.rs | 27 + crates/skippy-ffi/src/static_bindings.rs | 8 + crates/skippy-ffi/src/tests.rs | 30 +- .../src/binary/activation_codec.rs | 124 +- crates/skippy-protocol/src/binary/mod.rs | 1 + crates/skippy-protocol/src/lib.rs | 4 +- crates/skippy-protocol/src/validation.rs | 3 + crates/skippy-runtime/Cargo.toml | 1 + crates/skippy-runtime/src/kv_pages.rs | 161 +- crates/skippy-runtime/src/lib.rs | 13 +- crates/skippy-runtime/src/logging.rs | 414 ++- .../src/binary_transport/binary_messaging.rs | 5 +- .../src/binary_transport/options.rs | 3 + .../src/frontend/generation/queue.rs | 11 +- .../src/frontend/generation/server.rs | 6 +- .../token_generation/kv_restore.rs | 14 + .../src/frontend/tests/multimodal.rs | 1 + .../src/kv_integration/config.rs | 609 +++- .../src/kv_integration/exact_state.rs | 375 ++- .../skippy-server/src/kv_integration/mod.rs | 376 ++- .../src/kv_integration/records.rs | 8 + .../src/runtime_state/lane_lifecycle.rs | 18 + .../src/runtime_state/state_transfer.rs | 16 + crates/skippy-topology/src/lib.rs | 2 + crates/skippy-topology/src/phase_placement.rs | 249 ++ docs/CONFIGURATION_PR8_CLOSEOUT_AUDIT.md | 4 +- docs/README.md | 1 + docs/skippy/CACHEGEN_BACKEND_PLAN.md | 370 +++ docs/skippy/CONFIGURATION.md | 9 + docs/skippy/KV_CACHE_DISK.md | 250 ++ docs/skippy/PD_DISAGGREGATION_PLAN.md | 284 ++ docs/skippy/REMOTE_HANDOFF_RUNBOOK.md | 162 ++ ...chegen-lmcache-qwen3-0.6b-19k-summary.json | 76 + ...n-metal-device-qwen3-0.6b-19k-summary.json | 80 + ...en-metal-typed-qwen3-0.6b-19k-summary.json | 228 ++ ...n-quality-gate-qwen3-0.6b-19k-summary.json | 71 + evals/README.md | 44 + evals/agentic-replay.py | 1108 +++++++- evals/kv-restart-replay.py | 682 +++++ evals/test_agentic_replay_l3.py | 129 + scripts/build-llama.sh | 3 + scripts/ci-product-integration-smoke.sh | 156 +- scripts/ci-two-node-split-smoke.sh | 344 ++- scripts/remote-handoff-sweep.sh | 65 + scripts/tests/test_ci_lane_workflows.py | 13 +- .../test_ci_product_integration_smoke.py | 53 +- scripts/tests/test_ci_two_node_split_smoke.py | 2 +- scripts/tests/test_ci_workflow_artifacts.py | 15 +- scripts/tests/test_kv_restart_replay.py | 282 ++ .../tests/test_llama_native_full_replay.py | 9 +- .../tests/test_validate_ci_lane_results.py | 5 +- scripts/validate-ci-lane-results.py | 2 + ...-model-lifecycle-and-package-loading.patch | 23 +- ...r-graph-planning-and-stage-contracts.patch | 15 +- ...ippy-define-CacheGen-page-import-ABI.patch | 87 + ...y-expose-CacheGen-backend-capability.patch | 63 + ...atch-CacheGen-pages-into-resident-KV.patch | 554 ++++ ...code-CacheGen-pages-into-resident-KV.patch | 651 +++++ ...code-CacheGen-pages-into-resident-KV.patch | 481 ++++ ...ml-metal-align-staged-CacheGen-tiles.patch | 117 + ...-optimize-CacheGen-arithmetic-decode.patch | 126 + ...-decode-CacheGen-into-F32-KV-tensors.patch | 559 ++++ ...e-quantized-CacheGen-pages-on-device.patch | 501 ++++ ...3-ggml-metal-stage-CacheGen-directly.patch | 159 ++ ...de-packed-CacheGen-symbols-on-device.patch | 581 ++++ ...uda-stage-CacheGen-payloads-directly.patch | 115 + ...da-use-native-CacheGen-shuffle-masks.patch | 49 + tools/xtask/data/console_print_allowlist.json | 154 +- website/src/_data/docs.js | 1 + website/src/docs/pages/config-defaults.md | 5 +- website/src/docs/pages/config-reference.md | 15 +- website/src/docs/pages/kv-caching.md | 256 ++ .../src/docs/pages/openai-compatible-api.md | 5 + website/src/docs/pages/skippy-api.md | 23 +- 214 files changed, 37316 insertions(+), 471 deletions(-) create mode 100644 .github/workflows/ci-windows-product-smoke-slice.yml create mode 100644 crates/mesh-llm-commands/src/kv_cache.rs create mode 100644 crates/mesh-llm-commands/src/operational_logging/command_summary/kv_cache.rs create mode 100644 crates/mesh-llm-config/src/size.rs create mode 100644 crates/mesh-llm-config/src/wiring_status/runtime.rs create mode 100644 crates/mesh-llm-events/src/command_summary_grammar/descriptors/kv_cache.rs create mode 100644 crates/mesh-llm-host-runtime/src/api/routes/kv_cache.rs create mode 100644 crates/mesh-llm-host-runtime/src/mesh/owner_control/commands/kv_cache.rs create mode 100644 crates/mesh-llm-host-runtime/src/runtime/config_state_tests/kv_disk.rs create mode 100644 crates/mesh-llm-host-runtime/src/runtime/kv_disk_config.rs create mode 100644 crates/mesh-llm-host-runtime/src/runtime/local_memory_plan.rs create mode 100644 crates/skippy-bench/src/l2_tier.rs create mode 100644 crates/skippy-cache/examples/cachegen_cubecl_spike.rs create mode 100644 crates/skippy-cache/src/cachegen/archive.rs create mode 100644 crates/skippy-cache/src/cachegen/container.rs create mode 100644 crates/skippy-cache/src/cachegen/fixtures/generate_lmcache_compat.py create mode 100644 crates/skippy-cache/src/cachegen/fixtures/lmcache_b5d109e_bins16.bin create mode 100644 crates/skippy-cache/src/cachegen/fixtures/lmcache_b5d109e_bins32.bin create mode 100644 crates/skippy-cache/src/cachegen/fixtures/ryg_rans_golden.bin create mode 100644 crates/skippy-cache/src/cachegen/lmcache.rs create mode 100644 crates/skippy-cache/src/cachegen/mod.rs create mode 100644 crates/skippy-cache/src/cachegen/rans.rs create mode 100644 crates/skippy-cache/src/cachegen/reference.rs create mode 100644 crates/skippy-cache/src/fsinfo.rs create mode 100644 crates/skippy-cache/src/l2/mod.rs create mode 100644 crates/skippy-cache/src/l3.rs create mode 100644 crates/skippy-cache/src/l3/packed.rs create mode 100644 crates/skippy-cache/src/l3/tests.rs create mode 100644 crates/skippy-cache/src/l3_remote.rs create mode 100644 crates/skippy-cache/src/manager.rs create mode 100644 crates/skippy-cache/src/policy/accounting.rs create mode 100644 crates/skippy-cache/src/policy/admission.rs create mode 100644 crates/skippy-cache/src/policy/decay.rs create mode 100644 crates/skippy-cache/src/policy/lru_baseline.rs create mode 100644 crates/skippy-cache/src/policy/mod.rs create mode 100644 crates/skippy-cache/src/policy/score.rs create mode 100644 crates/skippy-cache/src/policy/tests.rs create mode 100644 crates/skippy-cache/src/policy/traces.rs create mode 100644 crates/skippy-cache/src/source.rs create mode 100644 crates/skippy-cache/src/tier.rs create mode 100644 crates/skippy-correctness/src/runner/cachegen_gate.rs create mode 100644 crates/skippy-correctness/src/runner/kv_page_growth.rs create mode 100644 crates/skippy-correctness/src/runner/remote_handoff.rs create mode 100644 crates/skippy-correctness/src/runner/remote_handoff/identity.rs create mode 100644 crates/skippy-topology/src/phase_placement.rs create mode 100644 docs/skippy/CACHEGEN_BACKEND_PLAN.md create mode 100644 docs/skippy/KV_CACHE_DISK.md create mode 100644 docs/skippy/PD_DISAGGREGATION_PLAN.md create mode 100644 docs/skippy/REMOTE_HANDOFF_RUNBOOK.md create mode 100644 docs/skippy/cachegen-lmcache-qwen3-0.6b-19k-summary.json create mode 100644 docs/skippy/cachegen-metal-device-qwen3-0.6b-19k-summary.json create mode 100644 docs/skippy/cachegen-metal-typed-qwen3-0.6b-19k-summary.json create mode 100644 docs/skippy/cachegen-quality-gate-qwen3-0.6b-19k-summary.json create mode 100644 evals/kv-restart-replay.py create mode 100644 evals/test_agentic_replay_l3.py create mode 100755 scripts/remote-handoff-sweep.sh create mode 100644 scripts/tests/test_kv_restart_replay.py create mode 100644 third_party/llama.cpp/patches/0034-skippy-define-CacheGen-page-import-ABI.patch create mode 100644 third_party/llama.cpp/patches/0035-skippy-expose-CacheGen-backend-capability.patch create mode 100644 third_party/llama.cpp/patches/0036-skippy-dispatch-CacheGen-pages-into-resident-KV.patch create mode 100644 third_party/llama.cpp/patches/0037-ggml-metal-decode-CacheGen-pages-into-resident-KV.patch create mode 100644 third_party/llama.cpp/patches/0038-ggml-cuda-decode-CacheGen-pages-into-resident-KV.patch create mode 100644 third_party/llama.cpp/patches/0039-ggml-metal-align-staged-CacheGen-tiles.patch create mode 100644 third_party/llama.cpp/patches/0040-ggml-optimize-CacheGen-arithmetic-decode.patch create mode 100644 third_party/llama.cpp/patches/0041-ggml-decode-CacheGen-into-F32-KV-tensors.patch create mode 100644 third_party/llama.cpp/patches/0042-ggml-restore-quantized-CacheGen-pages-on-device.patch create mode 100644 third_party/llama.cpp/patches/0043-ggml-metal-stage-CacheGen-directly.patch create mode 100644 third_party/llama.cpp/patches/0044-ggml-decode-packed-CacheGen-symbols-on-device.patch create mode 100644 third_party/llama.cpp/patches/0045-ggml-cuda-stage-CacheGen-payloads-directly.patch create mode 100644 third_party/llama.cpp/patches/0046-ggml-cuda-use-native-CacheGen-shuffle-masks.patch create mode 100644 website/src/docs/pages/kv-caching.md diff --git a/.agents/skills/manage-ci/references/current-inventory.md b/.agents/skills/manage-ci/references/current-inventory.md index be2862b804..f54c5f729f 100644 --- a/.agents/skills/manage-ci/references/current-inventory.md +++ b/.agents/skills/manage-ci/references/current-inventory.md @@ -219,7 +219,7 @@ it after the protected-main runner-contract update is active. | `ci-website-lane.yml` | Console and website graph; reusable from PRs and dispatchable for main/manual | | `ci-linux-lane.yml` | Linux host/runtime/product/Rust/SDK/smoke graph with one platform-local UI producer | | `ci-macos-lane.yml` | macOS host/runtime/product/platform/Swift/Metal graph with one platform-local UI producer | -| `ci-windows-lane.yml` | Windows host/runtime/product/platform graph with one platform-local UI producer | +| `ci-windows-lane.yml` | Windows host/runtime/product/platform graph plus the CPU durable-L3 product qualification, with one platform-local UI producer | | `ci-quality-slice.yml` | Contracts, format, Clippy and generated CLI inventory freshness; additive protected authority sentinel | | `ci-web-slice.yml` | Console quality, console Playwright E2E, public website build, and CLI explorer browser validation | | `ci-ui-artifact-slice.yml` | Immutable console distribution producer; release callers prepare one source/version-bound UI with complete file checksums, shared by all hosts and SDK resources | @@ -229,7 +229,7 @@ it after the protected-main runner-contract update is active. | `ci-{linux,macos,windows}-runtime-slice.yml` | Platform-pure native runtime producers. The Linux CPU row also runs the native runtime-event gate against the runtime it just built and uploads its evidence. | | `ci-{linux,macos,windows}-product-slice.yml` | Platform-pure composition-only product consumers | | `ci-platform-checks-slice.yml` | macOS portable/unit, Windows portable, and Windows log-store privacy ACL checks | -| `ci-linux-product-smoke-slice.yml`, `ci-macos-product-smoke-slice.yml` | Platform-local callers of the typed CPU/CUDA/Vulkan (`gpu-nvidia` self-hosted), conditional ROCm (`gpu-amd`), and Metal product-integration suite plus model-download. The suite stages the registry-pinned SmolLM2 Q8 and IBM Granite 4.0 H Q4 pair once, runs dense standalone/SDK/restart, then dense passive-client split routing and strict recurrent `KvRecurrent` validation. Each split phase persists strict-whitelist seed/worker node, mesh, and peer identity plus stage/model snapshots, then atomically reconciles exact two-observer, topology/run/model/package/manifest, two-stage contiguous-cut and bind-address, ready-status, and served-model agreement. A capped five-minute wall-clock deadline with parallel, bounded endpoint capture finalizes failure evidence before workflow cancellation; the status projection excludes invite tokens, nested fields, and unrelated paths. Product reconciliation independently verifies both evidence files, records their paths and SHA-256 digests in `phase-results.json`, rejects missing or modified evidence, and uploads every JSON snapshot/evidence file with logs on success or failure. Linux CUDA packages admit only the reviewed cudart, cuBLAS, cuBLASLt, and nvJitLink families for the declared CUDA major, retain NVIDIA object bytes, and include the toolkit distribution license. The Linux CUDA smoke verifies that closure with `LD_LIBRARY_PATH` unset; cudart and cuBLAS are not installed by apt, and the NVIDIA driver remains host-owned. Before inference, it records CUDA visibility variables, host driver-library resolution and NVIDIA device nodes, then runs the packaged benchmark's device-count probe without benchmark allocations, using inherited and strict packaged-library resolution. ROCm skips unless `MESH_ROCM_INFERENCE_RUNNER_ENABLED` is exactly `true`; accelerator product-integration rows remain outside the checked plan pending live qualification. | +| `ci-{linux,macos,windows}-product-smoke-slice.yml` | Platform-local callers of the typed CPU/CUDA/Vulkan (`gpu-nvidia` self-hosted), conditional ROCm (`gpu-amd`), Metal, and Windows CPU product-integration suite plus model-download. The suite stages the registry-pinned SmolLM2 Q8 and IBM Granite 4.0 H Q4 pair once, runs dense standalone/SDK/restart, then dense passive-client split routing and strict recurrent `KvRecurrent` validation. Each split phase persists strict-whitelist seed/worker node, mesh, and peer identity plus stage/model snapshots, then atomically reconciles exact two-observer, topology/run/model/package/manifest, two-stage contiguous-cut and bind-address, ready-status, and served-model agreement. A capped five-minute wall-clock deadline with parallel, bounded endpoint capture finalizes failure evidence before workflow cancellation; the status projection excludes invite tokens, nested fields, and unrelated paths. Product reconciliation independently verifies both evidence files, records their paths and SHA-256 digests in `phase-results.json`, rejects missing or modified evidence, and uploads every JSON snapshot/evidence file with logs on success or failure. Linux CUDA packages admit only the reviewed cudart, cuBLAS, cuBLASLt, and nvJitLink families for the declared CUDA major, retain NVIDIA object bytes, and include the toolkit distribution license. The Linux CUDA smoke verifies that closure with `LD_LIBRARY_PATH` unset; cudart and cuBLAS are not installed by apt, and the NVIDIA driver remains host-owned. Before inference, it records CUDA visibility variables, host driver-library resolution and NVIDIA device nodes, then runs the packaged benchmark's device-count probe without benchmark allocations, using inherited and strict packaged-library resolution. A separately reconciled durable-L3 phase preserves per-node roots and identities across full process restarts for both model families, requires active CLI-sourced disk configuration, persisted inventory, a post-restart L3 fill with cached tokens and exact output, then verifies status and clear. Windows CPU runs the durable phase alone on `windows-2022`; Linux CPU and macOS Metal run the complete suite. ROCm skips unless `MESH_ROCM_INFERENCE_RUNNER_ENABLED` is exactly `true`; accelerator product-integration rows remain outside the checked plan pending live qualification. | | `ci-linux-sdk-slice.yml`, `ci-macos-sdk-slice.yml` | Platform-local Rust/Kotlin/Swift smoke consumers; SDK producers are independent top-level calls and each smoke receives the lane-local immutable UI artifact | | `ci-runner-contract-slice.yml` | Provider/cache/plan trust and main runner-image checks | | `native-sdk-artifact.yml` | Typed native SDK producer | @@ -384,7 +384,7 @@ does not grant; GitHub rejects at run creation with a **zero-job `actionlint` cannot see it. Containerizing surfaced this because `packages: read` (needed to pull the private GHCR runner images) has to be granted at *every* hop, and -`ci-linux-product-smoke-slice.yml` / `ci-macos-product-smoke-slice.yml` sat at +the platform product-smoke slices sat at `contents: read` between granted parents and requesting children. `scripts/tests/test_ci_workflow_permission_contract.py` walks every local `uses: ./.github/workflows/X.yml` edge and asserts the caller's effective @@ -786,7 +786,7 @@ Current image references and historical null evidence remain unchanged. The `product-smoke` catalog role covers both the legacy `smoke.yml` job and the typed `product-integration-smoke.yml` job. The latter uses the same pinned -CPU image only for Linux CPU; accelerator and macOS paths retain their existing +CPU image only for Linux CPU; accelerator, macOS, and Windows paths retain their container opt-outs. The inventory has 9 images, 32 roles and 33 literal workflow image bindings. ### Qualified lean UI consumers diff --git a/.github/actions/plan-ci/action.yml b/.github/actions/plan-ci/action.yml index 6447c0ed76..7c60f81263 100644 --- a/.github/actions/plan-ci/action.yml +++ b/.github/actions/plan-ci/action.yml @@ -284,7 +284,7 @@ runs: rust_tests: .matrices.rust_tests, hosts: [.matrices.hosts[] | select(.platform == "linux")], runtime_products: [.matrices.runtime_products[] | select(.platform == "linux")], - smoke: [.matrices.smoke[] | select(.id != "metal-model-load")], + smoke: [.matrices.smoke[] | select(.id != "metal-model-load" and .id != "product-integration-metal" and .id != "product-integration-windows-cpu")], sdk: [.matrices.sdk[] | select(.platform == "linux")] } } @@ -309,7 +309,7 @@ runs: hosts: [.matrices.hosts[] | select(.platform == "macos")], runtime_products: [.matrices.runtime_products[] | select(.platform == "macos")], platform_checks: [.matrices.platform_checks[] | select(.platform == "macos")], - smoke: [.matrices.smoke[] | select(.id == "metal-model-load")], + smoke: [.matrices.smoke[] | select(.id == "metal-model-load" or .id == "product-integration-metal")], sdk: [.matrices.sdk[] | select(.platform == "macos")] } } @@ -326,7 +326,8 @@ runs: matrices: { hosts: [.matrices.hosts[] | select(.platform == "windows")], runtime_products: [.matrices.runtime_products[] | select(.platform == "windows")], - platform_checks: [.matrices.platform_checks[] | select(.platform == "windows")] + platform_checks: [.matrices.platform_checks[] | select(.platform == "windows")], + smoke: [.matrices.smoke[] | select(.id == "product-integration-windows-cpu")] } } | .required = ([.matrices[] | length] | add) > 0 diff --git a/.github/actions/restore-product-integration-inputs/action.yml b/.github/actions/restore-product-integration-inputs/action.yml index f39cf5fc15..a3f77dcdea 100644 --- a/.github/actions/restore-product-integration-inputs/action.yml +++ b/.github/actions/restore-product-integration-inputs/action.yml @@ -11,6 +11,10 @@ inputs: staged_binary_path: required: true description: Stable executable path used by the product suite. + binary_name: + required: false + default: mesh-llm + description: Host executable filename inside the composed product. model_cadence: required: true description: Authorized registry cadence for this invocation. @@ -50,6 +54,7 @@ runs: with: artifact_name: ${{ inputs.artifact_name }} artifact_path: ${{ inputs.artifact_path }} + binary_name: ${{ inputs.binary_name }} staged_binary_path: ${{ inputs.staged_binary_path }} - name: Resolve pinned product-integration fixtures diff --git a/.github/workflows/ci-linux-lane.yml b/.github/workflows/ci-linux-lane.yml index bfc9c53876..c97fb008f8 100644 --- a/.github/workflows/ci-linux-lane.yml +++ b/.github/workflows/ci-linux-lane.yml @@ -179,7 +179,7 @@ jobs: source_sha: ${{ inputs.source_sha }} smoke_matrix: ${{ toJson(fromJson(inputs.lane_plan_json).matrices.smoke) }} binary_target: target/release/mesh-llm - timeout_minutes: 30 + timeout_minutes: 45 secrets: HF_TOKEN: ${{ inputs.original_event_name == 'push' && secrets.HF_TOKEN || '' }} diff --git a/.github/workflows/ci-macos-lane.yml b/.github/workflows/ci-macos-lane.yml index dcc35d9923..6206afc6ac 100644 --- a/.github/workflows/ci-macos-lane.yml +++ b/.github/workflows/ci-macos-lane.yml @@ -170,7 +170,7 @@ jobs: architecture: ${{ fromJson(inputs.lane_plan_json).matrices.runtime_products[0].architecture }} smoke_matrix: ${{ toJson(fromJson(inputs.lane_plan_json).matrices.smoke) }} binary_target: target/release/mesh-llm - timeout_minutes: 30 + timeout_minutes: 45 summary: name: CI / macOS diff --git a/.github/workflows/ci-windows-lane.yml b/.github/workflows/ci-windows-lane.yml index e94a2938a7..3fe82a9e5f 100644 --- a/.github/workflows/ci-windows-lane.yml +++ b/.github/workflows/ci-windows-lane.yml @@ -36,6 +36,9 @@ on: required: false default: "[]" type: string + secrets: + HF_TOKEN: + required: false workflow_dispatch: inputs: *lane_inputs permissions: @@ -98,6 +101,18 @@ jobs: max_parallel: ${{ fromJson(inputs.lane_plan_json).budgets.windows_max_parallel }} fail_fast: ${{ inputs.original_event_name == 'pull_request' }} + product_smoke: + name: Windows product smoke + needs: [runtime_product] + if: ${{ !cancelled() && needs.runtime_product.result == 'success' && fromJson(inputs.lane_plan_json).matrices.smoke[0] != null }} + uses: ./.github/workflows/ci-windows-product-smoke-slice.yml + with: + source_sha: ${{ inputs.source_sha }} + smoke_matrix: ${{ toJson(fromJson(inputs.lane_plan_json).matrices.smoke) }} + timeout_minutes: 45 + secrets: + HF_TOKEN: ${{ secrets.HF_TOKEN }} + platform_checks: name: Windows checks if: ${{ fromJson(inputs.lane_plan_json).matrices.platform_checks[0] != null }} @@ -113,7 +128,7 @@ jobs: summary: name: CI / Windows - needs: [ui_artifact, hosts, native_runtimes, runtime_product, platform_checks] + needs: [ui_artifact, hosts, native_runtimes, runtime_product, product_smoke, platform_checks] if: ${{ !cancelled() }} runs-on: ubuntu-24.04 permissions: diff --git a/.github/workflows/ci-windows-product-smoke-slice.yml b/.github/workflows/ci-windows-product-smoke-slice.yml new file mode 100644 index 0000000000..e86c39945e --- /dev/null +++ b/.github/workflows/ci-windows-product-smoke-slice.yml @@ -0,0 +1,41 @@ +name: CI Windows Product Smoke Slice + +on: + workflow_call: + inputs: + source_sha: + description: Immutable product source revision. + required: true + type: string + smoke_matrix: + description: JSON Windows smoke rows from the versioned CI plan. + required: true + type: string + timeout_minutes: + description: Bounded smoke timeout. + required: false + default: 45 + type: number + secrets: + HF_TOKEN: + required: false + +permissions: + contents: read + packages: read + +jobs: + product_integration_cpu: + name: Windows CPU durable L3 qualification + if: ${{ contains(fromJson(inputs.smoke_matrix).*.id, 'product-integration-windows-cpu') }} + uses: ./.github/workflows/product-integration-smoke.yml + with: + source_sha: ${{ inputs.source_sha }} + artifact_name: ci-product-windows-amd64-cpu + artifact_path: ci-artifacts/windows + staged_binary_path: ci-artifacts/windows/mesh-llm.exe + platform: windows + backend: cpu + timeout_minutes: ${{ inputs.timeout_minutes }} + secrets: + HF_TOKEN: ${{ secrets.HF_TOKEN }} diff --git a/.github/workflows/main_windows.yml b/.github/workflows/main_windows.yml index 39dc952111..648888efd8 100644 --- a/.github/workflows/main_windows.yml +++ b/.github/workflows/main_windows.yml @@ -78,6 +78,8 @@ jobs: source_sha: ${{ needs.plan.outputs.source_sha }} original_event_name: push supersession_key: main-${{ github.sha }} + secrets: + HF_TOKEN: ${{ secrets.HF_TOKEN }} required: name: Main / Windows diff --git a/.github/workflows/product-integration-smoke.yml b/.github/workflows/product-integration-smoke.yml index 9e36278a25..63f91ace66 100644 --- a/.github/workflows/product-integration-smoke.yml +++ b/.github/workflows/product-integration-smoke.yml @@ -16,7 +16,7 @@ on: required: true type: string platform: - description: "Typed operating-system selector: linux or macos." + description: "Typed operating-system selector: linux, macos, or windows." required: true type: string backend: @@ -39,7 +39,7 @@ jobs: product_integration: name: Product integration (${{ inputs.platform }}/${{ inputs.backend }}) if: ${{ inputs.backend != 'rocm' || vars.MESH_ROCM_INFERENCE_RUNNER_ENABLED == 'true' }} - runs-on: ${{ inputs.platform == 'macos' && 'macos-15' || (inputs.backend == 'cuda' || inputs.backend == 'vulkan') && fromJSON('["self-hosted","Linux","X64","amd64","gpu-nvidia","mesh-llm-amd64","mesh-llm"]') || inputs.backend == 'rocm' && fromJSON('["self-hosted","Linux","X64","amd64","gpu-amd","mesh-llm-amd64","mesh-llm"]') || 'ubuntu-24.04' }} + runs-on: ${{ inputs.platform == 'windows' && 'windows-2022' || inputs.platform == 'macos' && 'macos-15' || (inputs.backend == 'cuda' || inputs.backend == 'vulkan') && fromJSON('["self-hosted","Linux","X64","amd64","gpu-nvidia","mesh-llm-amd64","mesh-llm"]') || inputs.backend == 'rocm' && fromJSON('["self-hosted","Linux","X64","amd64","gpu-amd","mesh-llm-amd64","mesh-llm"]') || 'ubuntu-24.04' }} timeout-minutes: ${{ inputs.timeout_minutes }} container: image: ${{ inputs.platform == 'linux' && inputs.backend == 'cpu' && 'ghcr.io/mesh-llm/mesh-llm-cuda-runner@sha256:8d93de6ba30173e825a16fdecf011f9c632edc6e1259df7289e491b0a05f829d' || '' }} @@ -58,7 +58,7 @@ jobs: run: | set -euo pipefail case "${PLATFORM}/${BACKEND}" in - linux/cpu|linux/cuda|linux/vulkan|linux/rocm|macos/metal) ;; + linux/cpu|linux/cuda|linux/vulkan|linux/rocm|macos/metal|windows/cpu) ;; *) echo "unsupported product integration platform/backend: ${PLATFORM}/${BACKEND}" >&2; exit 2 ;; esac @@ -121,6 +121,7 @@ jobs: with: artifact_name: ${{ inputs.artifact_name }} artifact_path: ${{ inputs.artifact_path }} + binary_name: ${{ inputs.platform == 'windows' && 'mesh-llm.exe' || 'mesh-llm' }} staged_binary_path: ${{ inputs.staged_binary_path }} model_cadence: ${{ (github.event_name == 'pull_request' || github.event_name == 'pull_request_target' || github.event.inputs.original_event_name == 'pull_request' || github.event.inputs.original_event_name == 'pull_request_target') && 'pull-request' || 'main' }} save_model_cache: ${{ github.ref == 'refs/heads/main' && github.event.inputs.original_event_name != 'pull_request' && github.event.inputs.original_event_name != 'pull_request_target' }} @@ -128,6 +129,7 @@ jobs: - name: Run typed product integration suite env: MESH_PRODUCT_INTEGRATION_PHASE_ROOT: ${{ runner.temp }}/product-integration-${{ inputs.platform }}-${{ inputs.backend }} + MESH_PRODUCT_INTEGRATION_DURABLE_ONLY: ${{ inputs.platform == 'windows' && '1' || '0' }} run: | scripts/ci-product-integration-smoke.sh \ "${{ inputs.staged_binary_path }}" \ @@ -150,5 +152,7 @@ jobs: ${{ runner.temp }}/product-integration-${{ inputs.platform }}-${{ inputs.backend }}/phase-results.json ${{ runner.temp }}/product-integration-${{ inputs.platform }}-${{ inputs.backend }}/*/split-evidence.json ${{ runner.temp }}/product-integration-${{ inputs.platform }}-${{ inputs.backend }}/*/split-evidence-snapshots/*.json + ${{ runner.temp }}/product-integration-${{ inputs.platform }}-${{ inputs.backend }}/*/durable-l3-evidence.json + ${{ runner.temp }}/product-integration-${{ inputs.platform }}-${{ inputs.backend }}/*/durable-*/*.json ${{ runner.temp }}/product-integration-${{ inputs.platform }}-${{ inputs.backend }}/*/*.log if-no-files-found: error diff --git a/.omo/specs/pr-ci-optimization.md b/.omo/specs/pr-ci-optimization.md index c75af9442e..460749fbc9 100644 --- a/.omo/specs/pr-ci-optimization.md +++ b/.omo/specs/pr-ci-optimization.md @@ -161,10 +161,11 @@ credentials may differ. - ci-{linux,macos,windows}-product-slice.yml: platform-local composition after matching host and runtime producers succeed. - ci-platform-checks-slice.yml: macOS portable/unit and Windows checks. -- ci-linux-product-smoke-slice.yml and ci-macos-product-smoke-slice.yml: - platform-local inference, backend, two-node, Metal and model-download - consumers using only composed artifacts. One Linux KV caching smoke job runs - a fixed dense SmolLM2 leg followed by a recurrent Qwen3.5 leg; both must pass. +- ci-{linux,macos,windows}-product-smoke-slice.yml: platform-local inference, + backend, two-node, Metal, Windows CPU and model-download consumers using only + composed artifacts. Product integration includes a digest-bound durable-L3 + phase for dense and recurrent models across a full process restart; Windows + runs that phase alone on a real product executor. - ci-linux-sdk-slice.yml and ci-macos-sdk-slice.yml: platform-local Rust/Kotlin/Swift consumers. Swift and Kotlin SDK artifacts are independent producers that start from the plan and static ABI respectively, before diff --git a/Cargo.lock b/Cargo.lock index 8d88f0e7a0..74a2ff3979 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,15 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "addr2line" +version = "0.25.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b5d307320b3181d6d7954e663bd7c774a838b8220fe0593c86d9fb09f498b4b" +dependencies = [ + "gimli", +] + [[package]] name = "adler2" version = "2.0.1" @@ -237,6 +246,15 @@ version = "0.7.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" +[[package]] +name = "ash" +version = "0.38.0+1.3.281" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bb44936d800fea8f016d7f2311c6a4f97aebd5dc86f09906139ec848cf3a46f" +dependencies = [ + "libloading 0.8.9", +] + [[package]] name = "askama" version = "0.16.0" @@ -262,7 +280,7 @@ dependencies = [ "memchr", "proc-macro2", "quote", - "rustc-hash", + "rustc-hash 2.1.3", "serde", "serde_derive", "syn 2.0.119", @@ -283,7 +301,7 @@ version = "0.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7db09fde9143e7ac4513358fb32ee32847125b63b18ea715afd487956da715da" dependencies = [ - "rustc-hash", + "rustc-hash 2.1.3", "serde", "serde_derive", "unicode-ident", @@ -622,6 +640,21 @@ dependencies = [ "tokio", ] +[[package]] +name = "backtrace" +version = "0.3.76" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb531853791a215d7c62a30daf0dde835f381ab5de4589cfe7c649d2cbe92bd6" +dependencies = [ + "addr2line", + "cfg-if 1.0.4", + "libc", + "miniz_oxide 0.8.9", + "object", + "rustc-demangle", + "windows-link", +] + [[package]] name = "base16ct" version = "1.0.0" @@ -661,6 +694,26 @@ version = "0.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "efbd3e1070bbdf4cd88a75264e18e8a26f7cb5c6949eadf0ceb85fb159cf08f8" +[[package]] +name = "bindgen" +version = "0.71.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f58bf3d7db68cfbac37cfc485a8d711e87e064c3d0fe0435b92f7a407f9d6b3" +dependencies = [ + "bitflags 2.13.1", + "cexpr", + "clang-sys", + "itertools 0.13.0", + "log", + "prettyplease 0.2.37", + "proc-macro2", + "quote", + "regex", + "rustc-hash 2.1.3", + "shlex 1.3.0", + "syn 2.0.119", +] + [[package]] name = "bip39" version = "2.2.2" @@ -681,6 +734,15 @@ dependencies = [ "bit-vec 0.6.3", ] +[[package]] +name = "bit-set" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34ddef2995421ab6a5c779542c81ee77c115206f4ad9d5a8e05f4ff49716a3dd" +dependencies = [ + "bit-vec 0.9.1", +] + [[package]] name = "bit-vec" version = "0.6.3" @@ -755,6 +817,9 @@ name = "bitflags" version = "2.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" +dependencies = [ + "serde_core", +] [[package]] name = "blake2" @@ -851,7 +916,7 @@ version = "3.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d4654961ad0494e4774c5c60b4cb4cd0ae9b9d92d039d901638b1dba97ebebf5" dependencies = [ - "darling", + "darling 0.24.1", "ident_case", "prettyplease 0.3.0", "proc-macro2", @@ -887,6 +952,20 @@ name = "bytemuck" version = "1.25.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0e56a716f1e132ff6bf4bdac1c944a3fcdc1cae65f70a4a2a1ac3b401d2d1f" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.4", +] [[package]] name = "byteorder" @@ -939,6 +1018,15 @@ dependencies = [ "thiserror 2.0.20", ] +[[package]] +name = "caseless" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b6fd507454086c8edfd769ca6ada439193cdb209c7681712ef6275cccbfe5d8" +dependencies = [ + "unicode-normalization", +] + [[package]] name = "castaway" version = "0.2.4" @@ -975,7 +1063,7 @@ dependencies = [ "find-msvc-tools", "jobserver", "libc", - "shlex", + "shlex 2.0.1", ] [[package]] @@ -984,6 +1072,15 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6d43a04d8753f35258c91f8ec639f792891f748a1edbd759cf1dcea3382ad83c" +[[package]] +name = "cexpr" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" +dependencies = [ + "nom 7.1.3", +] + [[package]] name = "cfg-if" version = "0.1.10" @@ -1064,6 +1161,33 @@ dependencies = [ "windows-link", ] +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + [[package]] name = "cipher" version = "0.4.4" @@ -1086,6 +1210,17 @@ dependencies = [ "inout 0.2.2", ] +[[package]] +name = "clang-sys" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "157a8ba7b480713b56f4c09fd13fc3e0a22a5dfab8097ba61cbc5feef950788a" +dependencies = [ + "glob", + "libc", + "libloading 0.8.9", +] + [[package]] name = "clap" version = "4.6.6" @@ -1159,6 +1294,17 @@ dependencies = [ "thiserror 2.0.20", ] +[[package]] +name = "codespan-reporting" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af491d569909a7e4dee0ad7db7f5341fef5c614d5b8ec8cf765732aba3cff681" +dependencies = [ + "serde", + "termcolor", + "unicode-width", +] + [[package]] name = "colorchoice" version = "1.0.5" @@ -1198,6 +1344,20 @@ dependencies = [ "static_assertions", ] +[[package]] +name = "comrak" +version = "0.39.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fefab951771fc3beeed0773ce66a4f7b706273fc6c4c95b08dd1615744abcf5" +dependencies = [ + "caseless", + "entities", + "memchr", + "slug", + "typed-arena", + "unicode_categories", +] + [[package]] name = "concurrent-queue" version = "2.5.0" @@ -1246,6 +1406,21 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" +[[package]] +name = "constcat" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "136d3e02915a2cea4d74caa8681e2d44b1c3254bdbf17d11d41d587ff858832c" + +[[package]] +name = "convert_case" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baaaa0ecca5b51987b9423ccdc971514dd8b0bb7b4060b983d3664dad3f1f89f" +dependencies = [ + "unicode-segmentation", +] + [[package]] name = "convert_case" version = "0.10.0" @@ -1535,6 +1710,336 @@ dependencies = [ "cmov", ] +[[package]] +name = "cubecl" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fd203fef6e359472e4cb8f6c0ef3eca062815a1edd560eb6d6c1d5c1397838d9" +dependencies = [ + "cubecl-core", + "cubecl-cpu", + "cubecl-cuda", + "cubecl-hip", + "cubecl-ir", + "cubecl-runtime", + "cubecl-std", + "cubecl-wgpu", + "half", +] + +[[package]] +name = "cubecl-common" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc956c2dcc993f16f748d03bfdbd900f75cab4d469f4e5d8924f8ee128e861ad" +dependencies = [ + "backtrace", + "bytemuck", + "cfg-if 1.0.4", + "cfg_aliases", + "ciborium", + "derive-new", + "derive_more", + "dirs", + "embassy-futures", + "embassy-time", + "float4", + "float8", + "futures-lite", + "half", + "hashbrown 0.16.1", + "log", + "num-traits", + "oneshot 0.2.1", + "parking_lot", + "portable-atomic", + "portable-atomic-util", + "rand 0.10.2", + "sanitize-filename", + "serde", + "serde_bytes", + "serde_json", + "spin 0.10.1", + "toml 1.1.5+spec-1.1.0", + "tynm", + "wasm-bindgen-futures", + "web-time", + "xxhash-rust", +] + +[[package]] +name = "cubecl-core" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7522e7acf25d7848032c7b5270780f9f2abe84e13c7ed6345aa27ba70c312ca" +dependencies = [ + "bitflags 2.13.1", + "bytemuck", + "cubecl-common", + "cubecl-ir", + "cubecl-macros", + "cubecl-runtime", + "cubecl-zspace", + "derive-new", + "derive_more", + "enumset", + "float-ord", + "half", + "hashbrown 0.16.1", + "log", + "num-traits", + "paste", + "serde", + "serde_json", + "variadics_please", +] + +[[package]] +name = "cubecl-cpp" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "411af04828cbf4583ecd388579fb30cd7a644eec19a334bb8f0d9d08522e1458" +dependencies = [ + "bytemuck", + "cubecl-common", + "cubecl-core", + "cubecl-opt", + "cubecl-runtime", + "derive-new", + "half", + "itertools 0.14.0", + "log", +] + +[[package]] +name = "cubecl-cpu" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f572143f54e42a49ee0635a4ec36005f639e62be029e4f49890c808985981b35" +dependencies = [ + "bytemuck", + "cubecl-common", + "cubecl-core", + "cubecl-opt", + "cubecl-runtime", + "cubecl-std", + "derive-new", + "half", + "log", + "paste", + "serde", + "sysinfo", + "tracel-llvm", + "tracel-llvm-bundler", +] + +[[package]] +name = "cubecl-cuda" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6b0a69ff45688d322ad8e92c8bf645167b9ca490fa8fa087fc6adac8c5e46be" +dependencies = [ + "bytemuck", + "cubecl-common", + "cubecl-core", + "cubecl-cpp", + "cubecl-runtime", + "cudarc", + "derive-new", + "half", + "log", + "serde", +] + +[[package]] +name = "cubecl-hip" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c6b510a9348f06ecd56b32cf10eb6241639e927a87f5c09ec61cc7a7610d5a3" +dependencies = [ + "bytemuck", + "cubecl-common", + "cubecl-core", + "cubecl-cpp", + "cubecl-hip-sys", + "cubecl-runtime", + "derive-new", + "half", + "log", + "paste", + "serde", +] + +[[package]] +name = "cubecl-hip-sys" +version = "7.14.6085000" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "760c605ca1b54d12ba9209d2a47992082ff66346fbd96d830af22ac43adfeebc" +dependencies = [ + "libc", + "regex", +] + +[[package]] +name = "cubecl-ir" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d28a897a40c8ee6c57a8b8d76d8823ba2e97f4c2b83db9338f17a0752132d486" +dependencies = [ + "cubecl-common", + "cubecl-macros-internal", + "derive-new", + "derive_more", + "enumset", + "float-ord", + "fnv", + "foldhash 0.2.0", + "half", + "hashbrown 0.16.1", + "num-traits", + "portable-atomic", + "serde", + "variadics_please", +] + +[[package]] +name = "cubecl-macros" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77aa6df563e8e6a0926d2e9eeacb968737940a0e1ae9be819f6a65a341006e12" +dependencies = [ + "cubecl-common", + "darling 0.23.0", + "derive-new", + "ident_case", + "inflections", + "prettyplease 0.2.37", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cubecl-macros-internal" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a515651a5a91e25d87f71f311f80a088e6cf815798b9922560a97636da6b8740" +dependencies = [ + "darling 0.23.0", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cubecl-opt" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0a599c6c3efefdcee8636994c90e58ef696c7efbed2d18b5b5ad8d5f29923abb" +dependencies = [ + "cubecl-common", + "cubecl-core", + "cubecl-ir", + "float-ord", + "log", + "num", + "petgraph", + "smallvec", + "stable-vec", + "type-map", +] + +[[package]] +name = "cubecl-runtime" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b68491bf5b3e997ae36bdc4e63b4ccd6d2f0e86b3b596a5d7a48d2b9e92622a0" +dependencies = [ + "ahash", + "async-channel", + "bytemuck", + "cfg-if 1.0.4", + "cfg_aliases", + "cubecl-common", + "cubecl-ir", + "cubecl-zspace", + "derive-new", + "derive_more", + "dirs", + "enumset", + "hashbrown 0.16.1", + "log", + "md5", + "serde", + "serde_json", + "spin 0.10.1", + "thiserror 2.0.20", + "toml 1.1.5+spec-1.1.0", + "wasm-bindgen-futures", + "web-time", +] + +[[package]] +name = "cubecl-std" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b391b584a4897683dd9fc0767f96337ac712b733126ce0f68a9ee8a2df073d2d" +dependencies = [ + "cubecl-common", + "cubecl-core", + "cubecl-runtime", + "half", + "num-traits", + "paste", + "serde", + "spin 0.10.1", + "variadics_please", +] + +[[package]] +name = "cubecl-wgpu" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4d3663c6e1a86187172b2cee495b6b533e7d91a2cb37e26f421649e315fe4c3a" +dependencies = [ + "async-channel", + "bytemuck", + "cfg-if 1.0.4", + "cfg_aliases", + "cubecl-common", + "cubecl-core", + "cubecl-ir", + "cubecl-runtime", + "derive-new", + "derive_more", + "half", + "hashbrown 0.16.1", + "log", + "sanitize-filename", + "wasm-bindgen-futures", + "wgpu", +] + +[[package]] +name = "cubecl-zspace" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04054d95698afab785a4dda9f949e297831dd9e4cc6d036a93e6603f88e88b07" +dependencies = [ + "derive-new", + "serde", + "smallvec", +] + +[[package]] +name = "cudarc" +version = "0.19.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "804764d10e844da09765a7b2ca9641a0851523d1702efb0d7299d73e31b86e80" +dependencies = [ + "libloading 0.9.0", +] + [[package]] name = "curve25519-dalek" version = "4.1.3" @@ -1579,14 +2084,60 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "darling" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9cdf337090841a411e2a7f3deb9187445851f91b309c0c0a29e05f74a00a48c0" +dependencies = [ + "darling_core 0.21.3", + "darling_macro 0.21.3", +] + +[[package]] +name = "darling" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d" +dependencies = [ + "darling_core 0.23.0", + "darling_macro 0.23.0", +] + [[package]] name = "darling" version = "0.24.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed17f5901b6630b993ca003def43f2f8ef4014fc13b047b57aad617ff32bc2ec" dependencies = [ - "darling_core", - "darling_macro", + "darling_core 0.24.1", + "darling_macro 0.24.1", +] + +[[package]] +name = "darling_core" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1247195ecd7e3c85f83c8d2a366e4210d588e802133e1e355180a9870b517ea4" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "darling_core" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9865a50f7c335f53564bb694ef660825eb8610e0a53d3e11bf1b0d3df31e03b0" +dependencies = [ + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn 2.0.119", ] [[package]] @@ -1602,13 +2153,35 @@ dependencies = [ "syn 3.0.4", ] +[[package]] +name = "darling_macro" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d38308df82d1080de0afee5d069fa14b0326a88c14f15c5ccda35b4a6c414c81" +dependencies = [ + "darling_core 0.21.3", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "darling_macro" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" +dependencies = [ + "darling_core 0.23.0", + "quote", + "syn 2.0.119", +] + [[package]] name = "darling_macro" version = "0.24.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2ac7135c3ef02b2f7833bbeb1be5ba7f966dcde8a87c6b87f65a778d71a02785" dependencies = [ - "darling_core", + "darling_core 0.24.1", "quote", "syn 3.0.4", ] @@ -1690,6 +2263,17 @@ version = "0.5.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" +[[package]] +name = "derive-new" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2cdc8d50f426189eef89dac62fabfa0abb27d5cc008f25bf4156a0203325becc" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "derive_more" version = "2.1.1" @@ -1713,6 +2297,12 @@ dependencies = [ "unicode-xid", ] +[[package]] +name = "deunicode" +version = "1.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "abd57806937c9cc163efc8ea3910e00a62e2aeb0b8119f1793a978088f8f6b04" + [[package]] name = "diatomic-waker" version = "0.2.3" @@ -1785,6 +2375,15 @@ dependencies = [ "syn 3.0.4", ] +[[package]] +name = "dlib" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab8ecd87370524b461f8557c119c405552c396ed91fc0a8eec68679eab26f94a" +dependencies = [ + "libloading 0.8.9", +] + [[package]] name = "dlopen2" version = "0.8.2" @@ -1850,6 +2449,62 @@ version = "1.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "252afb9ae5eaa683babdc6a068b3f5726eb19e05070c731f9b2a23a7c3e8ed34" +[[package]] +name = "embassy-futures" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc2d050bdc5c21e0862a89256ed8029ae6c290a93aecefc73084b3002cdebb01" + +[[package]] +name = "embassy-time" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "592b0c143ec626e821d4d90da51a2bd91d559d6c442b7c74a47d368c9e23d97a" +dependencies = [ + "cfg-if 1.0.4", + "critical-section", + "document-features", + "embassy-time-driver", + "embedded-hal 0.2.7", + "embedded-hal 1.0.0", + "embedded-hal-async", + "futures-core", +] + +[[package]] +name = "embassy-time-driver" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ee71af1b3a0deaa53eaf2d39252f83504c853646e472400b763060389b9fcc9" +dependencies = [ + "document-features", +] + +[[package]] +name = "embedded-hal" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35949884794ad573cf46071e41c9b60efb0cb311e3ca01f7af807af1debc66ff" +dependencies = [ + "nb 0.1.3", + "void", +] + +[[package]] +name = "embedded-hal" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "361a90feb7004eca4019fb28352a9465666b24f840f5c3cddf0ff13920590b89" + +[[package]] +name = "embedded-hal-async" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c4c685bbef7fe13c3c6dd4da26841ed3980ef33e841cddfa15ce8a8fb3f1884" +dependencies = [ + "embedded-hal 1.0.0", +] + [[package]] name = "embedded-io" version = "0.4.0" @@ -1883,6 +2538,12 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "869b0adbda23651a9c5c0c3d270aac9fcb52e8622a8f2b17e57802d7791962f2" +[[package]] +name = "entities" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5320ae4c3782150d900b79807611a59a99fc9a1d61d686faafc24b93fc8d7ca" + [[package]] name = "enum-assoc" version = "1.4.1" @@ -1915,6 +2576,28 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "enumset" +version = "1.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccc5801fd11762e24d1e420d01d2ac518f2a2ca4329d4fbb6639f2412b6204e0" +dependencies = [ + "enumset_derive", + "serde", +] + +[[package]] +name = "enumset_derive" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4bd536557b58c682b217b8fb199afdff47cd3eff260623f19e77074eb073d63a" +dependencies = [ + "darling 0.21.3", + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "equivalent" version = "1.0.2" @@ -1984,7 +2667,7 @@ version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b95f7c0680e4142284cf8b22c14a476e87d61b004a3a0861872b32ef7ead40a2" dependencies = [ - "bit-set", + "bit-set 0.5.3", "regex", ] @@ -2087,6 +2770,27 @@ dependencies = [ "zlib-rs", ] +[[package]] +name = "float-ord" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ce81f49ae8a0482e4c55ea62ebbd7e5a686af544c00b9d090bba3ff9be97b3d" + +[[package]] +name = "float4" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a5404bf31d22893d61cf24d4dda149d8e6b2ff07601c3cb3be651031f61a4ed" + +[[package]] +name = "float8" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2d1f04709a8ac06e8e8042875a3c466cc4832d3c1a18dbcb9dba3c6e83046bc" +dependencies = [ + "half", +] + [[package]] name = "flume" version = "0.12.0" @@ -2134,6 +2838,16 @@ dependencies = [ "autocfg", ] +[[package]] +name = "fs2" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9564fc758e15025b46aa6643b1b77d047d1a56a1aea6e01002ac0c7026876213" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "fs_extra" version = "1.3.0" @@ -2348,6 +3062,12 @@ dependencies = [ "polyval", ] +[[package]] +name = "gimli" +version = "0.32.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e629b9b98ef3dd8afe6ca2bd0f89306cec16d43d907889945bc5d6687f2f13c7" + [[package]] name = "git-version" version = "0.3.9" @@ -2368,6 +3088,17 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "gl_generator" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a95dfc23a2b4a9a2f5ab41d194f8bfda3cabec42af4e39f08c339eb2a0c124d" +dependencies = [ + "khronos_api", + "log", + "xml-rs", +] + [[package]] name = "glob" version = "0.3.4" @@ -2399,6 +3130,27 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "glow" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29038e1c483364cc6bb3cf78feee1816002e127c331a1eec55a4d202b9e1adb5" +dependencies = [ + "js-sys", + "slotmap", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "glutin_wgl_sys" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c4ee00b289aba7a9e5306d57c2d05499b2e5dc427f84ac708bd2c090212cf3e" +dependencies = [ + "gl_generator", +] + [[package]] name = "goblin" version = "0.8.2" @@ -2410,6 +3162,40 @@ dependencies = [ "scroll", ] +[[package]] +name = "gpu-allocator" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51255ea7cfaadb6c5f1528d43e92a82acb2b96c43365989a28b2d44ee38f8795" +dependencies = [ + "ash", + "hashbrown 0.16.1", + "log", + "presser", + "thiserror 2.0.20", + "windows", +] + +[[package]] +name = "gpu-descriptor" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b89c83349105e3732062a895becfc71a8f921bb71ecbbdd8ff99263e3b53a0ca" +dependencies = [ + "bitflags 2.13.1", + "gpu-descriptor-types", + "hashbrown 0.15.5", +] + +[[package]] +name = "gpu-descriptor-types" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdf242682df893b86f33a73828fb09ca4b2d3bb6cc95249707fc684d27484b91" +dependencies = [ + "bitflags 2.13.1", +] + [[package]] name = "h2" version = "0.4.19" @@ -2435,8 +3221,11 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if 1.0.4", "crunchy", + "num-traits", + "serde", "zerocopy", ] @@ -2540,6 +3329,12 @@ dependencies = [ "arrayvec", ] +[[package]] +name = "hexf-parse" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfa686283ad6dd069f105e5ab091b04c62850d3e4cf5d67debad1933f55023df" + [[package]] name = "hf-xet" version = "1.6.0" @@ -2764,6 +3559,7 @@ dependencies = [ "tokio", "tokio-rustls", "tower-service", + "webpki-roots 1.0.9", ] [[package]] @@ -3028,6 +3824,12 @@ dependencies = [ "rustversion", ] +[[package]] +name = "inflections" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a257582fdcde896fd96463bf2d40eefea0580021c0712a0e2b028b60b47a837a" + [[package]] name = "inout" version = "0.1.4" @@ -3054,7 +3856,7 @@ version = "0.3.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2bf84e73fa6f27f299dec58e13223cf70db80da872eb921d4f6138342a0eabc8" dependencies = [ - "darling", + "darling 0.24.1", "indoc", "proc-macro2", "quote", @@ -3119,8 +3921,8 @@ dependencies = [ "portable-atomic", "portmapper", "rand 0.10.2", - "reqwest", - "rustc-hash", + "reqwest 0.13.4", + "rustc-hash 2.1.3", "rustls", "rustls-pki-types", "serde", @@ -3191,7 +3993,7 @@ dependencies = [ "itoa", "n0-error", "portable-atomic", - "reqwest", + "reqwest 0.13.4", "rustls", "rustls-platform-verifier", "ryu", @@ -3246,7 +4048,7 @@ dependencies = [ "rand 0.10.2", "rcgen", "reloadable-state", - "reqwest", + "reqwest 0.13.4", "rustls", "rustls-cert-file-reader", "rustls-cert-reloadable-resolver", @@ -3277,6 +4079,15 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + [[package]] name = "itertools" version = "0.14.0" @@ -3429,6 +4240,23 @@ dependencies = [ "log", ] +[[package]] +name = "khronos-egl" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6aae1df220ece3c0ada96b8153459b67eebe9ae9212258bb0134ae60416fdf76" +dependencies = [ + "libc", + "libloading 0.8.9", + "pkg-config", +] + +[[package]] +name = "khronos_api" +version = "3.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2db585e1d738fc771bf08a151420d3ed193d9d895a36df7f6f8a9456b911ddc" + [[package]] name = "konst" version = "0.4.3" @@ -3464,6 +4292,16 @@ version = "0.2.189" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if 1.0.4", + "windows-link", +] + [[package]] name = "libloading" version = "0.9.0" @@ -3474,6 +4312,27 @@ dependencies = [ "windows-link", ] +[[package]] +name = "liblzma" +version = "0.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fe0a34ca854fd4f20c07f696fc8675aec78f87d88d29f5e10257a7490a1b2e1" +dependencies = [ + "liblzma-sys", + "num_cpus", +] + +[[package]] +name = "liblzma-sys" +version = "0.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a0dad045e4b1b7b170be4b60b54b780cafb4490165461bac7d1cf7b703f61d5f" +dependencies = [ + "cc", + "libc", + "pkg-config", +] + [[package]] name = "libm" version = "0.2.16" @@ -3543,7 +4402,7 @@ checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" name = "llama-quant-ffi" version = "0.76.1" dependencies = [ - "libloading", + "libloading 0.9.0", ] [[package]] @@ -3640,6 +4499,12 @@ version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" +[[package]] +name = "md5" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ebb8d8732c6a6df3d8f032a82911cfc747e00efb95cc46e8d0acd5b5b88570c" + [[package]] name = "mdns-sd" version = "0.21.1" @@ -3703,7 +4568,7 @@ dependencies = [ "mesh-llm-plugin-manager", "mesh-llm-system", "mesh-llm-tui", - "reqwest", + "reqwest 0.13.4", "serde", "serde_json", "serial_test", @@ -3804,7 +4669,7 @@ dependencies = [ "model-ref", "nix 0.31.3", "rand 0.10.2", - "reqwest", + "reqwest 0.13.4", "rpassword", "serde", "serde_json", @@ -3925,7 +4790,7 @@ dependencies = [ "hyper", "pathdiff", "percent-encoding", - "reqwest", + "reqwest 0.13.4", "serde", "serde_json", "sha2 0.11.0", @@ -4002,7 +4867,7 @@ dependencies = [ "prost", "rand 0.10.2", "regex-lite", - "reqwest", + "reqwest 0.13.4", "rmcp", "rpassword", "rustls", @@ -4013,6 +4878,7 @@ dependencies = [ "serde_json", "serial_test", "sha2 0.11.0", + "skippy-cache", "skippy-coordinator", "skippy-ffi", "skippy-model", @@ -4144,7 +5010,7 @@ dependencies = [ "futures-util", "hex", "mesh-llm-skills", - "reqwest", + "reqwest 0.13.4", "serde", "serde_json", "sha2 0.11.0", @@ -4208,7 +5074,7 @@ dependencies = [ "mesh-llm-build-info", "mesh-llm-hardware-profile", "mesh-llm-native-runtime", - "reqwest", + "reqwest 0.13.4", "serde", "serde_json", "sha2 0.11.0", @@ -4228,7 +5094,7 @@ dependencies = [ "mesh-llm-console-server", "mesh-llm-embedded-runtime", "mesh-llm-runtime-install", - "reqwest", + "reqwest 0.13.4", "serde", "serde_json", ] @@ -4254,13 +5120,13 @@ dependencies = [ "dirs", "hex", "libc", - "libloading", + "libloading 0.9.0", "mesh-llm-build-info", "mesh-llm-gpu-bench", "mesh-llm-native-runtime", "mesh-llm-release-footer", "mesh-llm-runtime-install", - "reqwest", + "reqwest 0.13.4", "semver", "serde", "serde_json", @@ -4275,7 +5141,7 @@ dependencies = [ name = "mesh-llm-test-harness" version = "0.76.1" dependencies = [ - "reqwest", + "reqwest 0.13.4", "serde_json", "thiserror 2.0.20", ] @@ -4320,7 +5186,7 @@ version = "0.76.1" dependencies = [ "async-trait", "mesh-llm-guardrails", - "reqwest", + "reqwest 0.13.4", "serde", "serde_json", "tokio", @@ -4336,7 +5202,7 @@ name = "mesh-native-serving-plugin-host" version = "0.76.1" dependencies = [ "anyhow", - "libloading", + "libloading 0.9.0", "mesh-native-serving-plugin-api", "skippy-server", "skippy-tokenizer", @@ -4458,7 +5324,7 @@ dependencies = [ "mesh-llm-hf-hub", "model-hf", "model-ref", - "reqwest", + "reqwest 0.13.4", "serde", "serde_json", "sha2 0.11.0", @@ -4577,6 +5443,32 @@ dependencies = [ "n0-future", ] +[[package]] +name = "naga" +version = "29.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2bf919621e7975acb27d881bae2fb993e0d45c8e0446e85e6272971e00dc8df" +dependencies = [ + "arrayvec", + "bit-set 0.9.1", + "bitflags 2.13.1", + "cfg-if 1.0.4", + "cfg_aliases", + "codespan-reporting", + "half", + "hashbrown 0.16.1", + "hexf-parse", + "indexmap", + "libm", + "log", + "num-traits", + "once_cell", + "rustc-hash 1.1.0", + "spirv", + "thiserror 2.0.20", + "unicode-ident", +] + [[package]] name = "napi" version = "3.12.2" @@ -4590,7 +5482,7 @@ dependencies = [ "napi-build", "napi-sys", "nohash-hasher", - "rustc-hash", + "rustc-hash 2.1.3", "tokio", ] @@ -4633,15 +5525,39 @@ version = "3.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "85fbf1fa9f1babfe396d74bbbf52b3643770243e8f5b0b46715d4caf7f0dfc9a" dependencies = [ - "libloading", + "libloading 0.9.0", +] + +[[package]] +name = "nb" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "801d31da0513b6ec5214e9bf433a77966320625a37860f910be265be6e18d06f" +dependencies = [ + "nb 1.1.0", ] +[[package]] +name = "nb" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d5439c4ad607c3c23abf66de8c8bf57ba8adcd1f129e699851a6e43935d339d" + [[package]] name = "ndk-context" version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "27b02d87554356db9e9a873add8782d4ea6e3e58ea071a9adb9a2e8ddb884a8b" +[[package]] +name = "ndk-sys" +version = "0.6.0+11769913" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee6cda3051665f1fb8d9e08fc35c96d5a244fb1be711a03b71118828afc9a873" +dependencies = [ + "jni-sys 0.3.1", +] + [[package]] name = "negentropy" version = "0.5.1" @@ -4887,7 +5803,7 @@ dependencies = [ "noq-proto", "noq-udp", "pin-project-lite", - "rustc-hash", + "rustc-hash 2.1.3", "rustls", "socket2", "thiserror 2.0.20", @@ -4914,7 +5830,7 @@ dependencies = [ "rand 0.10.2", "rand_pcg", "ring", - "rustc-hash", + "rustc-hash 2.1.3", "rustls", "rustls-pki-types", "rustls-platform-verifier", @@ -5112,6 +6028,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" dependencies = [ "autocfg", + "libm", +] + +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", ] [[package]] @@ -5246,6 +6173,31 @@ dependencies = [ "objc2-core-foundation", ] +[[package]] +name = "objc2-metal" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a0125f776a10d00af4152d74616409f0d4a2053a6f57fa5b7d6aa2854ac04794" +dependencies = [ + "bitflags 2.13.1", + "block2", + "objc2", + "objc2-foundation", +] + +[[package]] +name = "objc2-quartz-core" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96c1358452b371bf9f104e21ec536d37a650eb10f7ee379fff67d2e08d537f1f" +dependencies = [ + "bitflags 2.13.1", + "objc2", + "objc2-core-foundation", + "objc2-foundation", + "objc2-metal", +] + [[package]] name = "objc2-security" version = "0.3.2" @@ -5281,6 +6233,15 @@ dependencies = [ "objc2-security", ] +[[package]] +name = "object" +version = "0.37.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff76201f031d8863c38aa7f905eca4f53abbfa15f609db4277d44cd8938f33fe" +dependencies = [ + "memchr", +] + [[package]] name = "oid-registry" version = "0.8.1" @@ -5312,6 +6273,12 @@ version = "0.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "269bca4c2591a28585d6bf10d9ed0332b7d76900a1b02bec41bdc3a2cdcda107" +[[package]] +name = "oneshot" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfe21416a02c693fb9f980befcb230ecc70b0b3d1cc4abf88b9675c4c1457f0c" + [[package]] name = "opaque-debug" version = "0.3.1" @@ -5373,7 +6340,7 @@ dependencies = [ "bytes", "http", "opentelemetry", - "reqwest", + "reqwest 0.13.4", ] [[package]] @@ -5388,7 +6355,7 @@ dependencies = [ "opentelemetry-proto", "opentelemetry_sdk", "prost", - "reqwest", + "reqwest 0.13.4", "thiserror 2.0.20", ] @@ -5645,6 +6612,7 @@ dependencies = [ "fixedbitset 0.5.7", "hashbrown 0.15.5", "indexmap", + "serde", ] [[package]] @@ -5861,6 +6829,15 @@ dependencies = [ "serde", ] +[[package]] +name = "portable-atomic-util" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10ab3eb7f3becc3a1cbc4f2c6f20267996cfc1a6467a873763411b136a122715" +dependencies = [ + "portable-atomic", +] + [[package]] name = "portmapper" version = "0.19.3" @@ -5949,6 +6926,12 @@ dependencies = [ "num-traits", ] +[[package]] +name = "presser" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8cf8e6a8aa66ce33f63993ffc4ea4271eb5b0530a9002db8455ea6050c77bfa" + [[package]] name = "prettyplease" version = "0.2.37" @@ -6001,6 +6984,12 @@ dependencies = [ "windows", ] +[[package]] +name = "profiling" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d595e54a326bc53c1c197b32d295e14b169e3cfeaa8dc82b529f947fba6bcf5" + [[package]] name = "proptest" version = "1.11.0" @@ -6033,7 +7022,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "03da047801ff44bb6a4d407d4860c05fd70bb81714e6b2f3812603d5b145b042" dependencies = [ "heck", - "itertools", + "itertools 0.14.0", "log", "multimap", "petgraph", @@ -6052,7 +7041,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" dependencies = [ "anyhow", - "itertools", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.119", @@ -6163,7 +7152,7 @@ dependencies = [ "pin-project-lite", "quinn-proto", "quinn-udp", - "rustc-hash", + "rustc-hash 2.1.3", "rustls", "socket2", "thiserror 2.0.20", @@ -6185,7 +7174,7 @@ dependencies = [ "rand 0.10.2", "rand_pcg", "ring", - "rustc-hash", + "rustc-hash 2.1.3", "rustls", "rustls-pki-types", "slab", @@ -6334,6 +7323,12 @@ dependencies = [ "rand_core 0.9.5", ] +[[package]] +name = "range-alloc" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca45419789ae5a7899559e9512e58ca889e41f04f1f2445e9f4b290ceccd1d08" + [[package]] name = "ratatui" version = "0.30.2" @@ -6360,7 +7355,7 @@ dependencies = [ "compact_str", "critical-section", "hashbrown 0.17.1", - "itertools", + "itertools 0.14.0", "kasuari", "lru", "palette", @@ -6425,7 +7420,7 @@ dependencies = [ "hashbrown 0.17.1", "indoc", "instability", - "itertools", + "itertools 0.14.0", "line-clipping", "ratatui-core", "serde", @@ -6435,6 +7430,24 @@ dependencies = [ "unicode-width", ] +[[package]] +name = "raw-window-handle" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20675572f6f24e9e76ef639bc5552774ed45f1c30e2951e1e99c59888861c539" + +[[package]] +name = "raw-window-metal" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40d213455a5f1dc59214213c7330e074ddf8114c9a42411eb890c767357ce135" +dependencies = [ + "objc2", + "objc2-core-foundation", + "objc2-foundation", + "objc2-quartz-core", +] + [[package]] name = "rcgen" version = "0.14.10" @@ -6550,6 +7563,53 @@ dependencies = [ "tokio", ] +[[package]] +name = "renderdoc-sys" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19b30a45b0cd0bcca8037f3d0dc3421eaf95327a17cad11964fb8179b4fc4832" + +[[package]] +name = "reqwest" +version = "0.12.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures-channel", + "futures-core", + "futures-util", + "h2", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls", + "rustls-pki-types", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tokio-rustls", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", + "webpki-roots 1.0.9", +] + [[package]] name = "reqwest" version = "0.13.4" @@ -6605,7 +7665,7 @@ dependencies = [ "anyhow", "async-trait", "http", - "reqwest", + "reqwest 0.13.4", "thiserror 2.0.20", "tower-service", ] @@ -6649,7 +7709,7 @@ dependencies = [ "pin-project-lite", "process-wrap", "rand 0.10.2", - "reqwest", + "reqwest 0.13.4", "rmcp-macros", "schemars", "serde", @@ -6670,7 +7730,7 @@ version = "3.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cdf1c49bd4d52014b94db0877410db273c2008f01628b0252a2e9460ad9b7fda" dependencies = [ - "darling", + "darling 0.24.1", "proc-macro2", "quote", "serde_json", @@ -6711,6 +7771,18 @@ dependencies = [ "smallvec", ] +[[package]] +name = "rustc-demangle" +version = "0.1.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b74b56ffa8bb2830709a538c2cbcae9aa062db0d2a42563bfb09bdaae44020eb" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + [[package]] name = "rustc-hash" version = "2.1.3" @@ -6929,6 +8001,15 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "sanitize-filename" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc984f4f9ceb736a7bb755c3e3bd17dc56370af2600c9780dcc48c66453da34d" +dependencies = [ + "regex", +] + [[package]] name = "schannel" version = "0.1.29" @@ -7305,6 +8386,12 @@ dependencies = [ "os_str_bytes", ] +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + [[package]] name = "shlex" version = "2.0.1" @@ -7400,9 +8487,10 @@ dependencies = [ "model-artifact", "model-hf", "model-ref", - "reqwest", + "reqwest 0.13.4", "serde", "serde_json", + "skippy-cache", "skippy-protocol", "skippy-runtime", "skippy-topology", @@ -7414,7 +8502,13 @@ version = "0.76.1" dependencies = [ "anyhow", "blake3", + "cubecl", + "fs2", + "libc", + "serde", + "serde_json", "skippy-protocol", + "windows-sys 0.61.2", ] [[package]] @@ -7429,15 +8523,17 @@ name = "skippy-correctness" version = "0.76.1" dependencies = [ "anyhow", + "blake3", "clap", "hex", "model-artifact", "model-hf", "model-ref", - "reqwest", + "reqwest 0.13.4", "serde", "serde_json", "sha2 0.11.0", + "skippy-cache", "skippy-protocol", "skippy-runtime", ] @@ -7446,7 +8542,7 @@ dependencies = [ name = "skippy-ffi" version = "0.76.1" dependencies = [ - "libloading", + "libloading 0.9.0", ] [[package]] @@ -7555,6 +8651,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.11.0", + "skippy-cache", "skippy-ffi", "skippy-model", "tempfile", @@ -7632,11 +8729,33 @@ version = "0.4.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" +[[package]] +name = "slotmap" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bdd58c3c93c3d278ca835519292445cb4b0d4dc59ccfdf7ceadaab3f8aeb4038" +dependencies = [ + "version_check", +] + +[[package]] +name = "slug" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "882a80f72ee45de3cc9a5afeb2da0331d58df69e4e7d8eeb5d3c7784ae67e724" +dependencies = [ + "deunicode", + "wasm-bindgen", +] + [[package]] name = "smallvec" version = "1.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b9be42f50aa861c555654aa3a37f52f4b1074bacf4e48fe0ef7fa584e80f1f0f" +dependencies = [ + "serde", +] [[package]] name = "smawk" @@ -7696,6 +8815,19 @@ name = "spin" version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "023a211cb3138dbc438680b32560ad89f699977624c9f8dbb95a47d5b4c07dd3" +dependencies = [ + "lock_api", + "portable-atomic", +] + +[[package]] +name = "spirv" +version = "0.4.0+sdk-1.4.341.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9571ea910ebd84c86af4b3ed27f9dbdc6ad06f17c5f96146b2b671e2976744f" +dependencies = [ + "bitflags 2.13.1", +] [[package]] name = "spki" @@ -7720,6 +8852,12 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "stable-vec" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "047720fbc1c4cc7daa75c925adba9ce3fa8e1f7b039337c281b142d4ecc6d86b" + [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -7908,6 +9046,15 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "termcolor" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06794f8f6c5c898b3275aebefa6b8a1cb24cd2c6c79397ab15774837a0bc5755" +dependencies = [ + "winapi-util", +] + [[package]] name = "termina" version = "0.3.3" @@ -8187,7 +9334,7 @@ dependencies = [ "pem 3.0.6", "proc-macro2", "rcgen", - "reqwest", + "reqwest 0.13.4", "ring", "rustls", "serde", @@ -8468,6 +9615,85 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" +[[package]] +name = "tracel-llvm" +version = "20.1.4-7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "982535db9eb1a30ac0f2c50239a0eec3e5cf50993a88e92b04747bd2f4d365b2" +dependencies = [ + "tracel-mlir-rs", + "tracel-mlir-sys", +] + +[[package]] +name = "tracel-llvm-bundler" +version = "20.1.4-7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c75b8e477cb8d49d907afab029ca74d48459f5b88c27bdb4c6cd6acb5e61977" +dependencies = [ + "anyhow", + "bytes", + "constcat", + "dirs", + "liblzma", + "regex", + "reqwest 0.12.28", + "serde", + "serde_json", + "sha2 0.10.9", + "tar", + "walkdir", +] + +[[package]] +name = "tracel-mlir-rs" +version = "20.1.4-7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77a478a35efd68d0ba73f747adfb7923b121c64e7f5be9cd8364ca1dcb772d5c" +dependencies = [ + "tracel-mlir-rs-macros", + "tracel-mlir-sys", +] + +[[package]] +name = "tracel-mlir-rs-macros" +version = "20.1.4-7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a94f36868c3b10b1825945223d99d106c73f4d249f063caa4651deeb9379344" +dependencies = [ + "comrak", + "convert_case 0.8.0", + "proc-macro2", + "quote", + "regex", + "syn 2.0.119", + "tracel-llvm-bundler", + "tracel-tblgen-rs", + "unindent", +] + +[[package]] +name = "tracel-mlir-sys" +version = "20.1.4-7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02f26d31af0c225a6d2e3d65d012fd6de848c9fc776897b152ee83b7d1bd15c4" +dependencies = [ + "tracel-llvm-bundler", +] + +[[package]] +name = "tracel-tblgen-rs" +version = "20.1.4-7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d2581070380418ccc33b500f3739e4d4869421fdb477fcea51ff97c6253a52" +dependencies = [ + "bindgen", + "cc", + "paste", + "thiserror 2.0.20", + "tracel-llvm-bundler", +] + [[package]] name = "tracing" version = "0.1.44" @@ -8587,6 +9813,30 @@ version = "2.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5283634e518fe9e82c7b20520bb4bc209009fd16c82077c802f8111ecbb0117a" +[[package]] +name = "tynm" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a21cdb0fc8f85c98b1ec812bc4cd69faf6c0fa2fc17d44ea3c2cdd38dc08e999" +dependencies = [ + "nom 8.0.0", +] + +[[package]] +name = "type-map" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb30dbbd9036155e74adad6812e9898d03ec374946234fbcebd5dfc7b9187b90" +dependencies = [ + "rustc-hash 2.1.3", +] + +[[package]] +name = "typed-arena" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6af6ae20167a9ece4bcb41af5b80f8a1f1df981f6391189ce00fd257af04126a" + [[package]] name = "typed-path" version = "0.12.3" @@ -8661,7 +9911,7 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "16b380a1238663e5f8a691f9039c73e1cdae598a30e9855f541d29b08b53e9a5" dependencies = [ - "itertools", + "itertools 0.14.0", "unicode-segmentation", "unicode-width", ] @@ -8678,6 +9928,12 @@ version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + [[package]] name = "uniffi" version = "0.32.0" @@ -8809,6 +10065,12 @@ dependencies = [ "weedle2", ] +[[package]] +name = "unindent" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" + [[package]] name = "universal-hash" version = "0.5.1" @@ -8903,6 +10165,17 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" +[[package]] +name = "variadics_please" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41b6d82be61465f97d42bd1d15bf20f3b0a3a0905018f38f9d6f6962055b0b5c" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "vcpkg" version = "0.2.15" @@ -8915,6 +10188,12 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "void" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a02e4885ed3bc0f2de90ea6dd45ebcbb66dacffe03547fadbb0eeae2770887d" + [[package]] name = "vtparse" version = "0.6.2" @@ -9044,6 +10323,18 @@ dependencies = [ "web-sys", ] +[[package]] +name = "wayland-sys" +version = "0.31.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8eab23fefc9e41f8e841df4a9c707e8a8c4ed26e944ef69297184de2785e3be" +dependencies = [ + "dlib", + "log", + "once_cell", + "pkg-config", +] + [[package]] name = "web-sys" version = "0.3.104" @@ -9178,6 +10469,174 @@ dependencies = [ "wezterm-dynamic", ] +[[package]] +name = "wgpu" +version = "29.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76e8840e1ba2881d4cbb18d2147627a56af426ff064c0401eb0c8410c6325d07" +dependencies = [ + "arrayvec", + "bitflags 2.13.1", + "bytemuck", + "cfg-if 1.0.4", + "cfg_aliases", + "document-features", + "hashbrown 0.16.1", + "js-sys", + "log", + "naga", + "parking_lot", + "portable-atomic", + "profiling", + "raw-window-handle", + "smallvec", + "static_assertions", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", + "wgpu-core", + "wgpu-hal", + "wgpu-types", +] + +[[package]] +name = "wgpu-core" +version = "29.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f519832254e56965a9940c4af57dcb75f702b6f6fa4a0b172f685395843a4d7" +dependencies = [ + "arrayvec", + "bit-set 0.9.1", + "bit-vec 0.9.1", + "bitflags 2.13.1", + "bytemuck", + "cfg_aliases", + "document-features", + "hashbrown 0.16.1", + "indexmap", + "log", + "naga", + "once_cell", + "parking_lot", + "portable-atomic", + "profiling", + "raw-window-handle", + "rustc-hash 1.1.0", + "smallvec", + "thiserror 2.0.20", + "wgpu-core-deps-apple", + "wgpu-core-deps-emscripten", + "wgpu-core-deps-windows-linux-android", + "wgpu-hal", + "wgpu-naga-bridge", + "wgpu-types", +] + +[[package]] +name = "wgpu-core-deps-apple" +version = "29.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f5e39e26c4c0e07589e67d18546cf79ff45383659fc72fca4dd293358a0347f3" +dependencies = [ + "wgpu-hal", +] + +[[package]] +name = "wgpu-core-deps-emscripten" +version = "29.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01e09be551dc939498bdd5f6b2c66e55ab275dad25825267a08605a80fc9f0af" +dependencies = [ + "wgpu-hal", +] + +[[package]] +name = "wgpu-core-deps-windows-linux-android" +version = "29.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e592c1bbef6ad047647ae6e666ebd8cee7a32bb4544d9700ec96cbf73230257" +dependencies = [ + "wgpu-hal", +] + +[[package]] +name = "wgpu-hal" +version = "29.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97ace1c17727311c22a46e4e3faf56ea6de81af99dcc839bdfb54857b94d448d" +dependencies = [ + "android_system_properties", + "arrayvec", + "ash", + "bit-set 0.9.1", + "bitflags 2.13.1", + "block2", + "bytemuck", + "cfg-if 1.0.4", + "cfg_aliases", + "glow", + "glutin_wgl_sys", + "gpu-allocator", + "gpu-descriptor", + "hashbrown 0.16.1", + "js-sys", + "khronos-egl", + "libc", + "libloading 0.8.9", + "log", + "naga", + "ndk-sys", + "objc2", + "objc2-core-foundation", + "objc2-foundation", + "objc2-metal", + "objc2-quartz-core", + "once_cell", + "ordered-float", + "parking_lot", + "portable-atomic", + "portable-atomic-util", + "profiling", + "range-alloc", + "raw-window-handle", + "raw-window-metal", + "renderdoc-sys", + "smallvec", + "thiserror 2.0.20", + "wasm-bindgen", + "wayland-sys", + "web-sys", + "wgpu-naga-bridge", + "wgpu-types", + "windows", + "windows-core", + "windows-result", +] + +[[package]] +name = "wgpu-naga-bridge" +version = "29.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95226013f547544b223281cd16a4fb549aa9dcb562adbda0faae4c73ffbbc161" +dependencies = [ + "naga", + "wgpu-types", +] + +[[package]] +name = "wgpu-types" +version = "29.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "84bf84cd9ca8ca45e2b223a3868f1adf9bfc0c66aeac212e76ee7e40fdadf8f5" +dependencies = [ + "bitflags 2.13.1", + "bytemuck", + "js-sys", + "log", + "raw-window-handle", + "web-sys", +] + [[package]] name = "whoami" version = "2.1.3" @@ -9707,7 +11166,7 @@ dependencies = [ "more-asserts", "rand 0.10.2", "redb", - "reqwest", + "reqwest 0.13.4", "reqwest-middleware", "serde", "serde_json", @@ -9742,7 +11201,7 @@ dependencies = [ "futures-util", "getrandom 0.4.3", "heapify", - "itertools", + "itertools 0.14.0", "lz4_flex", "more-asserts", "rand 0.10.2", @@ -9771,7 +11230,7 @@ dependencies = [ "chrono", "gearhash", "http", - "itertools", + "itertools 0.14.0", "more-asserts", "rand 0.10.2", "serde", @@ -9811,10 +11270,10 @@ dependencies = [ "konst", "libc", "more-asserts", - "oneshot", + "oneshot 0.1.13", "pin-project", "rand 0.10.2", - "reqwest", + "reqwest 0.13.4", "serde", "serde_json", "shellexpand", @@ -9859,6 +11318,12 @@ dependencies = [ "sha2 0.11.0", ] +[[package]] +name = "xxhash-rust" +version = "0.8.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aee1b19627c7c60102ab80d3a9cbe18de90bfe03bfa6c3715447681f0e8c8af6" + [[package]] name = "yasna" version = "0.6.0" diff --git a/ci/ci.md b/ci/ci.md index 754bcd7a7e..191fb3f4fd 100644 --- a/ci/ci.md +++ b/ci/ci.md @@ -387,8 +387,7 @@ runtime producers are not duplicated. that join only their matching immutable host and runtime artifacts. - `ci-platform-checks-slice.yml` — macOS portable/unit, Windows portable, and focused Windows log-store privacy ACL checks. -- `ci-linux-product-smoke-slice.yml` and - `ci-macos-product-smoke-slice.yml` — platform-local callers of the typed +- `ci-{linux,macos,windows}-product-smoke-slice.yml` — platform-local callers of the typed product-integration suite and the model-download consumer. The suite stages its registry-derived pair exactly once: dense SmolLM2-135M Q8 and recurrent IBM Granite 4.0 H 350M Q4. Its ordered phases cover dense standalone, @@ -401,7 +400,13 @@ runtime producers are not duplicated. topology/run/model/package/manifest identity, the same exact two-stage contiguous cut on distinct nodes and bind addresses, two matching `ready` statuses, and the same sole served model. It atomically records - `split-evidence.json`. Readiness uses a capped five-minute wall-clock deadline + `split-evidence.json`. A separately reconciled `durable-l3` phase enables a + fixed, CLI-sourced disk tier for both model families, preserves each node's + root and identity across a full process restart, requires a post-restart L3 + fill with cached tokens and exact output, exercises `kv-cache status` and + `clear`, and records digest-bound evidence. Windows CPU runs this phase alone + on a real Windows product executor; Linux CPU and macOS Metal include it in + the complete suite. Readiness uses a capped five-minute wall-clock deadline and parallel endpoint captures bounded to two seconds by default; timeout or process-exit diagnostics retain the final snapshots, failed reconciliation, and both server log tails. The status projection never persists invite @@ -412,10 +417,9 @@ runtime producers are not duplicated. evidence. All JSON snapshots and reconciled evidence upload with the phase logs on success or failure. The existing Qwen3.5 recurrent job remains required until Granite passes that live contract. The typed runner supports - CPU, CUDA, Metal, Vulkan, and ROCm, - but only CPU is selected during the first qualification stage; the existing - CUDA inference and Metal model-load signals remain required until their typed - product rows pass live qualification in that order. CUDA and Metal request + CPU, CUDA, Metal, Vulkan, ROCm, and Windows CPU. CPU, Metal, and Windows CPU + are selected; the existing CUDA inference signal remains required until its + typed product row passes live qualification. CUDA and Metal request their explicit accelerator device and reject unsupported typed selections. CUDA inference uses the approved `gpu-nvidia` ephemeral self-hosted scale set, including for diff --git a/crates/mesh-client/src/client/control_plane.rs b/crates/mesh-client/src/client/control_plane.rs index 4f9d2b3868..d5d5c9ae51 100644 --- a/crates/mesh-client/src/client/control_plane.rs +++ b/crates/mesh-client/src/client/control_plane.rs @@ -6,6 +6,7 @@ use crate::proto::node::{ OwnerControlDrainModelResponse, OwnerControlEnsureModelRequest, OwnerControlEnsureModelResponse, OwnerControlEnvelope, OwnerControlError, OwnerControlErrorCode, OwnerControlGetConfigRequest, OwnerControlHandshake, + OwnerControlKvCacheOperation, OwnerControlKvCacheRequest, OwnerControlKvCacheResponse, OwnerControlLoadModelRequest, OwnerControlLoadModelResponse, OwnerControlModelRef, OwnerControlRefreshInventory, OwnerControlRefreshInventoryRequest, OwnerControlRequest, OwnerControlResponse, OwnerControlUnloadModelRequest, OwnerControlUnloadModelResponse, @@ -34,6 +35,9 @@ const OWNER_CONTROL_REQUEST_WRITE_TIMEOUT_SECS: u64 = 2; const OWNER_CONTROL_SERVER_UNARY_DEADLINE_SECS_FOR_CLIENT_MARGIN: u64 = 5; const OWNER_CONTROL_UNARY_RESPONSE_TIMEOUT_SECS: u64 = OWNER_CONTROL_SERVER_UNARY_DEADLINE_SECS_FOR_CLIENT_MARGIN + 5; +const OWNER_CONTROL_SERVER_KV_CACHE_DEADLINE_SECS_FOR_CLIENT_MARGIN: u64 = 30; +const OWNER_CONTROL_KV_CACHE_RESPONSE_TIMEOUT_SECS: u64 = + OWNER_CONTROL_SERVER_KV_CACHE_DEADLINE_SECS_FOR_CLIENT_MARGIN + 5; const OWNER_CONTROL_SERVER_SCAN_DEADLINE_SECS_FOR_CLIENT_MARGIN: u64 = 30; const OWNER_CONTROL_INVENTORY_RESPONSE_TIMEOUT_SECS: u64 = OWNER_CONTROL_SERVER_SCAN_DEADLINE_SECS_FOR_CLIENT_MARGIN + 5; @@ -316,6 +320,27 @@ fn map_legacy_lifecycle_unsupported( } } +fn map_kv_cache_unsupported(error: ControlPlaneClientError) -> ControlPlaneClientError { + const LEGACY_UNKNOWN_COMMAND_MESSAGE: &str = + "owner control request requires exactly one command variant"; + match error { + // Only a server that does not know the command at all. A supported + // server answers a malformed request with BadRequest, and reporting + // that as "unsupported" would be a false capability result. + ControlPlaneClientError::Remote(mut remote) + if remote.code == OwnerControlErrorCode::UnknownCommand + || (remote.code == OwnerControlErrorCode::BadRequest + && remote.message == LEGACY_UNKNOWN_COMMAND_MESSAGE) => + { + remote.code = OwnerControlErrorCode::ControlUnsupported; + remote.message = + "remote owner-control endpoint does not support kv-cache operations".to_string(); + ControlPlaneClientError::Remote(remote) + } + other => other, + } +} + enum LifecycleCommand { Load(OwnerControlLoadModelRequest), Unload(OwnerControlUnloadModelRequest), @@ -442,6 +467,7 @@ impl OwnerControlClient { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }, ) .await?; @@ -478,6 +504,7 @@ impl OwnerControlClient { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }, ) .await?; @@ -513,6 +540,7 @@ impl OwnerControlClient { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }, ) .await?; @@ -532,6 +560,42 @@ impl OwnerControlClient { }) } + pub async fn kv_cache( + &self, + operation: OwnerControlKvCacheOperation, + target_bytes: Option, + model_identity: Option, + ) -> Result { + let response = self + .send_unary_request( + std::time::Duration::from_secs(OWNER_CONTROL_KV_CACHE_RESPONSE_TIMEOUT_SECS), + |request_id, requester_node_id, target_node_id| OwnerControlRequest { + request_id, + kv_cache: Some(OwnerControlKvCacheRequest { + requester_node_id, + target_node_id, + operation: operation as i32, + target_bytes, + model_identity, + }), + ..Default::default() + }, + ) + .await + .map_err(map_kv_cache_unsupported)?; + let response = response.kv_cache.ok_or_else(|| { + ControlPlaneClientError::Protocol( + "owner-control kv_cache response missing payload".to_string(), + ) + })?; + serde_json::from_slice::(&response.status_json).map_err(|error| { + ControlPlaneClientError::Protocol(format!( + "owner-control kv_cache response has invalid status JSON: {error}" + )) + })?; + Ok(response) + } + pub async fn load_model( &self, model_ref: String, @@ -705,6 +769,7 @@ impl OwnerControlClient { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), response: None, error: None, diff --git a/crates/mesh-client/tests/control_plane_client.rs b/crates/mesh-client/tests/control_plane_client.rs index 269c81efbf..952efca743 100644 --- a/crates/mesh-client/tests/control_plane_client.rs +++ b/crates/mesh-client/tests/control_plane_client.rs @@ -3,7 +3,8 @@ use iroh::{Endpoint, EndpointAddr, SecretKey}; use mesh_client::proto::node::{ CompactModelMetadata, ConfigApplyMode, NodeConfigSnapshot, NodeGpuConfig, NodeModelEntry, OwnerControlEnvelope, OwnerControlErrorCode, OwnerControlInventoryEntry, - OwnerControlRefreshInventory, OwnerControlRefreshInventoryDisposition, + OwnerControlKvCacheOperation, OwnerControlRefreshInventory, + OwnerControlRefreshInventoryDisposition, }; use mesh_client::protocol::{ ALPN_CONTROL_V1, MAX_CONTROL_FRAME_BYTES, NODE_PROTOCOL_GENERATION, @@ -170,6 +171,7 @@ async fn spawn_success_server_with_refresh_detail( unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), error: None, }, @@ -205,6 +207,7 @@ async fn spawn_success_server_with_refresh_detail( unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), error: None, }, @@ -247,6 +250,39 @@ async fn spawn_success_server_with_refresh_detail( unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, + }), + error: None, + }, + ) + .await; + let _ = send.finish(); + return; + } + if let Some(kv_cache) = request.kv_cache { + assert_eq!( + OwnerControlKvCacheOperation::try_from(kv_cache.operation), + Ok(OwnerControlKvCacheOperation::Prune) + ); + assert_eq!(kv_cache.target_bytes, Some(1024)); + assert_eq!(kv_cache.model_identity.as_deref(), Some("blake3:model")); + write_control_envelope( + &mut send, + OwnerControlEnvelope { + r#gen: NODE_PROTOCOL_GENERATION, + handshake: None, + request: None, + response: Some(mesh_client::proto::node::OwnerControlResponse { + request_id: request.request_id, + kv_cache: Some( + mesh_client::proto::node::OwnerControlKvCacheResponse { + status_json: + br#"{"version":1,"effective":{"state":"active"}}"# + .to_vec(), + freed_bytes: Some(512), + }, + ), + ..Default::default() }), error: None, }, @@ -291,6 +327,7 @@ async fn spawn_success_server_with_refresh_detail( unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), error: None, }, @@ -555,6 +592,19 @@ async fn control_plane_client_apply_config_get_watch_refresh_and_close() { "hf://mesh/refresh-model-GGUF:Q4_K_M" ); + let cache = control + .kv_cache( + OwnerControlKvCacheOperation::Prune, + Some(1024), + Some("blake3:model".to_string()), + ) + .await + .expect("kv_cache should succeed"); + assert_eq!(cache.freed_bytes, Some(512)); + let status: serde_json::Value = + serde_json::from_slice(&cache.status_json).expect("valid status JSON"); + assert_eq!(status["version"], 1); + let mut watch = control .watch_config(true) .await diff --git a/crates/mesh-client/tests/protocol_wire.rs b/crates/mesh-client/tests/protocol_wire.rs index 9e7db5b4e7..cd67498fc6 100644 --- a/crates/mesh-client/tests/protocol_wire.rs +++ b/crates/mesh-client/tests/protocol_wire.rs @@ -560,6 +560,7 @@ fn owner_control_envelope_roundtrip() { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), response: None, error: None, @@ -584,6 +585,7 @@ fn owner_control_unknown_command_rejects_with_structured_error() { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), response: None, error: None, diff --git a/crates/mesh-llm-cli/src/lib.rs b/crates/mesh-llm-cli/src/lib.rs index a06acd7a25..3ea9289b0f 100644 --- a/crates/mesh-llm-cli/src/lib.rs +++ b/crates/mesh-llm-cli/src/lib.rs @@ -18,7 +18,7 @@ pub use inventory::{ pub use parser::{ AuthCommand, BinaryFlavor, Cli, Command, ConfigCommand, DiscoveryScope, DoctorCommand, - GpuCommand, MeshDiscoveryMode, MeshGuardrailCliMode, NormalizedRuntimeArgs, PluginCommand, - RuntimeSurface, SkillAgentArg, SkillCommand, TrustCommand, TrustPolicy, + GpuCommand, KvCacheCommand, MeshDiscoveryMode, MeshGuardrailCliMode, NormalizedRuntimeArgs, + PluginCommand, RuntimeSurface, SkillAgentArg, SkillCommand, TrustCommand, TrustPolicy, legacy_runtime_surface_warning, normalize_runtime_surface_args, validate_discovery_mode_args, }; diff --git a/crates/mesh-llm-cli/src/parser.rs b/crates/mesh-llm-cli/src/parser.rs index 4fe7b4f196..14e3e61a88 100644 --- a/crates/mesh-llm-cli/src/parser.rs +++ b/crates/mesh-llm-cli/src/parser.rs @@ -6,8 +6,8 @@ mod validation; pub use commands::{ AuthCommand, BinaryFlavor, Cli, Command, ConfigCommand, DiscoveryScope, DoctorCommand, - GpuCommand, MeshDiscoveryMode, MeshGuardrailCliMode, PluginCommand, SkillAgentArg, - SkillCommand, TrustCommand, TrustPolicy, + GpuCommand, KvCacheCommand, MeshDiscoveryMode, MeshGuardrailCliMode, PluginCommand, + SkillAgentArg, SkillCommand, TrustCommand, TrustPolicy, }; pub use logging_help::logging_help; pub use normalization::{ @@ -100,7 +100,7 @@ pub fn assert_mesh_requirements_docs_examples_parse() { #[cfg(test)] mod tests { - use super::{Cli, Command}; + use super::{Cli, Command, KvCacheCommand}; use crate::models::ModelsCommand; use clap::Parser; @@ -142,6 +142,74 @@ mod tests { ); } + #[test] + fn disk_cache_runtime_flags_parse_without_rewriting_config() { + let normalized = super::normalize_runtime_surface_args( + [ + "mesh-llm", + "serve", + "--kv-cache-disk", + "32GiB", + "--kv-cache-disk-dir", + "/fast/mesh-kv", + "--kv-cache-min-free", + "16GiB", + ] + .into_iter() + .map(std::ffi::OsString::from) + .collect::>(), + ); + let cli = Cli::try_parse_from(normalized.normalized) + .expect("disk cache runtime flags should parse"); + + assert_eq!(cli.kv_cache_disk.as_deref(), Some("32GiB")); + assert_eq!( + cli.kv_cache_disk_dir, + Some(std::path::PathBuf::from("/fast/mesh-kv")) + ); + assert_eq!(cli.kv_cache_min_free.as_deref(), Some("16GiB")); + } + + #[test] + fn kv_cache_lifecycle_commands_parse_exact_model_identity_and_confirmation() { + let cli = Cli::try_parse_from([ + "mesh-llm", + "kv-cache", + "prune", + "--target", + "16GiB", + "--model-identity", + "blake3:model", + "--yes", + "--endpoint", + "node-a", + "--endpoint", + "node-b", + "--json", + ]) + .expect("kv-cache prune should parse"); + match cli.command { + Some(Command::KvCache { + command: + KvCacheCommand::Prune { + target, + model_identity, + yes, + endpoints, + json, + .. + }, + }) => { + assert_eq!(target.as_deref(), Some("16GiB")); + assert_eq!(model_identity.as_deref(), Some("blake3:model")); + assert!(yes); + assert_eq!(endpoints, vec!["node-a", "node-b"]); + assert!(json); + } + other => panic!("unexpected command: {other:?}"), + } + } + #[test] fn models_package_parses_experimental_publication() { let cli = Cli::parse_from([ diff --git a/crates/mesh-llm-cli/src/parser/commands.rs b/crates/mesh-llm-cli/src/parser/commands.rs index 47ee3ba06a..8a68ff6303 100644 --- a/crates/mesh-llm-cli/src/parser/commands.rs +++ b/crates/mesh-llm-cli/src/parser/commands.rs @@ -750,6 +750,18 @@ pub struct Cli { #[arg(long)] pub config: Option, + /// Node-local disk prompt cache: off, auto, or an explicit IEC size such as 32GiB. + #[arg(long, value_name = "off|auto|SIZE")] + pub kv_cache_disk: Option, + + /// Absolute node-local disk prompt-cache directory. + #[arg(long, value_name = "ABSOLUTE_PATH")] + pub kv_cache_disk_dir: Option, + + /// Minimum free storage to preserve, with an IEC suffix such as 16GiB. + #[arg(long, value_name = "SIZE")] + pub kv_cache_min_free: Option, + /// Path to the owner keystore used to attest this node. #[arg(long)] pub owner_key: Option, @@ -833,6 +845,12 @@ pub enum Command { #[command(subcommand)] command: ConfigCommand, }, + /// Inspect and manage the node-local durable prompt cache. + #[command(name = "kv-cache")] + KvCache { + #[command(subcommand)] + command: KvCacheCommand, + }, /// Diagnose local mesh, runtime, and split-readiness problems. Doctor { /// Print machine-readable JSON for the default doctor report. @@ -1093,6 +1111,53 @@ pub enum Command { ExternalPlugin(Vec), } +#[derive(Subcommand, Debug)] +pub enum KvCacheCommand { + /// Show the configured and effective cache state. + Status { + /// Authenticated owner-control endpoint; repeat for multiple owned nodes. + #[arg(long = "endpoint")] + endpoints: Vec, + #[arg(long, default_value = "3131")] + port: u16, + #[arg(long)] + json: bool, + }, + /// Evict least-recently-used inactive entries. + Prune { + /// Optional target size with an IEC suffix (for example 16GiB). + #[arg(long)] + target: Option, + /// Exact numerical model identity; display names are not accepted. + #[arg(long)] + model_identity: Option, + #[arg(long)] + yes: bool, + /// Authenticated owner-control endpoint; repeat for multiple owned nodes. + #[arg(long = "endpoint")] + endpoints: Vec, + #[arg(long, default_value = "3131")] + port: u16, + #[arg(long)] + json: bool, + }, + /// Clear inactive entries while inference falls back to cold prefill. + Clear { + /// Exact numerical model identity; omit to clear the full root. + #[arg(long)] + model_identity: Option, + #[arg(long)] + yes: bool, + /// Authenticated owner-control endpoint; repeat for multiple owned nodes. + #[arg(long = "endpoint")] + endpoints: Vec, + #[arg(long, default_value = "3131")] + port: u16, + #[arg(long)] + json: bool, + }, +} + #[derive(Subcommand, Debug)] pub enum ConfigCommand { /// Validate a config TOML file without starting a node. diff --git a/crates/mesh-llm-commands/src/gpus/tune/apply_write_tests.rs b/crates/mesh-llm-commands/src/gpus/tune/apply_write_tests.rs index 14da1a70c0..da3f333f3e 100644 --- a/crates/mesh-llm-commands/src/gpus/tune/apply_write_tests.rs +++ b/crates/mesh-llm-commands/src/gpus/tune/apply_write_tests.rs @@ -90,7 +90,7 @@ fn gpu_tune_apply_preserves_comments_and_writes_nested_fields() { .lines() .any(|line| line.trim() == "flash_attention = \"enabled\"") ); - assert!(!prefix.lines().any(|line| line.trim() == "ubatch = 128")); + assert!(!prefix.lines().any(|line| line.trim() == "ubatch = 512")); assert!( model_fit_section .lines() @@ -109,7 +109,7 @@ fn gpu_tune_apply_preserves_comments_and_writes_nested_fields() { assert!( model_fit_section .lines() - .any(|line| line.trim() == "ubatch = 128") + .any(|line| line.trim() == "ubatch = 512") ); assert!( !model_fit_section @@ -270,7 +270,7 @@ fn gpu_tune_replace_existing_writes_nested_recommendations_over_legacy_manual_fi assert_eq!(model_fit.cache_type_v.as_deref(), Some("q8_0")); assert_eq!(model_fit.ctx_size, Some(65_536)); assert_eq!(model_fit.batch, Some(512)); - assert_eq!(model_fit.ubatch, Some(128)); + assert_eq!(model_fit.ubatch, Some(512)); } #[test] diff --git a/crates/mesh-llm-commands/src/gpus/tune/planning.rs b/crates/mesh-llm-commands/src/gpus/tune/planning.rs index 1a7aa1f427..bb1d331bf3 100644 --- a/crates/mesh-llm-commands/src/gpus/tune/planning.rs +++ b/crates/mesh-llm-commands/src/gpus/tune/planning.rs @@ -1,7 +1,7 @@ use super::*; const BUILTIN_BATCH: u32 = 512; -const BUILTIN_UBATCH: u32 = 128; +const BUILTIN_UBATCH: u32 = 512; const BUILTIN_SAFETY_MARGIN_GB: f64 = 2.0; const LARGE_MODEL_MIN_BYTES: u64 = 50 * 1024 * 1024 * 1024; const MIN_AUTO_CONTEXT_LENGTH: u32 = 512; diff --git a/crates/mesh-llm-commands/src/gpus/tune/recommendation_defaults_tests.rs b/crates/mesh-llm-commands/src/gpus/tune/recommendation_defaults_tests.rs index 6adbbe9f36..d0e1087b27 100644 --- a/crates/mesh-llm-commands/src/gpus/tune/recommendation_defaults_tests.rs +++ b/crates/mesh-llm-commands/src/gpus/tune/recommendation_defaults_tests.rs @@ -1,6 +1,6 @@ use mesh_llm_config::{ FlashAttentionType, HardwareConfig, IntegerOrString, MeshConfig, ModelConfigDefaults, - ModelConfigEntry, ModelFitConfig, + ModelConfigEntry, ModelFitConfig, ThroughputConfig, }; use super::*; @@ -21,11 +21,40 @@ fn gpu_tune_recommends_stable_defaults() { assert_applied_flash_attention(&plan, TuneFlashAttentionValue::Enabled); assert_applied_context(&plan, 131_072); assert_applied_batch(&plan, 512); - assert_applied_ubatch(&plan, 128); + assert_applied_ubatch(&plan, 512); assert_applied_gpu_layers(&plan, TuneGpuLayersValue::All); assert_applied_fit_target(&plan, 22 * 1024); } +#[test] +fn gpu_tune_does_not_shadow_throughput_profile_ubatch() { + let config = MeshConfig { + defaults: Some(ModelConfigDefaults { + throughput: Some(ThroughputConfig { + tuning_profile: Some("throughput".to_string()), + ..ThroughputConfig::default() + }), + ..ModelConfigDefaults::default() + }), + models: vec![ModelConfigEntry { + model: "hf://mesh/example.gguf".to_string(), + ..ModelConfigEntry::default() + }], + ..MeshConfig::default() + }; + let plan = build_tune_plan(TuneRecommendationInput { + apply_mode: TuneApplyMode::ApplyMissing, + config: &config, + target: &recommendation_target(true), + metadata: &sample_metadata(8 * gib(), 32, 131_072, 0), + hardware: &gpu_hardware(24 * gib()), + survey: &survey_with_gpu(24 * gib(), 64 * gib()), + }); + + assert_applied_batch(&plan, 512); + assert_preserved(&plan, TuneField::Ubatch, "throughput tuning profile"); +} + #[test] fn gpu_tune_uses_q4_policy_for_large_models() { let plan = build_tune_plan(TuneRecommendationInput { diff --git a/crates/mesh-llm-commands/src/gpus/tune/recommendation_writes.rs b/crates/mesh-llm-commands/src/gpus/tune/recommendation_writes.rs index 8bd6ead853..d3dee9278c 100644 --- a/crates/mesh-llm-commands/src/gpus/tune/recommendation_writes.rs +++ b/crates/mesh-llm-commands/src/gpus/tune/recommendation_writes.rs @@ -132,6 +132,16 @@ pub(crate) fn push_batch_statuses( existing_ubatch_source(model_entry, defaults), ), ] { + if field == TuneField::Ubatch + && source.is_none() + && effective_tuning_profile(model_entry, defaults).is_some() + { + plan.field_statuses.push(TuneFieldStatus::Preserved { + field, + reason: "the effective throughput tuning profile remains authoritative".to_string(), + }); + continue; + } if let Some(source) = source && apply_mode != TuneApplyMode::ReplaceExisting { @@ -158,6 +168,16 @@ pub(crate) fn push_batch_statuses( } } +fn effective_tuning_profile<'a>( + model_entry: Option<&'a ModelConfigEntry>, + defaults: Option<&'a ModelConfigDefaults>, +) -> Option<&'a str> { + model_entry + .and_then(|entry| entry.throughput.as_ref()) + .and_then(|throughput| throughput.tuning_profile.as_deref()) + .or_else(|| defaults?.throughput.as_ref()?.tuning_profile.as_deref()) +} + pub(crate) fn push_gpu_layers_status( plan: &mut TunePlan, apply_mode: TuneApplyMode, diff --git a/crates/mesh-llm-commands/src/kv_cache.rs b/crates/mesh-llm-commands/src/kv_cache.rs new file mode 100644 index 0000000000..ab7e22630e --- /dev/null +++ b/crates/mesh-llm-commands/src/kv_cache.rs @@ -0,0 +1,217 @@ +use std::io::{self, IsTerminal, Write}; + +use anyhow::{Context, Result, bail}; +use mesh_llm_cli::KvCacheCommand; +use serde_json::{Value, json}; + +const KV_CACHE_HTTP_TIMEOUT_SECS: u64 = 60; + +pub async fn dispatch_kv_cache_command(command: &KvCacheCommand) -> Result<()> { + match command { + KvCacheCommand::Status { + endpoints, + port, + json, + } => { + let value = if endpoints.is_empty() { + request(*port, reqwest::Method::GET, "/api/runtime/kv-cache", None).await? + } else { + remote_request(*port, endpoints, "status", None, None).await? + }; + print_response(&value, *json) + } + KvCacheCommand::Prune { + target, + model_identity, + yes, + endpoints, + port, + json: json_output, + } => { + confirm_destructive("prune inactive disk prompt-cache entries", *yes)?; + let target_bytes = target + .as_deref() + .map(mesh_llm_config::parse_iec_size) + .transpose() + .context("invalid --target")?; + let body = json!({ + "target_bytes": target_bytes, + "model_identity": model_identity, + }); + let value = if endpoints.is_empty() { + request( + *port, + reqwest::Method::POST, + "/api/runtime/kv-cache/prune", + Some(body), + ) + .await? + } else { + remote_request( + *port, + endpoints, + "prune", + target_bytes, + model_identity.clone(), + ) + .await? + }; + print_response(&value, *json_output) + } + KvCacheCommand::Clear { + model_identity, + yes, + endpoints, + port, + json: json_output, + } => { + confirm_destructive("clear inactive disk prompt-cache entries", *yes)?; + let value = if endpoints.is_empty() { + request( + *port, + reqwest::Method::DELETE, + "/api/runtime/kv-cache", + Some(json!({ "model_identity": model_identity })), + ) + .await? + } else { + remote_request(*port, endpoints, "clear", None, model_identity.clone()).await? + }; + print_response(&value, *json_output) + } + } +} + +async fn remote_request( + port: u16, + endpoints: &[String], + operation: &str, + target_bytes: Option, + model_identity: Option, +) -> Result { + request( + port, + reqwest::Method::POST, + "/api/runtime/control/kv-cache", + Some(json!({ + "endpoints": endpoints, + "operation": operation, + "target_bytes": target_bytes, + "model_identity": model_identity, + })), + ) + .await +} + +async fn request( + port: u16, + method: reqwest::Method, + path: &str, + body: Option, +) -> Result { + let client = reqwest::Client::builder() + // The runtime bounds a remote batch at 45 seconds and returns one + // result per requested endpoint, including explicit timeout receipts. + // Leave transport and JSON-decoding margin outside that server bound. + .timeout(std::time::Duration::from_secs(KV_CACHE_HTTP_TIMEOUT_SECS)) + .build()?; + let url = format!("http://127.0.0.1:{port}{path}"); + let mut request = client.request(method, &url); + if let Some(body) = body { + request = request.json(&body); + } + let response = request + .send() + .await + .with_context(|| format!("request disk prompt-cache endpoint {url}"))?; + let status = response.status(); + let value = response + .json::() + .await + .with_context(|| format!("decode disk prompt-cache response from {url}"))?; + if !status.is_success() { + bail!( + "disk prompt-cache request failed ({status}): {}", + value + .get("error") + .and_then(Value::as_str) + .unwrap_or("unknown error") + ); + } + Ok(value) +} + +fn confirm_destructive(action: &str, yes: bool) -> Result<()> { + if yes { + return Ok(()); + } + if !io::stdin().is_terminal() { + bail!("{action} requires --yes when stdin is not interactive"); + } + eprint!("{action}? [y/N] "); + io::stderr().flush()?; + let mut response = String::new(); + io::stdin().read_line(&mut response)?; + if !matches!(response.trim().to_ascii_lowercase().as_str(), "y" | "yes") { + bail!("operation cancelled"); + } + Ok(()) +} + +fn print_response(value: &Value, json_output: bool) -> Result<()> { + if json_output { + println!("{}", serde_json::to_string(value)?); + return Ok(()); + } + if let Some(results) = value.get("results").and_then(Value::as_array) { + for result in results { + let node = result + .get("target_node_id") + .and_then(Value::as_str) + .unwrap_or("unknown"); + if let Some(error) = result.get("error").filter(|value| !value.is_null()) { + println!( + "Node {node}: error: {}", + error + .get("message") + .and_then(Value::as_str) + .unwrap_or("unknown error") + ); + } else if let Some(freed) = result.get("freed_bytes").and_then(Value::as_u64) { + println!("Node {node}: freed {freed} bytes"); + } else { + let state = result + .get("status") + .and_then(|status| status.get("effective")) + .and_then(|effective| effective.get("state")) + .and_then(Value::as_str) + .unwrap_or("unknown"); + println!("Node {node}: {state}"); + } + } + return Ok(()); + } + if let Some(freed) = value.get("freed_bytes").and_then(Value::as_u64) { + println!("Freed {freed} bytes"); + return Ok(()); + } + let configured = &value["configured"]; + let effective = &value["effective"]; + println!( + "Disk prompt cache: {} ({})", + effective["state"].as_str().unwrap_or("unknown"), + configured["mode"].as_str().unwrap_or("unknown") + ); + println!( + "Root: {}", + configured["directory"].as_str().unwrap_or("unknown") + ); + if let Some(usage) = value.get("usage").filter(|usage| !usage.is_null()) { + println!( + "Used: {} / {} bytes", + usage["used_bytes"].as_u64().unwrap_or(0), + usage["budget_bytes"].as_u64().unwrap_or(0) + ); + } + Ok(()) +} diff --git a/crates/mesh-llm-commands/src/lib.rs b/crates/mesh-llm-commands/src/lib.rs index 209071d410..7bdd5ce9f9 100644 --- a/crates/mesh-llm-commands/src/lib.rs +++ b/crates/mesh-llm-commands/src/lib.rs @@ -6,6 +6,7 @@ pub mod benchmark; pub mod config; pub mod doctor; pub mod gpus; +pub mod kv_cache; pub mod model_package; pub mod operational_logging; pub mod plugin; diff --git a/crates/mesh-llm-commands/src/operational_logging.rs b/crates/mesh-llm-commands/src/operational_logging.rs index cb1f68e347..4ac214e596 100644 --- a/crates/mesh-llm-commands/src/operational_logging.rs +++ b/crates/mesh-llm-commands/src/operational_logging.rs @@ -144,6 +144,7 @@ pub fn command_family(command: &Command) -> CliCommandFamily { Command::Serve | Command::Client | Command::Runtime { .. } + | Command::KvCache { .. } | Command::Load { .. } | Command::Unload { .. } | Command::Status { .. } diff --git a/crates/mesh-llm-commands/src/operational_logging/command_summary.rs b/crates/mesh-llm-commands/src/operational_logging/command_summary.rs index 799720e1e9..f50529172a 100644 --- a/crates/mesh-llm-commands/src/operational_logging/command_summary.rs +++ b/crates/mesh-llm-commands/src/operational_logging/command_summary.rs @@ -2,6 +2,7 @@ mod administration; mod auth; mod benchmark; mod dispatch; +mod kv_cache; mod models; mod runtime; diff --git a/crates/mesh-llm-commands/src/operational_logging/command_summary/dispatch.rs b/crates/mesh-llm-commands/src/operational_logging/command_summary/dispatch.rs index 14a0a980ea..c6a1d7344a 100644 --- a/crates/mesh-llm-commands/src/operational_logging/command_summary/dispatch.rs +++ b/crates/mesh-llm-commands/src/operational_logging/command_summary/dispatch.rs @@ -2,7 +2,7 @@ use mesh_llm_cli::Command; use super::{ DEFAULT_AGENT_PORT, DEFAULT_LOCAL_PORT, ModelPrepareSummary, SummaryAssembly, administration, - auth, benchmark, models, runtime, + auth, benchmark, kv_cache, models, runtime, }; pub(super) fn format_command(command: &Command, assembly: &mut SummaryAssembly) { @@ -11,6 +11,7 @@ pub(super) fn format_command(command: &Command, assembly: &mut SummaryAssembly) Command::Client => assembly.command.push_str(" client"), Command::Models { command } => models::format_models(command, assembly), Command::Runtime { command } => runtime::format_runtime(command.as_ref(), assembly), + Command::KvCache { command } => kv_cache::format_kv_cache(command, assembly), Command::Plugin { command } => administration::format_plugin(command, assembly), Command::Auth { command } => auth::format_auth(command, assembly), Command::Benchmark { command } => benchmark::format_benchmark(command, assembly), diff --git a/crates/mesh-llm-commands/src/operational_logging/command_summary/kv_cache.rs b/crates/mesh-llm-commands/src/operational_logging/command_summary/kv_cache.rs new file mode 100644 index 0000000000..07541c74d3 --- /dev/null +++ b/crates/mesh-llm-commands/src/operational_logging/command_summary/kv_cache.rs @@ -0,0 +1,49 @@ +use mesh_llm_cli::KvCacheCommand; + +use super::{DEFAULT_LOCAL_PORT, SummaryAssembly}; + +pub(super) fn format_kv_cache(command: &KvCacheCommand, assembly: &mut SummaryAssembly) { + assembly.command.push_str(" kv-cache"); + match command { + KvCacheCommand::Status { + endpoints, + port, + json, + } => { + assembly.command.push_str(" status"); + assembly.redact("--endpoint", !endpoints.is_empty()); + assembly.port(*port, DEFAULT_LOCAL_PORT); + assembly.flag("json", *json); + } + KvCacheCommand::Prune { + target, + model_identity, + yes, + endpoints, + port, + json, + } => { + assembly.command.push_str(" prune"); + assembly.redact("--target", target.is_some()); + assembly.redact("--model-identity", model_identity.is_some()); + assembly.redact("--endpoint", !endpoints.is_empty()); + assembly.port(*port, DEFAULT_LOCAL_PORT); + assembly.flag("yes", *yes); + assembly.flag("json", *json); + } + KvCacheCommand::Clear { + model_identity, + yes, + endpoints, + port, + json, + } => { + assembly.command.push_str(" clear"); + assembly.redact("--model-identity", model_identity.is_some()); + assembly.redact("--endpoint", !endpoints.is_empty()); + assembly.port(*port, DEFAULT_LOCAL_PORT); + assembly.flag("yes", *yes); + assembly.flag("json", *json); + } + } +} diff --git a/crates/mesh-llm-commands/src/operational_logging/command_summary_tests.rs b/crates/mesh-llm-commands/src/operational_logging/command_summary_tests.rs index fb50b3f720..7a3efbcdfa 100644 --- a/crates/mesh-llm-commands/src/operational_logging/command_summary_tests.rs +++ b/crates/mesh-llm-commands/src/operational_logging/command_summary_tests.rs @@ -131,6 +131,40 @@ fn command_summary_covers_plugin_config_and_doctor_values_without_defaults() { ); } +#[test] +fn command_summary_records_kv_cache_subcommands_and_redacts_values() { + assert_eq!( + parsed_summary(&[ + "mesh-llm", + "kv-cache", + "status", + "--endpoint", + "secret-endpoint", + "--port", + "4444", + "--json", + ]), + "mesh-llm kv-cache status --json --port 4444 --endpoint [REDACTED]" + ); + assert_eq!( + parsed_summary(&[ + "mesh-llm", + "kv-cache", + "prune", + "--target", + "4GiB", + "--model-identity", + "private-model", + "--yes", + ]), + "mesh-llm kv-cache prune --yes --target [REDACTED] --model-identity [REDACTED]" + ); + assert_eq!( + parsed_summary(&["mesh-llm", "kv-cache", "clear", "--yes"]), + "mesh-llm kv-cache clear --yes" + ); +} + #[test] fn command_summary_covers_auth_trust_and_nested_command_families() { let auth = parsed_summary(&[ diff --git a/crates/mesh-llm-config/src/lib.rs b/crates/mesh-llm-config/src/lib.rs index 700277adb3..a8d9d7e870 100644 --- a/crates/mesh-llm-config/src/lib.rs +++ b/crates/mesh-llm-config/src/lib.rs @@ -5,6 +5,7 @@ mod hardware_validation; mod model; mod model_validation; mod plugin_validation; +mod size; mod store; mod validate; mod validation_support; @@ -41,6 +42,7 @@ pub use plugin_validation::{ PluginSettingConstraint, PluginSettingSchema, PluginValueKind, PluginValueSchema, SUPPORTED_PLUGIN_CONFIG_SCHEMA_VERSION, }; +pub use size::{IecSizeParseError, parse_iec_size}; pub use store::{ ConfigStore, config_path, config_to_toml, load_config, parse_config_toml, parse_config_toml_structural, @@ -61,10 +63,10 @@ pub use wiring_validation::wiring_manifest_diagnostics; #[cfg(test)] mod tests { use super::{ - ConfigStore, ConfigValueSource, GpuAssignment, LifecycleLogParserMode, - LocalServingNodeConfig, MeshConfig, SpeculativeConfig, built_in_config_schema, - canonicalize_built_in_config_identifier, parse_config_toml, - resolve_lifecycle_log_parser_override, validate_config, + ConfigApplyMode, ConfigRestartScope, ConfigStore, ConfigValueSource, GpuAssignment, + KvDiskTierMode, LifecycleLogParserMode, LocalServingNodeConfig, MeshConfig, + SpeculativeConfig, built_in_config_schema, canonicalize_built_in_config_identifier, + parse_config_toml, resolve_lifecycle_log_parser_override, validate_config, }; use std::collections::{BTreeMap, BTreeSet}; use std::fs; @@ -292,6 +294,70 @@ skippy_abi = "0.1.25" ); } + #[test] + fn disk_cache_config_parses_defaults_and_fixed_mode() { + let defaults = parse_config_toml("").expect("empty config should parse"); + assert_eq!(defaults.runtime.kv_cache.disk.mode, None); + assert_eq!( + defaults.runtime.kv_cache.disk.effective_mode(), + KvDiskTierMode::Off + ); + assert_eq!(defaults.runtime.kv_cache.disk.minimum_free_mib, None); + assert_eq!( + defaults.runtime.kv_cache.disk.effective_minimum_free_mib(), + 16_384 + ); + assert_eq!(defaults.runtime.kv_cache.disk.budget_mib, None); + + let fixed = parse_config_toml( + r#" +[runtime.kv_cache.disk] +mode = "fixed" +directory = "/fast-disk/mesh-kv-cache" +budget_mib = 32768 +minimum_free_mib = 16384 +"#, + ) + .expect("fixed disk-cache config should parse"); + assert_eq!( + fixed.runtime.kv_cache.disk.mode, + Some(KvDiskTierMode::Fixed) + ); + assert_eq!(fixed.runtime.kv_cache.disk.budget_mib, Some(32_768)); + } + + #[test] + fn disk_cache_config_rejects_invalid_mode_budget_and_path_combinations() { + for (raw, expected) in [ + ( + "[runtime.kv_cache.disk]\nmode = \"fixed\"\n", + "budget_mib is required", + ), + ( + "[runtime.kv_cache.disk]\nmode = \"auto\"\nbudget_mib = 1024\n", + "budget_mib is only valid", + ), + ( + "[runtime.kv_cache.disk]\nmode = \"fixed\"\nbudget_mib = 0\n", + "must be greater than zero", + ), + ( + "[runtime.kv_cache.disk]\ndirectory = \"relative/cache\"\n", + "must be an absolute path", + ), + ( + "[runtime.kv_cache.disk]\nminimum_free_mib = 1023\n", + "must be at least 1024", + ), + ] { + let error = parse_config_toml(raw).expect_err("invalid disk-cache config must fail"); + assert!( + error.to_string().contains(expected), + "unexpected error: {error}" + ); + } + } + #[test] fn native_runtime_override_rejects_unknown_backend_selection() { let err = parse_config_toml( @@ -341,6 +407,35 @@ selection = "vulcan" } } + #[test] + fn disk_cache_schema_marks_static_and_live_fields_correctly() { + let schema = built_in_config_schema(); + let setting = |path: &str| { + schema + .settings + .iter() + .find(|setting| setting.path.render() == path) + .unwrap_or_else(|| panic!("missing schema setting {path}")) + }; + for path in [ + "runtime.kv_cache.disk.mode", + "runtime.kv_cache.disk.directory", + ] { + assert_eq!(setting(path).apply_mode, ConfigApplyMode::StaticOnLoad); + assert_eq!( + setting(path).restart_scope, + ConfigRestartScope::ProcessRestart + ); + } + for path in [ + "runtime.kv_cache.disk.budget_mib", + "runtime.kv_cache.disk.minimum_free_mib", + ] { + assert_eq!(setting(path).apply_mode, ConfigApplyMode::DynamicApply); + assert_eq!(setting(path).restart_scope, ConfigRestartScope::None); + } + } + #[test] fn config_store_add_model_preserves_existing_fields() { let temp_dir = TempDir::new().unwrap(); @@ -884,6 +979,8 @@ gpu_id = "pci:0000:65:00.0" ("GpuConfig", 1), ("RuntimeConfig", 1), ("NativeRuntimeConfig", 1), + ("RuntimeKvCacheConfig", 1), + ("KvDiskTierConfig", 1), ("MeshRequirementsConfig", 1), ("ModelConfigEntry", 1), ("ModelFitConfig", 2), @@ -912,6 +1009,8 @@ gpu_id = "pci:0000:65:00.0" "OwnerControlConfig", "RuntimeConfig", "NativeRuntimeConfig", + "RuntimeKvCacheConfig", + "KvDiskTierConfig", "TelemetryConfig", "TelemetryMetricsConfig", "AuditConfig", diff --git a/crates/mesh-llm-config/src/model.rs b/crates/mesh-llm-config/src/model.rs index 434776daed..9665c2f1ec 100644 --- a/crates/mesh-llm-config/src/model.rs +++ b/crates/mesh-llm-config/src/model.rs @@ -118,6 +118,9 @@ pub struct RuntimeConfig { pub reconcile_model_target_demand_upgrades: bool, #[serde(default)] pub native_runtime: NativeRuntimeConfig, + /// Node-wide resident and durable KV-cache policy. + #[serde(default)] + pub kv_cache: RuntimeKvCacheConfig, #[serde(default = "default_model_target_demand_upgrade_min_requests")] pub model_target_demand_upgrade_min_requests: u64, #[serde(default = "default_model_target_demand_upgrade_max_age_secs")] @@ -139,6 +142,7 @@ impl Default for RuntimeConfig { reconcile_model_targets: false, reconcile_model_target_demand_upgrades: false, native_runtime: NativeRuntimeConfig::default(), + kv_cache: RuntimeKvCacheConfig::default(), model_target_demand_upgrade_min_requests: DEFAULT_MODEL_TARGET_DEMAND_UPGRADE_MIN_REQUESTS, model_target_demand_upgrade_max_age_secs: @@ -179,6 +183,49 @@ impl LifecycleLogParserMode { } } +pub const DEFAULT_KV_DISK_MINIMUM_FREE_MIB: u64 = 16 * 1024; +pub const MIN_KV_DISK_MINIMUM_FREE_MIB: u64 = 1024; + +#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum KvDiskTierMode { + #[default] + Off, + Auto, + Fixed, +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq, Eq)] +pub struct RuntimeKvCacheConfig { + #[serde(default)] + pub disk: KvDiskTierConfig, +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq, Eq)] +pub struct KvDiskTierConfig { + #[serde(default)] + pub mode: Option, + /// Absolute node-local root. Absence resolves to `$MESH_LLM_HOME/kv-cache`. + #[serde(default)] + pub directory: Option, + /// Hard whole-node cap, required only in fixed mode. + #[serde(default)] + pub budget_mib: Option, + #[serde(default)] + pub minimum_free_mib: Option, +} + +impl KvDiskTierConfig { + pub fn effective_mode(&self) -> KvDiskTierMode { + self.mode.unwrap_or_default() + } + + pub fn effective_minimum_free_mib(&self) -> u64 { + self.minimum_free_mib + .unwrap_or(DEFAULT_KV_DISK_MINIMUM_FREE_MIB) + } +} + #[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq, Eq)] pub struct NativeRuntimeConfig { #[serde(default)] diff --git a/crates/mesh-llm-config/src/model/built_in_schema/control_behavior/runtime_controls.rs b/crates/mesh-llm-config/src/model/built_in_schema/control_behavior/runtime_controls.rs index b42f433c06..4fa148a886 100644 --- a/crates/mesh-llm-config/src/model/built_in_schema/control_behavior/runtime_controls.rs +++ b/crates/mesh-llm-config/src/model/built_in_schema/control_behavior/runtime_controls.rs @@ -1,8 +1,8 @@ use super::shared::{ - absent_condition, equals_bool_condition, present_condition, push_allowed_pattern_constraint, - push_dependency_disable, push_non_empty_constraint, push_range_constraint, - push_requires_constraint, set_numeric, set_runtime_native_backend_options, set_static_options, - set_text_format, + absent_condition, equals_bool_condition, equals_string_condition, present_condition, + push_allowed_pattern_constraint, push_dependency_disable, push_non_empty_constraint, + push_range_constraint, push_requires_constraint, set_numeric, + set_runtime_native_backend_options, set_static_options, set_text_format, }; use super::*; @@ -110,6 +110,23 @@ pub(super) fn apply_runtime_controls_behavior(setting: &mut ConfigSettingSchema, "runtime.activity.response" | "runtime.activity.advertisement" => { set_static_options(setting) } + "runtime.kv_cache.disk.mode" => set_static_options(setting), + "runtime.kv_cache.disk.directory" => { + set_text_format(setting, ConfigTextFormat::Path); + push_non_empty_constraint(setting); + } + "runtime.kv_cache.disk.budget_mib" => { + set_numeric(setting, Some(1.0), None, Some(1.0), Some("MiB")); + control_behavior_mut(setting) + .enable_when + .push(equals_string_condition( + "runtime.kv_cache.disk.mode", + "fixed", + )); + } + "runtime.kv_cache.disk.minimum_free_mib" => { + set_numeric(setting, Some(1024.0), None, Some(1.0), Some("MiB")); + } _ => {} } } diff --git a/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs b/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs index ccd02db1ee..741bd946c2 100644 --- a/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs +++ b/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs @@ -113,6 +113,26 @@ fn build_built_in_config_schema() -> ConfigSchema { "runtime.native_runtime.selection", one_of([string_enum(["recommended"]), ConfigValueSchema::String]), ), + kv_disk_setting( + "runtime.kv_cache.disk.mode", + string_enum(["off", "auto", "fixed"]), + false, + ), + kv_disk_setting( + "runtime.kv_cache.disk.directory", + ConfigValueSchema::Path, + false, + ), + kv_disk_setting( + "runtime.kv_cache.disk.budget_mib", + ConfigValueSchema::Integer, + true, + ), + kv_disk_setting( + "runtime.kv_cache.disk.minimum_free_mib", + ConfigValueSchema::Integer, + true, + ), runtime_setting( "runtime.model_target_demand_upgrade_min_requests", ConfigValueSchema::Integer, diff --git a/crates/mesh-llm-config/src/model/built_in_schema/presentation.rs b/crates/mesh-llm-config/src/model/built_in_schema/presentation.rs index f11eb93e24..0023497e5c 100644 --- a/crates/mesh-llm-config/src/model/built_in_schema/presentation.rs +++ b/crates/mesh-llm-config/src/model/built_in_schema/presentation.rs @@ -47,6 +47,12 @@ const RUNTIME_POLICY_CATEGORY: CategoryPresentation = CategoryPresentation { summary: "Runtime reconciliation behavior applied by the local process", order: 10, }; +const PROMPT_CACHE_CATEGORY: CategoryPresentation = CategoryPresentation { + id: "prompt-cache", + label: "Prompt Cache", + summary: "Node-local durable prefix cache storage and capacity", + order: 15, +}; const MEMORY_CATEGORY: CategoryPresentation = CategoryPresentation { id: "memory", label: "Memory", @@ -106,6 +112,7 @@ struct SettingPresentation { fn setting_presentation_for_path(rendered: &str) -> Option { logging_presentation(rendered) + .or_else(|| kv_disk_presentation(rendered)) .or_else(|| process_setting_presentation(rendered)) .or_else(|| native_runtime_presentation(rendered)) .or_else(|| runtime_defaults_presentation(rendered)) @@ -114,6 +121,56 @@ fn setting_presentation_for_path(rendered: &str) -> Option .or_else(|| model_and_plugin_presentation(rendered)) } +fn kv_disk_presentation(rendered: &str) -> Option { + match rendered { + "runtime.kv_cache.disk.mode" => Some( + sp( + "Disk prompt cache", + "Keep compatible prompt state on local disk across model reloads and process restarts.", + PROMPT_CACHE_CATEGORY, + 10, + ) + .hint("segmented") + .choices(&[ + ("off", "Off", "Do not read or write the disk cache."), + ("auto", "Automatic", "Derive a bounded budget from available local storage."), + ("fixed", "Fixed", "Use the configured hard budget."), + ]), + ), + "runtime.kv_cache.disk.directory" => Some( + sp( + "Cache directory", + "Absolute node-local directory used for durable prompt-cache data.", + PROMPT_CACHE_CATEGORY, + 20, + ) + .placeholder("/fast-disk/mesh-kv-cache") + .hint("text"), + ), + "runtime.kv_cache.disk.budget_mib" => Some( + sp( + "Fixed cache budget", + "Hard whole-node disk-cache budget used only in fixed mode.", + PROMPT_CACHE_CATEGORY, + 30, + ) + .unit("MiB") + .hint("number"), + ), + "runtime.kv_cache.disk.minimum_free_mib" => Some( + sp( + "Minimum free space", + "Free local storage Mesh preserves before accepting cache writes.", + PROMPT_CACHE_CATEGORY, + 40, + ) + .unit("MiB") + .hint("number"), + ), + _ => None, + } +} + fn process_setting_presentation(rendered: &str) -> Option { match rendered { "gpu.assignment" => Some(sp( @@ -445,7 +502,7 @@ fn runtime_defaults_presentation(rendered: &str) -> Option .hint("range")), "defaults.model_fit.ubatch" => Some(sp( "Micro-batch size", - "Set the default decode micro-batch size.", + "Set the default micro-batch (physical prefill chunk) size. Values at or below 128 keep the CUDA SSM sequential-scan fallback; larger values enable the SSD chunked kernel for recurrent models.", MEMORY_CATEGORY, 50, ) diff --git a/crates/mesh-llm-config/src/model/built_in_schema/setting_schema.rs b/crates/mesh-llm-config/src/model/built_in_schema/setting_schema.rs index e6d82c9b4a..57fe6c0318 100644 --- a/crates/mesh-llm-config/src/model/built_in_schema/setting_schema.rs +++ b/crates/mesh-llm-config/src/model/built_in_schema/setting_schema.rs @@ -77,6 +77,30 @@ fn activity_runtime_setting(path: &str, value_schema: ConfigValueSchema) -> Conf setting } +fn kv_disk_setting( + path: &str, + value_schema: ConfigValueSchema, + dynamic: bool, +) -> ConfigSettingSchema { + let mut setting = basic_setting(path, value_schema); + setting.control_surfaces = vec![ + ConfigControlSurface::ConfigFile, + ConfigControlSurface::Cli, + ConfigControlSurface::OwnerControl, + ConfigControlSurface::Api, + ConfigControlSurface::Ui, + ]; + setting.visibility = ConfigVisibility::User; + if dynamic { + setting.apply_mode = ConfigApplyMode::DynamicApply; + setting.restart_scope = ConfigRestartScope::None; + } else { + setting.apply_mode = ConfigApplyMode::StaticOnLoad; + setting.restart_scope = ConfigRestartScope::ProcessRestart; + } + setting +} + fn plugin_setting(path: &str, value_schema: ConfigValueSchema) -> ConfigSettingSchema { let mut setting = basic_setting(path, value_schema); setting.control_surfaces = vec![ diff --git a/crates/mesh-llm-config/src/size.rs b/crates/mesh-llm-config/src/size.rs new file mode 100644 index 0000000000..9f7d3a6eed --- /dev/null +++ b/crates/mesh-llm-config/src/size.rs @@ -0,0 +1,76 @@ +use std::fmt; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct IecSizeParseError { + message: String, +} + +impl fmt::Display for IecSizeParseError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(&self.message) + } +} + +impl std::error::Error for IecSizeParseError {} + +/// Parse a positive integer size with an explicit IEC suffix. +/// +/// Bare numbers, decimal values, SI suffixes, zero, and overflow are rejected +/// so a public cache setting can never silently change units or mean +/// "unbounded". +pub fn parse_iec_size(input: &str) -> Result { + let value = input.trim(); + let (digits, multiplier) = [ + ("KiB", 1024_u64), + ("MiB", 1024_u64.pow(2)), + ("GiB", 1024_u64.pow(3)), + ("TiB", 1024_u64.pow(4)), + ] + .into_iter() + .find_map(|(suffix, multiplier)| { + value + .strip_suffix(suffix) + .map(|digits| (digits, multiplier)) + }) + .ok_or_else(|| error("size must use an explicit IEC suffix: KiB, MiB, GiB, or TiB"))?; + if digits.is_empty() || !digits.bytes().all(|byte| byte.is_ascii_digit()) { + return Err(error( + "size must be a positive whole number with an IEC suffix", + )); + } + let units = digits + .parse::() + .map_err(|_| error("size is too large"))?; + if units == 0 { + return Err(error("size must be greater than zero")); + } + units + .checked_mul(multiplier) + .ok_or_else(|| error("size is too large")) +} + +fn error(message: &str) -> IecSizeParseError { + IecSizeParseError { + message: message.to_string(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_explicit_iec_sizes() { + assert_eq!(parse_iec_size("1KiB").unwrap(), 1024); + assert_eq!(parse_iec_size("32MiB").unwrap(), 32 * 1024 * 1024); + assert_eq!(parse_iec_size("32GiB").unwrap(), 32 * 1024 * 1024 * 1024); + assert_eq!(parse_iec_size("1TiB").unwrap(), 1024_u64.pow(4)); + } + + #[test] + fn rejects_ambiguous_or_unbounded_sizes() { + for value in ["", "0", "0GiB", "32", "32GB", "1.5GiB", "-1GiB"] { + assert!(parse_iec_size(value).is_err(), "accepted {value:?}"); + } + } +} diff --git a/crates/mesh-llm-config/src/validate.rs b/crates/mesh-llm-config/src/validate.rs index d5717c7c1b..db4511486c 100644 --- a/crates/mesh-llm-config/src/validate.rs +++ b/crates/mesh-llm-config/src/validate.rs @@ -277,6 +277,81 @@ fn validate_runtime_config(config: &RuntimeConfig) -> Vec { "runtime.drain_timeout_secs must not exceed runtime.drain_timeout_max_secs", )); } + diagnostics.extend(validate_kv_disk_config(&config.kv_cache.disk)); + diagnostics +} + +fn validate_kv_disk_config(config: &KvDiskTierConfig) -> Vec { + let mut diagnostics = Vec::new(); + match config.mode { + Some(KvDiskTierMode::Fixed) if config.budget_mib.is_none() => { + diagnostics.push(validation_diagnostic( + "runtime.kv_cache.disk.budget_mib", + "runtime.kv_cache.disk.budget_mib is required when mode = \"fixed\"", + )); + } + Some(KvDiskTierMode::Off | KvDiskTierMode::Auto) | None if config.budget_mib.is_some() => { + diagnostics.push(validation_diagnostic( + "runtime.kv_cache.disk.budget_mib", + "runtime.kv_cache.disk.budget_mib is only valid when mode = \"fixed\"", + )); + } + _ => {} + } + if config.budget_mib == Some(0) { + diagnostics.push(validation_diagnostic( + "runtime.kv_cache.disk.budget_mib", + "runtime.kv_cache.disk.budget_mib must be greater than zero", + )); + } + if config + .minimum_free_mib + .is_some_and(|minimum| minimum < MIN_KV_DISK_MINIMUM_FREE_MIB) + { + diagnostics.push(validation_diagnostic( + "runtime.kv_cache.disk.minimum_free_mib", + format!( + "runtime.kv_cache.disk.minimum_free_mib must be at least {MIN_KV_DISK_MINIMUM_FREE_MIB}" + ), + )); + } + for (path, value) in [ + ("runtime.kv_cache.disk.budget_mib", config.budget_mib), + ( + "runtime.kv_cache.disk.minimum_free_mib", + config.minimum_free_mib, + ), + ] { + if value.is_some_and(|mib| mib.checked_mul(1024 * 1024).is_none()) { + diagnostics.push(validation_diagnostic( + path, + format!("{path} is too large to represent as bytes"), + )); + } + } + if let Some(directory) = config.directory.as_deref() { + let rendered = directory.to_string_lossy(); + // A drive letter is absolute only where the host understands one. On + // Unix `C:\cache` is a relative path, and accepting it would put a + // durable cache under the process working directory. + let windows_absolute = cfg!(windows) + && rendered.as_bytes().get(1) == Some(&b':') + && rendered + .as_bytes() + .get(2) + .is_some_and(|byte| matches!(byte, b'/' | b'\\')); + if rendered.trim().is_empty() { + diagnostics.push(validation_diagnostic( + "runtime.kv_cache.disk.directory", + "runtime.kv_cache.disk.directory must not be empty", + )); + } else if !directory.is_absolute() && !windows_absolute { + diagnostics.push(validation_diagnostic( + "runtime.kv_cache.disk.directory", + "runtime.kv_cache.disk.directory must be an absolute path", + )); + } + } diagnostics } diff --git a/crates/mesh-llm-config/src/wiring_status.rs b/crates/mesh-llm-config/src/wiring_status.rs index a924b901c5..caa7c63640 100644 --- a/crates/mesh-llm-config/src/wiring_status.rs +++ b/crates/mesh-llm-config/src/wiring_status.rs @@ -17,6 +17,7 @@ //! validate` agrees with real runtime support instead of drifting from it. mod checkpoint; +mod runtime; mod topology; /// The wiring status of one canonical config path, matching the `Status` @@ -510,6 +511,10 @@ pub const WIRING_MANIFEST: &[WiringEntry] = &[ reason: "", behavior: WiringBehavior::None, }, + runtime::KV_CACHE_DISK_MODE, + runtime::KV_CACHE_DISK_DIRECTORY, + runtime::KV_CACHE_DISK_BUDGET_MIB, + runtime::KV_CACHE_DISK_MINIMUM_FREE_MIB, WiringEntry { path: "runtime.model_target_demand_upgrade_min_requests", status: WiringStatus::Wired, diff --git a/crates/mesh-llm-config/src/wiring_status/runtime.rs b/crates/mesh-llm-config/src/wiring_status/runtime.rs new file mode 100644 index 0000000000..990eaee6ea --- /dev/null +++ b/crates/mesh-llm-config/src/wiring_status/runtime.rs @@ -0,0 +1,19 @@ +use super::{WiringBehavior, WiringEntry, WiringStatus}; + +pub(super) const KV_CACHE_DISK_MODE: WiringEntry = wired_kv_cache("runtime.kv_cache.disk.mode"); +pub(super) const KV_CACHE_DISK_DIRECTORY: WiringEntry = + wired_kv_cache("runtime.kv_cache.disk.directory"); +pub(super) const KV_CACHE_DISK_BUDGET_MIB: WiringEntry = + wired_kv_cache("runtime.kv_cache.disk.budget_mib"); +pub(super) const KV_CACHE_DISK_MINIMUM_FREE_MIB: WiringEntry = + wired_kv_cache("runtime.kv_cache.disk.minimum_free_mib"); + +const fn wired_kv_cache(path: &'static str) -> WiringEntry { + WiringEntry { + path, + status: WiringStatus::Wired, + owner: "#1576", + reason: "", + behavior: WiringBehavior::None, + } +} diff --git a/crates/mesh-llm-events/src/command_summary_grammar/descriptors.rs b/crates/mesh-llm-events/src/command_summary_grammar/descriptors.rs index fbf079db91..7aa0b112bd 100644 --- a/crates/mesh-llm-events/src/command_summary_grammar/descriptors.rs +++ b/crates/mesh-llm-events/src/command_summary_grammar/descriptors.rs @@ -1,5 +1,7 @@ #[path = "descriptors/auth.rs"] mod auth; +#[path = "descriptors/kv_cache.rs"] +mod kv_cache; #[path = "descriptors/models.rs"] mod models; #[path = "descriptors/plugins_benchmark.rs"] @@ -68,5 +70,6 @@ pub(super) const DESCRIPTOR_GROUPS: &[&[Descriptor]] = &[ models::DESCRIPTORS, plugins_benchmark::BENCHMARK_DESCRIPTORS, runtime::DESCRIPTORS, + kv_cache::DESCRIPTORS, auth::DESCRIPTORS, ]; diff --git a/crates/mesh-llm-events/src/command_summary_grammar/descriptors/kv_cache.rs b/crates/mesh-llm-events/src/command_summary_grammar/descriptors/kv_cache.rs new file mode 100644 index 0000000000..bb79d5e34f --- /dev/null +++ b/crates/mesh-llm-events/src/command_summary_grammar/descriptors/kv_cache.rs @@ -0,0 +1,25 @@ +use super::{Descriptor, JSON, RawKind, YES_JSON, descriptor}; + +pub(super) const DESCRIPTORS: &[Descriptor] = &[ + descriptor( + &["mesh-llm", "kv-cache", "status"], + JSON, + &["--endpoint"], + true, + RawKind::None, + ), + descriptor( + &["mesh-llm", "kv-cache", "prune"], + YES_JSON, + &["--target", "--model-identity", "--endpoint"], + true, + RawKind::None, + ), + descriptor( + &["mesh-llm", "kv-cache", "clear"], + YES_JSON, + &["--model-identity", "--endpoint"], + true, + RawKind::None, + ), +]; diff --git a/crates/mesh-llm-events/src/command_summary_grammar/raw_options.rs b/crates/mesh-llm-events/src/command_summary_grammar/raw_options.rs index 64d43bf468..668dc44539 100644 --- a/crates/mesh-llm-events/src/command_summary_grammar/raw_options.rs +++ b/crates/mesh-llm-events/src/command_summary_grammar/raw_options.rs @@ -110,5 +110,15 @@ fn matches_port_prefix(prefix: &[&str]) -> bool { | ["mesh-llm", "runtime", "remote-model", "--json"] | ["mesh-llm", "runtime", "apply-config"] | ["mesh-llm", "runtime", "apply-config", "--json"] + | ["mesh-llm", "kv-cache", "status"] + | ["mesh-llm", "kv-cache", "status", "--json"] + | ["mesh-llm", "kv-cache", "prune"] + | ["mesh-llm", "kv-cache", "prune", "--yes"] + | ["mesh-llm", "kv-cache", "prune", "--json"] + | ["mesh-llm", "kv-cache", "prune", "--yes", "--json"] + | ["mesh-llm", "kv-cache", "clear"] + | ["mesh-llm", "kv-cache", "clear", "--yes"] + | ["mesh-llm", "kv-cache", "clear", "--json"] + | ["mesh-llm", "kv-cache", "clear", "--yes", "--json"] ) } diff --git a/crates/mesh-llm-events/src/command_summary_grammar/vocabulary.rs b/crates/mesh-llm-events/src/command_summary_grammar/vocabulary.rs index 31ddf3aa8a..e454a52f59 100644 --- a/crates/mesh-llm-events/src/command_summary_grammar/vocabulary.rs +++ b/crates/mesh-llm-events/src/command_summary_grammar/vocabulary.rs @@ -39,6 +39,7 @@ pub(super) fn is_static_summary_token(token: &str) -> bool { | "installed" | "cleanup" | "prune" + | "clear" | "certify" | "show" | "download" @@ -48,6 +49,7 @@ pub(super) fn is_static_summary_token(token: &str) -> bool { | "import-prompts" | "model-prepare" | "runtime" + | "kv-cache" | "guardrails" | "bootstrap" | "remove" @@ -184,6 +186,7 @@ pub(super) fn is_redacted_marker(token: &str) -> bool { | "--cache-dir" | "--mesh-version" | "--endpoint" + | "--model-identity" | "--profile" | "--instance-id" | "--expected-revision" diff --git a/crates/mesh-llm-host-runtime/Cargo.toml b/crates/mesh-llm-host-runtime/Cargo.toml index 6106c7c772..aec04a639d 100644 --- a/crates/mesh-llm-host-runtime/Cargo.toml +++ b/crates/mesh-llm-host-runtime/Cargo.toml @@ -53,6 +53,7 @@ model-ref = { path = "../model-ref", version = "0.76.1" } model-resolver = { path = "../model-resolver", version = "0.76.1" } openai-frontend = { path = "../openai-frontend", version = "0.76.1" } skippy-protocol = { path = "../skippy-protocol", version = "0.76.1" } +skippy-cache = { path = "../skippy-cache", version = "0.76.1" } skippy-coordinator = { path = "../skippy-coordinator", version = "0.76.1" } skippy-runtime = { path = "../skippy-runtime", version = "0.76.1" } skippy-ffi = { path = "../skippy-ffi", version = "0.76.1", default-features = false } diff --git a/crates/mesh-llm-host-runtime/src/api/mod.rs b/crates/mesh-llm-host-runtime/src/api/mod.rs index f231a698c4..ebedea7218 100644 --- a/crates/mesh-llm-host-runtime/src/api/mod.rs +++ b/crates/mesh-llm-host-runtime/src/api/mod.rs @@ -24,6 +24,7 @@ //! POST /api/runtime/control/get-config — run local owner-control get-config against an explicit endpoint //! POST /api/runtime/control/refresh-inventory — run local owner-control refresh-inventory against an explicit endpoint //! POST /api/runtime/control/apply-config — run local owner-control apply-config against an explicit endpoint +//! POST /api/runtime/control/kv-cache — run authenticated L3 KV cache operations against owned nodes //! POST /api/runtime/models — load a local model //! DELETE /api/runtime/models/{model} — unload a local model //! DELETE /api/runtime/instances/{instance_id} — unload one local runtime instance diff --git a/crates/mesh-llm-host-runtime/src/api/routes/kv_cache.rs b/crates/mesh-llm-host-runtime/src/api/routes/kv_cache.rs new file mode 100644 index 0000000000..8896b09e8d --- /dev/null +++ b/crates/mesh-llm-host-runtime/src/api/routes/kv_cache.rs @@ -0,0 +1,306 @@ +use serde::{Deserialize, Serialize}; +use skippy_cache::{ + L3ActivitySnapshot, L3EffectiveState, L3EffectiveStatus, L3InventoryEntry, L3StateReason, + MANIFEST_VERSION, StoreReconciliation, StoreUsage, +}; +use tokio::net::TcpStream; + +use super::runtime::ensure_loopback_control_caller; +use crate::api::http::{respond_error, respond_json}; +use crate::runtime::kv_disk_config::{ + KvDiskConfigSources, node_kv_disk_cache, node_kv_disk_manager, +}; + +#[derive(Debug, Serialize)] +struct KvCacheConfiguredPayload { + mode: mesh_llm_config::KvDiskTierMode, + directory: std::path::PathBuf, + budget_bytes: Option, + minimum_free_bytes: u64, + sources: KvDiskConfigSources, +} + +#[derive(Debug, Serialize)] +struct KvCacheEffectivePayload { + state: String, + reason: Option, + manager: Option, +} + +#[derive(Debug, Serialize)] +pub(crate) struct KvCacheStatusPayload { + version: u32, + format_version: u32, + configured: KvCacheConfiguredPayload, + effective: KvCacheEffectivePayload, + usage: Option, + activity: Option, + reconciliation: Option, + inventory: Vec, +} + +#[derive(Debug, Default, Deserialize)] +struct KvCachePruneRequest { + target_bytes: Option, + model_identity: Option, +} + +#[derive(Debug, Default, Deserialize)] +struct KvCacheClearRequest { + model_identity: Option, +} + +#[derive(Clone, Debug)] +pub(crate) enum KvCacheOperation { + Status, + Prune { + target_bytes: Option, + model_identity: Option, + }, + Clear { + model_identity: Option, + }, +} + +#[derive(Debug)] +pub(crate) struct KvCacheOperationResult { + pub(crate) status_json: Vec, + pub(crate) freed_bytes: Option, +} + +/// Trim and reject a blank model filter, matching `decode_operation` on the +/// owner-control path so both entry points share one contract. An untrimmed +/// identity would compare unequal against every stored manifest and report a +/// successful prune that freed nothing. +fn normalize_model_identity(identity: Option) -> Result, String> { + match identity { + Some(identity) => { + let trimmed = identity.trim(); + if trimmed.is_empty() { + return Err("model_identity must not be empty".to_string()); + } + Ok(Some(trimmed.to_string())) + } + None => Ok(None), + } +} + +pub(super) async fn handle_status(stream: &mut TcpStream) -> anyhow::Result<()> { + if !ensure_loopback_control_caller(stream).await? { + return Ok(()); + } + match status_payload() { + Ok(status) => respond_json(stream, 200, &status).await, + Err(error) => respond_error(stream, 500, &error.to_string()).await, + } +} + +pub(super) async fn handle_prune(stream: &mut TcpStream, body: &str) -> anyhow::Result<()> { + if !ensure_loopback_control_caller(stream).await? { + return Ok(()); + } + let request = if body.trim().is_empty() { + KvCachePruneRequest::default() + } else { + match serde_json::from_str::(body) { + Ok(request) => request, + Err(error) => return respond_error(stream, 400, &error.to_string()).await, + } + }; + let Some(manager) = node_kv_disk_manager() else { + return respond_error(stream, 409, "disk prompt cache is not active").await; + }; + let model_identity = match normalize_model_identity(request.model_identity) { + Ok(identity) => identity, + Err(error) => return respond_error(stream, 400, &error).await, + }; + // budget_bytes == 0 means "no cap", so the 85%-of-budget default target + // computes to 0 and a prune would remove every unpinned manifest. That is + // a clear, which the CLI gates behind explicit confirmation. + if request.target_bytes.is_none() && manager.limits().budget_bytes == 0 { + return respond_error( + stream, + 400, + "prune requires target_bytes when the disk cache has no budget; \ + use the clear endpoint to remove every entry", + ) + .await; + } + match execute_operation(KvCacheOperation::Prune { + target_bytes: request.target_bytes, + model_identity, + }) + .await + { + Ok(result) => respond_operation(stream, result).await, + Err(error) => respond_error(stream, 500, &error.to_string()).await, + } +} + +pub(super) async fn handle_clear(stream: &mut TcpStream, body: &str) -> anyhow::Result<()> { + if !ensure_loopback_control_caller(stream).await? { + return Ok(()); + } + let request = if body.trim().is_empty() { + KvCacheClearRequest::default() + } else { + match serde_json::from_str::(body) { + Ok(request) => request, + Err(error) => return respond_error(stream, 400, &error.to_string()).await, + } + }; + if node_kv_disk_manager().is_none() { + return respond_error(stream, 409, "disk prompt cache is not active").await; + } + let model_identity = match normalize_model_identity(request.model_identity) { + Ok(identity) => identity, + Err(error) => return respond_error(stream, 400, &error).await, + }; + match execute_operation(KvCacheOperation::Clear { model_identity }).await { + Ok(result) => respond_operation(stream, result).await, + Err(error) => respond_error(stream, 500, &error.to_string()).await, + } +} + +async fn respond_operation( + stream: &mut TcpStream, + result: KvCacheOperationResult, +) -> anyhow::Result<()> { + let status = serde_json::from_slice::(&result.status_json)?; + respond_json( + stream, + 200, + &serde_json::json!({ + "freed_bytes": result.freed_bytes.unwrap_or(0), + "status": status, + }), + ) + .await +} + +pub(crate) async fn execute_operation( + operation: KvCacheOperation, +) -> anyhow::Result { + let freed_bytes = match operation { + KvCacheOperation::Status => None, + KvCacheOperation::Prune { + target_bytes, + model_identity, + } => { + let manager = node_kv_disk_manager() + .ok_or_else(|| anyhow::anyhow!("disk prompt cache is not active"))?; + let target = target_bytes.unwrap_or_else(|| { + manager + .limits() + .budget_bytes + .saturating_mul(85) + .checked_div(100) + .unwrap_or(0) + }); + Some( + tokio::task::spawn_blocking(move || match model_identity.as_deref() { + Some(identity) if identity.trim().is_empty() => { + anyhow::bail!("model_identity must not be empty") + } + Some(identity) => manager.prune_model_to(identity, target), + None => manager.prune_to(target), + }) + .await??, + ) + } + KvCacheOperation::Clear { model_identity } => { + let manager = node_kv_disk_manager() + .ok_or_else(|| anyhow::anyhow!("disk prompt cache is not active"))?; + Some( + tokio::task::spawn_blocking(move || match model_identity.as_deref() { + Some(identity) if identity.trim().is_empty() => { + anyhow::bail!("model_identity must not be empty") + } + Some(identity) => manager.clear_model(identity), + None => manager.clear(), + }) + .await??, + ) + } + }; + Ok(KvCacheOperationResult { + status_json: serde_json::to_vec(&status_payload()?)?, + freed_bytes, + }) +} + +fn status_payload() -> anyhow::Result { + let cache = + node_kv_disk_cache().ok_or_else(|| anyhow::anyhow!("disk cache not initialized"))?; + let configured = cache.configured; + let (effective, usage, activity, reconciliation, inventory) = match cache.manager { + Some(manager) => { + let manager_status = manager.effective_status(); + ( + KvCacheEffectivePayload { + state: match manager_status.state { + L3EffectiveState::Active => "active", + L3EffectiveState::ReadOnlyLowSpace => "read_only_low_space", + L3EffectiveState::Degraded => "degraded", + } + .to_string(), + reason: manager_status.reason.map(|reason| { + match reason { + L3StateReason::ReadOnlyLowSpace => "read_only_low_space", + L3StateReason::InsufficientSpace => "insufficient_space", + L3StateReason::StorageError => "storage_error", + } + .to_string() + }), + manager: Some(manager_status), + }, + Some(manager.usage()?), + Some(manager.activity()?), + Some(manager.reconciliation()), + manager.inventory()?, + ) + } + None if !configured.enabled() => ( + KvCacheEffectivePayload { + state: "off".to_string(), + reason: None, + manager: None, + }, + None, + None, + None, + Vec::new(), + ), + None => ( + KvCacheEffectivePayload { + state: "degraded".to_string(), + reason: Some(if configured.budget_bytes == Some(0) { + "budget_below_entry_floor".to_string() + } else { + "storage_unavailable".to_string() + }), + manager: None, + }, + None, + None, + None, + Vec::new(), + ), + }; + Ok(KvCacheStatusPayload { + version: 1, + format_version: MANIFEST_VERSION, + configured: KvCacheConfiguredPayload { + mode: configured.mode, + directory: configured.directory, + budget_bytes: configured.budget_bytes, + minimum_free_bytes: configured.minimum_free_bytes, + sources: configured.sources, + }, + effective, + usage, + activity, + reconciliation, + inventory, + }) +} diff --git a/crates/mesh-llm-host-runtime/src/api/routes/mod.rs b/crates/mesh-llm-host-runtime/src/api/routes/mod.rs index ada23f6bad..6df103e683 100644 --- a/crates/mesh-llm-host-runtime/src/api/routes/mod.rs +++ b/crates/mesh-llm-host-runtime/src/api/routes/mod.rs @@ -3,6 +3,7 @@ mod control_apply_diagnostics; mod diagnostics; mod discover; mod health; +pub(crate) mod kv_cache; pub(crate) mod logs; mod mcp; mod mesh_hook; @@ -89,6 +90,7 @@ pub(super) const DISPATCH_REQUEST: DispatchRequestFn = | ("GET", "/api/runtime/stages") | ("GET", "/api/runtime/config-schema") | ("GET", "/api/runtime/config-control-state") + | ("GET", "/api/runtime/kv-cache") | ("GET", "/api/runtime/control-bootstrap") | ("GET", "/api/runtime/intents") | ("GET", "/api/runtime/activity") @@ -96,6 +98,7 @@ pub(super) const DISPATCH_REQUEST: DispatchRequestFn = | ("POST", "/api/runtime/control/scan-refresh") | ("POST", "/api/runtime/control/refresh-inventory") | ("POST", "/api/runtime/control/apply-config") + | ("POST", "/api/runtime/control/kv-cache") | ("POST", "/api/runtime/control/load-model") | ("POST", "/api/runtime/control/unload-model") | ("POST", "/api/runtime/control/ensure-model") @@ -103,6 +106,8 @@ pub(super) const DISPATCH_REQUEST: DispatchRequestFn = | ("POST", "/api/runtime/config/validate") | ("POST", "/api/runtime/mesh-guardrails") | ("POST", "/api/runtime/models") + | ("POST", "/api/runtime/kv-cache/prune") + | ("DELETE", "/api/runtime/kv-cache") | ("PUT", "/api/runtime/activity/override") | ("DELETE", "/api/runtime/activity/override") | ("GET", "/api/events") => { diff --git a/crates/mesh-llm-host-runtime/src/api/routes/runtime.rs b/crates/mesh-llm-host-runtime/src/api/routes/runtime.rs index 0aa4a28cb6..2a2db62cfc 100644 --- a/crates/mesh-llm-host-runtime/src/api/routes/runtime.rs +++ b/crates/mesh-llm-host-runtime/src/api/routes/runtime.rs @@ -13,6 +13,7 @@ use crate::crypto::{ OwnerKeychainLoadError, keystore_metadata, load_keystore, load_owner_keypair_from_keychain, }; use crate::plugin::validate_config_diagnostics_with_installed_plugin_schemas; +use futures_util::{FutureExt, StreamExt, stream}; use mesh_client::{ ClientBuilder, ControlPlaneBootstrapOptions, ControlPlaneClientError, ControlPlaneConnection, InviteToken, OwnerControlRemoteError, client::control_plane::OwnerControlScanRefreshResult, @@ -32,6 +33,10 @@ use tokio::io::AsyncWriteExt; use tokio::net::TcpStream; use zeroize::Zeroizing; +const CONTROL_KV_CACHE_MAX_ENDPOINTS: usize = 256; +const CONTROL_KV_CACHE_CONCURRENCY: usize = 8; +const CONTROL_KV_CACHE_BATCH_TIMEOUT_SECS: u64 = 45; + pub(super) async fn handle( stream: &mut TcpStream, state: &MeshApi, @@ -43,7 +48,7 @@ pub(super) async fn handle( "GET" => handle_get(stream, state, path_only).await, "POST" => handle_post(stream, state, path_only, body).await, "PUT" => handle_put(stream, state, path_only, body).await, - "DELETE" => handle_delete(stream, state, path_only).await, + "DELETE" => handle_delete(stream, state, path_only, body).await, _ => Ok(()), } } @@ -66,6 +71,7 @@ async fn handle_get( "/api/runtime/config-control-state" => { handle_runtime_config_control_state(stream, state).await } + "/api/runtime/kv-cache" => super::kv_cache::handle_status(stream).await, "/api/runtime/control-bootstrap" => handle_control_bootstrap(stream, state).await, "/api/runtime/intents" => handle_get_intents(stream, state).await, "/api/runtime/activity" => super::runtime_activity::handle_get(stream, state).await, @@ -91,6 +97,7 @@ async fn handle_post( "/api/runtime/control/apply-config" => { handle_control_apply_config(stream, state, body).await } + "/api/runtime/control/kv-cache" => handle_control_kv_cache(stream, state, body).await, "/api/runtime/control/load-model" => handle_control_load_model(stream, state, body).await, "/api/runtime/control/unload-model" => { handle_control_unload_model(stream, state, body).await @@ -102,6 +109,7 @@ async fn handle_post( "/api/runtime/config/validate" => handle_runtime_config_validate(stream, body).await, "/api/runtime/mesh-guardrails" => handle_set_mesh_guardrails(stream, state, body).await, "/api/runtime/models" => handle_load_model(stream, state, body).await, + "/api/runtime/kv-cache/prune" => super::kv_cache::handle_prune(stream, body).await, _ => Ok(()), } } @@ -124,11 +132,13 @@ async fn handle_delete( stream: &mut TcpStream, state: &MeshApi, path_only: &str, + body: &str, ) -> anyhow::Result<()> { match path_only { "/api/runtime/activity/override" => { super::runtime_activity::handle_delete(stream, state).await } + "/api/runtime/kv-cache" => super::kv_cache::handle_clear(stream, body).await, p if p.starts_with("/api/runtime/instances/") => { handle_unload_instance(stream, state, p).await } @@ -191,6 +201,14 @@ struct RawApplyConfigRequest { config: serde_json::Value, } +#[derive(Debug, Deserialize)] +struct ControlKvCacheRequest { + endpoints: Vec, + operation: String, + target_bytes: Option, + model_identity: Option, +} + #[derive(Debug, Deserialize)] struct ValidateConfigRequest { toml: String, @@ -313,6 +331,25 @@ struct LocalControlErrorPayload { current_revision: Option, } +#[derive(Debug, Serialize)] +struct LocalControlKvCachePayload { + results: Vec, +} + +#[derive(Debug, Serialize)] +struct LocalControlKvCacheResult { + endpoint: String, + #[serde(skip_serializing_if = "Option::is_none")] + target_node_id: Option, + operation: String, + #[serde(skip_serializing_if = "Option::is_none")] + freed_bytes: Option, + #[serde(skip_serializing_if = "Option::is_none")] + status: Option, + #[serde(skip_serializing_if = "Option::is_none")] + error: Option, +} + #[derive(Clone, Debug, Default, Serialize)] pub(crate) struct ConfigControlStatePayload { #[serde(default)] @@ -532,6 +569,150 @@ async fn handle_control_apply_config( } } +async fn handle_control_kv_cache( + stream: &mut TcpStream, + state: &MeshApi, + body: &str, +) -> anyhow::Result<()> { + if !ensure_loopback_control_caller(stream).await? { + return Ok(()); + } + let request: ControlKvCacheRequest = match serde_json::from_str(body) { + Ok(request) => request, + Err(_) => return respond_error(stream, 400, "Invalid JSON body").await, + }; + if request.endpoints.is_empty() { + return respond_error( + stream, + 400, + "at least one owner-control endpoint is required", + ) + .await; + } + if request.endpoints.len() > CONTROL_KV_CACHE_MAX_ENDPOINTS { + return respond_error( + stream, + 400, + "at most 256 owner-control endpoints are allowed", + ) + .await; + } + let operation = match request.operation.as_str() { + "status" => mesh_client::proto::node::OwnerControlKvCacheOperation::Status, + "prune" => mesh_client::proto::node::OwnerControlKvCacheOperation::Prune, + "clear" => mesh_client::proto::node::OwnerControlKvCacheOperation::Clear, + _ => return respond_error(stream, 400, "unknown kv-cache operation").await, + }; + let endpoints = request.endpoints; + let timeout_endpoints = endpoints.clone(); + let operation_label = request.operation.clone(); + let mut operations = stream::iter(endpoints.into_iter().enumerate()) + .map(|(index, endpoint)| { + execute_control_kv_cache_target( + state, + endpoint, + request.operation.clone(), + operation, + request.target_bytes, + request.model_identity.clone(), + ) + .map(move |result| (index, result)) + }) + .buffer_unordered(CONTROL_KV_CACHE_CONCURRENCY); + let deadline = tokio::time::Instant::now() + + std::time::Duration::from_secs(CONTROL_KV_CACHE_BATCH_TIMEOUT_SECS); + let mut indexed_results = Vec::with_capacity(timeout_endpoints.len()); + loop { + match tokio::time::timeout_at(deadline, operations.next()).await { + Ok(Some(result)) => indexed_results.push(result), + Ok(None) => break, + Err(_) => break, + } + } + drop(operations); + if indexed_results.len() < timeout_endpoints.len() { + let completed = indexed_results + .iter() + .map(|(index, _)| *index) + .collect::>(); + indexed_results.extend( + timeout_endpoints + .into_iter() + .enumerate() + .filter(|(index, _)| !completed.contains(index)) + .map(|(index, endpoint)| { + ( + index, + LocalControlKvCacheResult { + endpoint, + target_node_id: None, + operation: operation_label.clone(), + freed_bytes: None, + status: None, + error: Some(LocalControlErrorPayload { + code: "control_timeout".to_string(), + message: format!( + "kv-cache batch exceeded its {CONTROL_KV_CACHE_BATCH_TIMEOUT_SECS}-second deadline before a receipt was returned" + ), + legacy_retry_allowed: false, + current_revision: None, + }), + }, + ) + }), + ); + } + indexed_results.sort_by_key(|(index, _)| *index); + let results = indexed_results + .into_iter() + .map(|(_, result)| result) + .collect(); + respond_json(stream, 200, &LocalControlKvCachePayload { results }).await +} + +async fn execute_control_kv_cache_target( + state: &MeshApi, + endpoint: String, + operation_label: String, + operation: mesh_client::proto::node::OwnerControlKvCacheOperation, + target_bytes: Option, + model_identity: Option, +) -> LocalControlKvCacheResult { + let mut result = LocalControlKvCacheResult { + endpoint: endpoint.clone(), + target_node_id: None, + operation: operation_label, + freed_bytes: None, + status: None, + error: None, + }; + match connect_owner_control_client(state, &endpoint).await { + Ok(client) => { + result.target_node_id = Some(hex::encode(client.target_node_id())); + let response = client + .kv_cache(operation, target_bytes, model_identity) + .await; + client.close().await; + match response { + Ok(response) => { + result.freed_bytes = response.freed_bytes; + match serde_json::from_slice(&response.status_json) { + Ok(status) => result.status = Some(status), + Err(error) => { + result.error = Some(control_error_from_anyhow(anyhow::anyhow!( + "invalid kv-cache status from target: {error}" + ))); + } + } + } + Err(error) => result.error = Some(control_error_from_client(error)), + } + } + Err(error) => result.error = Some(error), + } + result +} + async fn handle_control_load_model( stream: &mut TcpStream, state: &MeshApi, @@ -769,7 +950,7 @@ fn required_control_endpoint(endpoint: Option) -> Result anyhow::Result { +pub(super) async fn ensure_loopback_control_caller(stream: &mut TcpStream) -> anyhow::Result { ensure_loopback_control_caller_for_peer_addr(stream, stream.peer_addr()).await } diff --git a/crates/mesh-llm-host-runtime/src/api/tests/support.rs b/crates/mesh-llm-host-runtime/src/api/tests/support.rs index 456dd7dcaa..4e772f2d9d 100644 --- a/crates/mesh-llm-host-runtime/src/api/tests/support.rs +++ b/crates/mesh-llm-host-runtime/src/api/tests/support.rs @@ -342,6 +342,7 @@ async fn spawn_owner_control_test_server() -> OwnerControlTestServer { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), error: None, }; @@ -426,6 +427,7 @@ async fn spawn_owner_control_apply_test_server( unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), error: None, }, diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/certification.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/certification.rs index 899fb4a4fe..46032f3219 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/certification.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/certification.rs @@ -579,7 +579,7 @@ mod tests { use crate::inference::skippy::materialization::{StagePackageInfo, StagePackageLayerInfo}; use serde_json::json; use tokio::io::{AsyncReadExt, AsyncWriteExt}; - use tokio::net::TcpListener; + use tokio::net::{TcpListener, TcpStream}; #[test] fn certification_ranges_split_two_stage_package() { @@ -835,6 +835,34 @@ mod tests { format!("http://{addr}") } + async fn read_complete_http_request(stream: &mut TcpStream) -> Vec { + let mut request = Vec::new(); + loop { + let mut chunk = [0u8; 4096]; + let n = stream.read(&mut chunk).await.unwrap(); + assert!(n > 0, "unexpected EOF while reading certification request"); + request.extend_from_slice(&chunk[..n]); + let Some(header_end) = request.windows(4).position(|window| window == b"\r\n\r\n") + else { + continue; + }; + let body_start = header_end + 4; + let headers = String::from_utf8_lossy(&request[..header_end]); + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + .unwrap_or(0); + if request.len() >= body_start + content_length { + return request; + } + } + } + /// Mimics a node that only recognizes `served_model_id` — anything else 404s, /// the same way a real host does when a client asks for a model name it /// doesn't advertise. @@ -844,9 +872,8 @@ mod tests { tokio::spawn(async move { for _ in 0..3 { let (mut stream, _) = listener.accept().await.unwrap(); - let mut buf = [0u8; 4096]; - let n = stream.read(&mut buf).await.unwrap(); - let request = String::from_utf8_lossy(&buf[..n]); + let request = read_complete_http_request(&mut stream).await; + let request = String::from_utf8_lossy(&request); let response = if request.starts_with("GET") { let body = json!({ "object": "list", @@ -900,4 +927,23 @@ mod tests { assert_eq!(gate.status, CertificationGateStatus::Passed, "{gate:?}"); } } + + #[tokio::test] + async fn certification_stub_reads_a_body_split_from_its_headers() { + let model = "hf://meshllm/split-request@abc123"; + let api_base = spawn_certification_stub_server(model.to_string()).await; + let addr = api_base.strip_prefix("http://").unwrap(); + let body = json!({"model": model, "messages": []}).to_string(); + let headers = format!( + "POST /v1/chat/completions HTTP/1.1\r\nHost: {addr}\r\nContent-Length: {}\r\n\r\n", + body.len() + ); + let mut stream = TcpStream::connect(addr).await.unwrap(); + stream.write_all(headers.as_bytes()).await.unwrap(); + tokio::task::yield_now().await; + stream.write_all(body.as_bytes()).await.unwrap(); + let mut response = Vec::new(); + stream.read_to_end(&mut response).await.unwrap(); + assert!(String::from_utf8_lossy(&response).starts_with("HTTP/1.1 200 OK")); + } } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs index 2b83af5df9..cc8733cc27 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs @@ -591,6 +591,7 @@ fn embedded_openai_args_from( linear_proposal_ingress: serving_hooks.linear_proposal_ingress(), kv_lifecycle_observer: serving_hooks.kv_lifecycle_observer(), openai_guardrails: None, + l3_manager: crate::runtime::kv_disk_config::node_kv_disk_manager(), }) } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs index 11e120a260..9193a99f52 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs @@ -384,7 +384,7 @@ tuning_profile = "throughput" assert_eq!(resolved.model_fit.kv_offload, "true"); assert_eq!(resolved.throughput.tuning_profile, "throughput"); assert_eq!(resolved.model_fit.batch, 1024); - assert_eq!(resolved.model_fit.ubatch, 256); + assert_eq!(resolved.model_fit.ubatch, 1024); assert_eq!(resolved.throughput.parallel, 2); assert_eq!(resolved.throughput.continuous_batching, "true"); assert_eq!(resolved.hardware.fit_target_mib, Some(10_752)); diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs index 16edcd6964..be415bf193 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs @@ -689,6 +689,7 @@ impl ResolvedEmbeddedOpenAiArgs { linear_proposal_ingress: None, kv_lifecycle_observer: None, openai_guardrails: None, + l3_manager: crate::runtime::kv_disk_config::node_kv_disk_manager(), } } } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs index 5b0d2d4098..f49d29e7e7 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs @@ -8,7 +8,13 @@ use crate::plugin::{MeshConfig, ReasoningBudget, ReasoningEnabled, RequestDefaul pub(super) const BUILTIN_CTX_SIZE: u32 = 4096; pub(super) const BUILTIN_BATCH: u32 = 512; -pub(super) const BUILTIN_UBATCH: u32 = 128; +/// Matches llama.cpp's own default (`LLAMA_SERVER_DEFAULT_N_UBATCH = 512`) and clears +/// the CUDA SSM SSD kernel gate (`n_tok > SSM_SSD_MIN_TOKENS`, 128, strict), which the +/// previous 128 default missed by exactly one token — forcing every recurrent (mamba) +/// prefill onto the sequential scan fallback. Measured on granite-4.0-h-1b: TTFT p50 +/// 0.670 → 0.415 s (C1) and 6.38 → 3.97 s (C8); decode 22.2 → 39.4 tok/s at C8. +/// See WHITE_UBATCH_512_FALSIFICATION_2026_09_08 in the 2026-09-08 competitive bench. +pub(super) const BUILTIN_UBATCH: u32 = 512; pub(super) const BUILTIN_PARALLEL: usize = 32; pub(super) const BUILTIN_PREFILL_CHUNK_SIZE: usize = 64; pub(super) const BUILTIN_PREFILL_ADAPTIVE_START: usize = 64; diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs index 81f09effab..d186d0fe2e 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs @@ -330,6 +330,7 @@ impl StageControlState { native_mtp_enabled: effective_load.native_mtp_enabled, continuous_batching: effective_load.continuous_batching, openai: None, + l3_manager: crate::runtime::kv_disk_config::node_kv_disk_manager(), }); self.stages.insert( key.clone(), diff --git a/crates/mesh-llm-host-runtime/src/mesh/owner_control/commands/kv_cache.rs b/crates/mesh-llm-host-runtime/src/mesh/owner_control/commands/kv_cache.rs new file mode 100644 index 0000000000..661b9aac19 --- /dev/null +++ b/crates/mesh-llm-host-runtime/src/mesh/owner_control/commands/kv_cache.rs @@ -0,0 +1,122 @@ +use crate::api::routes::kv_cache::{KvCacheOperation, execute_operation}; +use crate::mesh::owner_control_error_envelope; +use crate::proto::node::{ + OwnerControlEnvelope, OwnerControlErrorCode, OwnerControlKvCacheOperation, + OwnerControlKvCacheRequest, OwnerControlKvCacheResponse, OwnerControlResponse, +}; +use crate::protocol::NODE_PROTOCOL_GENERATION; + +pub(crate) async fn execute( + request_id: u64, + request: OwnerControlKvCacheRequest, +) -> OwnerControlEnvelope { + let operation = match decode_operation(&request) { + Ok(operation) => operation, + Err(message) => { + return owner_control_error_envelope( + OwnerControlErrorCode::BadRequest, + Some(request_id), + None, + message, + ); + } + }; + + match execute_operation(operation).await { + Ok(result) => OwnerControlEnvelope { + r#gen: NODE_PROTOCOL_GENERATION, + handshake: None, + request: None, + response: Some(OwnerControlResponse { + request_id, + kv_cache: Some(OwnerControlKvCacheResponse { + status_json: result.status_json, + freed_bytes: result.freed_bytes, + }), + ..Default::default() + }), + error: None, + }, + Err(error) => owner_control_error_envelope( + OwnerControlErrorCode::ControlUnavailable, + Some(request_id), + None, + error.to_string(), + ), + } +} + +fn decode_operation(request: &OwnerControlKvCacheRequest) -> Result { + let model_identity = request + .model_identity + .as_ref() + .map(|identity| identity.trim().to_string()) + .transpose_nonempty("model_identity")?; + match OwnerControlKvCacheOperation::try_from(request.operation).ok() { + Some(OwnerControlKvCacheOperation::Status) => { + if request.target_bytes.is_some() || model_identity.is_some() { + return Err("kv-cache status does not accept mutation parameters".to_string()); + } + Ok(KvCacheOperation::Status) + } + Some(OwnerControlKvCacheOperation::Prune) => Ok(KvCacheOperation::Prune { + target_bytes: request.target_bytes, + model_identity, + }), + Some(OwnerControlKvCacheOperation::Clear) => { + if request.target_bytes.is_some() { + return Err("kv-cache clear does not accept target_bytes".to_string()); + } + Ok(KvCacheOperation::Clear { model_identity }) + } + Some(OwnerControlKvCacheOperation::Unspecified) | None => { + Err("kv-cache operation is unknown or unspecified".to_string()) + } + } +} + +trait NonemptyStringOption { + fn transpose_nonempty(self, field: &str) -> Result, String>; +} + +impl NonemptyStringOption for Option { + fn transpose_nonempty(self, field: &str) -> Result, String> { + match self { + Some(value) if value.is_empty() => Err(format!("{field} must not be empty")), + value => Ok(value), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn request(operation: OwnerControlKvCacheOperation) -> OwnerControlKvCacheRequest { + OwnerControlKvCacheRequest { + requester_node_id: vec![1; 32], + target_node_id: vec![2; 32], + operation: operation as i32, + target_bytes: None, + model_identity: None, + } + } + + #[test] + fn rejects_parameters_that_do_not_apply_to_operation() { + let mut status = request(OwnerControlKvCacheOperation::Status); + status.target_bytes = Some(1); + assert!(decode_operation(&status).is_err()); + + let mut clear = request(OwnerControlKvCacheOperation::Clear); + clear.target_bytes = Some(1); + assert!(decode_operation(&clear).is_err()); + } + + #[test] + fn exact_model_identity_must_not_be_blank() { + let mut prune = request(OwnerControlKvCacheOperation::Prune); + prune.model_identity = Some(" ".to_string()); + assert!(decode_operation(&prune).is_err()); + } +} diff --git a/crates/mesh-llm-host-runtime/src/mesh/owner_control/commands/mod.rs b/crates/mesh-llm-host-runtime/src/mesh/owner_control/commands/mod.rs index 6e7d051d8e..626bccbb61 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/owner_control/commands/mod.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/owner_control/commands/mod.rs @@ -1,9 +1,12 @@ +pub(crate) mod kv_cache; pub(crate) mod model_lifecycle; pub(crate) mod scan_refresh; +const OWNER_CONTROL_KV_CACHE_DEADLINE_SECS: u64 = 30; + use crate::proto::node::{ OwnerControlApplyConfigRequest, OwnerControlDrainModelRequest, OwnerControlEnsureModelRequest, - OwnerControlGetConfigRequest, OwnerControlLoadModelRequest, + OwnerControlGetConfigRequest, OwnerControlKvCacheRequest, OwnerControlLoadModelRequest, OwnerControlRefreshInventoryRequest, OwnerControlRequest, OwnerControlUnloadModelRequest, OwnerControlWatchConfigRequest, }; @@ -93,6 +96,10 @@ pub(crate) enum OwnedNodeCommand { request_id: u64, request: OwnerControlDrainModelRequest, }, + KvCache { + request_id: u64, + request: OwnerControlKvCacheRequest, + }, } impl OwnedNodeCommand { @@ -140,7 +147,13 @@ impl OwnedNodeCommand { request, }); } - request.drain_model.map(|request| Self::DrainModel { + if let Some(request) = request.drain_model { + return Some(Self::DrainModel { + request_id, + request, + }); + } + request.kv_cache.map(|request| Self::KvCache { request_id, request, }) @@ -155,7 +168,8 @@ impl OwnedNodeCommand { | Self::LoadModel { request_id, .. } | Self::UnloadModel { request_id, .. } | Self::EnsureModel { request_id, .. } - | Self::DrainModel { request_id, .. } => *request_id, + | Self::DrainModel { request_id, .. } + | Self::KvCache { request_id, .. } => *request_id, } } @@ -169,6 +183,7 @@ impl OwnedNodeCommand { Self::UnloadModel { request, .. } => &request.requester_node_id, Self::EnsureModel { request, .. } => &request.requester_node_id, Self::DrainModel { request, .. } => &request.requester_node_id, + Self::KvCache { request, .. } => &request.requester_node_id, } } @@ -182,6 +197,7 @@ impl OwnedNodeCommand { Self::UnloadModel { request, .. } => &request.target_node_id, Self::EnsureModel { request, .. } => &request.target_node_id, Self::DrainModel { request, .. } => &request.target_node_id, + Self::KvCache { request, .. } => &request.target_node_id, } } @@ -194,7 +210,8 @@ impl OwnedNodeCommand { | Self::LoadModel { .. } | Self::UnloadModel { .. } | Self::EnsureModel { .. } - | Self::DrainModel { .. } => OwnedNodeCommandExecutionShape::Unary, + | Self::DrainModel { .. } + | Self::KvCache { .. } => OwnedNodeCommandExecutionShape::Unary, } } @@ -206,6 +223,9 @@ impl OwnedNodeCommand { | Self::UnloadModel { .. } | Self::EnsureModel { .. } | Self::DrainModel { .. } => OwnedNodeCommandDeadline::Unary(Duration::from_secs(5)), + Self::KvCache { .. } => OwnedNodeCommandDeadline::Unary(Duration::from_secs( + OWNER_CONTROL_KV_CACHE_DEADLINE_SECS, + )), Self::ScanRefresh { .. } => OwnedNodeCommandDeadline::Scan(Duration::from_secs( OWNER_CONTROL_SCAN_DEADLINE_SECS, )), @@ -245,6 +265,7 @@ mod tests { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }) .expect("typed command"); @@ -261,6 +282,37 @@ mod tests { ); } + #[test] + fn kv_cache_command_is_authenticated_unary_work() { + let command = OwnedNodeCommand::decode(OwnerControlRequest { + request_id: 42, + kv_cache: Some(crate::proto::node::OwnerControlKvCacheRequest { + requester_node_id: vec![1; 32], + target_node_id: vec![2; 32], + operation: crate::proto::node::OwnerControlKvCacheOperation::Status as i32, + target_bytes: None, + model_identity: None, + }), + ..Default::default() + }) + .expect("typed command"); + + assert_eq!(command.request_id(), 42); + assert_eq!(command.requester_node_id(), [1; 32]); + assert_eq!(command.target_node_id(), [2; 32]); + assert_eq!( + command.execution_shape(), + OwnedNodeCommandExecutionShape::Unary + ); + assert_eq!( + command.deadline(), + OwnedNodeCommandDeadline::Unary(Duration::from_secs( + OWNER_CONTROL_KV_CACHE_DEADLINE_SECS + )) + ); + assert!(!command.is_model_lifecycle()); + } + #[tokio::test(start_paused = true)] async fn slow_scan_within_deadline_completes() { let deadline = OwnedNodeCommandDeadline::Scan(Duration::from_secs(30)); @@ -329,6 +381,7 @@ mod tests { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), error: None, }; @@ -368,6 +421,7 @@ mod tests { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, } } diff --git a/crates/mesh-llm-host-runtime/src/mesh/owner_control/commands/model_lifecycle.rs b/crates/mesh-llm-host-runtime/src/mesh/owner_control/commands/model_lifecycle.rs index cd49c0a0b1..47fe2ff4fa 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/owner_control/commands/model_lifecycle.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/owner_control/commands/model_lifecycle.rs @@ -276,6 +276,7 @@ fn success_lifecycle_envelope( unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }; match operation { LifecycleOperation::Load => { diff --git a/crates/mesh-llm-host-runtime/src/mesh/owner_control/commands/scan_refresh.rs b/crates/mesh-llm-host-runtime/src/mesh/owner_control/commands/scan_refresh.rs index f167450427..7fba72da2b 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/owner_control/commands/scan_refresh.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/owner_control/commands/scan_refresh.rs @@ -54,6 +54,7 @@ async fn success_envelope( unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), error: None, } diff --git a/crates/mesh-llm-host-runtime/src/mesh/owner_control/mod.rs b/crates/mesh-llm-host-runtime/src/mesh/owner_control/mod.rs index abe684fec4..cdc562d652 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/owner_control/mod.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/owner_control/mod.rs @@ -626,6 +626,7 @@ impl Node { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), error: None, } @@ -672,6 +673,7 @@ impl Node { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), error: None, }, @@ -1070,6 +1072,13 @@ impl Node { } self.send_owner_control_envelope(send, envelope).await?; } + OwnedNodeCommand::KvCache { + request_id, + request, + } => { + let envelope = commands::kv_cache::execute(request_id, request).await; + self.send_owner_control_envelope(send, envelope).await?; + } } anyhow::Ok(()) }; diff --git a/crates/mesh-llm-host-runtime/src/mesh/owner_control_response.rs b/crates/mesh-llm-host-runtime/src/mesh/owner_control_response.rs index a3e6e0c453..b5d9033e9d 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/owner_control_response.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/owner_control_response.rs @@ -21,6 +21,7 @@ pub(super) fn apply_response_envelope( unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), error: None, } diff --git a/crates/mesh-llm-host-runtime/src/mesh/owner_lifecycle_cache/tests.rs b/crates/mesh-llm-host-runtime/src/mesh/owner_lifecycle_cache/tests.rs index 6e16bfaea3..12f8ebdcc1 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/owner_lifecycle_cache/tests.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/owner_lifecycle_cache/tests.rs @@ -27,6 +27,7 @@ fn lifecycle_envelope(request_id: u64, marker: &str) -> crate::proto::node::Owne unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), error: None, } diff --git a/crates/mesh-llm-host-runtime/src/mesh/tests/control_plane_unique.rs b/crates/mesh-llm-host-runtime/src/mesh/tests/control_plane_unique.rs index cc8214295a..770fba27d7 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/tests/control_plane_unique.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/tests/control_plane_unique.rs @@ -52,6 +52,7 @@ async fn control_plane_legacy_compat_new_client_prefers_control_alpn() -> Result unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), response: None, error: None, @@ -146,6 +147,7 @@ async fn control_plane_validation_error_preserves_request_id() -> Result<()> { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), response: None, error: None, diff --git a/crates/mesh-llm-host-runtime/src/models/profile.rs b/crates/mesh-llm-host-runtime/src/models/profile.rs index 6ed3c90783..c379988670 100644 --- a/crates/mesh-llm-host-runtime/src/models/profile.rs +++ b/crates/mesh-llm-host-runtime/src/models/profile.rs @@ -90,6 +90,7 @@ fn resolve_parameter_size( parameter_count: Option, ) -> Option { source_size + .and_then(non_empty) .or_else(|| parameter_count.and_then(parameter_size_from_count)) .or_else(|| parameter_size_from_text(model_name)) } @@ -190,6 +191,13 @@ mod tests { resolve_parameter_size("model-7B", None, None).as_deref(), Some("7B") ); + for blank in ["", " ", "\t\n"] { + assert_eq!( + resolve_parameter_size("model-7B", Some(blank.to_string()), Some(8_000_000_000),) + .as_deref(), + Some("8B") + ); + } } #[test] diff --git a/crates/mesh-llm-host-runtime/src/network/openai/moa_gateway/self_fill.rs b/crates/mesh-llm-host-runtime/src/network/openai/moa_gateway/self_fill.rs index f94c01332c..2f7ce3a467 100644 --- a/crates/mesh-llm-host-runtime/src/network/openai/moa_gateway/self_fill.rs +++ b/crates/mesh-llm-host-runtime/src/network/openai/moa_gateway/self_fill.rs @@ -1,9 +1,10 @@ //! Same-model committees: distinct physical clones, reserved for the turn. +use super::pool::canonical_base_name; use super::workers::{LocalModelBackend, RemoteModelBackend, ReservedModelBackend}; use crate::inference::election::{InferenceTarget, ModelTargets}; use crate::mesh; use crate::network::affinity::AffinityRouter; -use crate::network::openai::routing_rank::rank_targets_by_context; +use crate::network::openai::routing_rank::rank_aliased_targets_by_context; use crate::network::reservations::RoutingReservation; use mesh_mixture_of_agents as moa; use std::sync::Arc; @@ -11,6 +12,7 @@ use std::sync::Arc; /// Measured self-MoA width; fleet capacity must not increase fan-out cost. const SELF_FILL_TARGET_WORKERS: usize = 2; +#[cfg(test)] async fn select_clones( node: &mesh::Node, name: &str, @@ -18,6 +20,29 @@ async fn select_clones( candidates: Vec, affinity: Option<&AffinityRouter>, ) -> Vec<(InferenceTarget, Option)> { + select_aliased_clones( + node, + &canonical_base_name(name), + required_tokens, + candidates + .into_iter() + .map(|target| (target, name.to_string())) + .collect(), + affinity, + ) + .await + .into_iter() + .map(|(target, _, reservation)| (target, reservation)) + .collect() +} + +async fn select_aliased_clones( + node: &mesh::Node, + reservation_key: &str, + required_tokens: Option, + candidates: Vec<(InferenceTarget, String)>, + affinity: Option<&AffinityRouter>, +) -> Vec<(InferenceTarget, String, Option)> { use crate::proto::node::InferenceAdmissionState; let deprioritized: std::collections::HashSet<_> = node @@ -32,36 +57,48 @@ async fn select_clones( // Preserve admission priority before context/throughput ranking. Local and // legacy peers stay healthy; hosts_for_model already excludes paused peers. let (mut healthy, mut spillover): (Vec<_>, Vec<_>) = candidates.into_iter().partition( - |target| !matches!(target, InferenceTarget::Remote(id) if deprioritized.contains(id)), + |(target, _)| !matches!(target, InferenceTarget::Remote(id) if deprioritized.contains(id)), ); let mut selected = Vec::with_capacity(SELF_FILL_TARGET_WORKERS); while selected.len() < SELF_FILL_TARGET_WORKERS { // Exhaust context-eligible healthy endpoints before considering spillover, // even when every healthy clone already has reservations from other turns. - let mut ranked = rank_targets_by_context(node, name, required_tokens, &healthy).await; + let mut ranked = rank_aliased_targets_by_context(node, required_tokens, &healthy).await; if ranked.ordered.is_empty() { - ranked = rank_targets_by_context(node, name, required_tokens, &spillover).await; + ranked = rank_aliased_targets_by_context(node, required_tokens, &spillover).await; } let Some(preferred) = ranked.ordered.first() else { break; }; + let physical = ranked + .ordered + .iter() + .map(|(target, _)| target.clone()) + .collect::>(); + let preferred_target = &preferred.0; let (target, reservation) = affinity .and_then(|router| { router.reserve_route( - name, - &ranked.ordered, + reservation_key, + &physical, ranked.equivalent_prefix, - preferred, + preferred_target, false, ) }) .map(|(target, guard)| (target, Some(guard))) - .unwrap_or_else(|| (preferred.clone(), None)); + .unwrap_or_else(|| (preferred_target.clone(), None)); + let alias = ranked + .ordered + .iter() + .find(|(candidate, _)| candidate == &target) + .map(|(_, alias)| alias.clone()) + .expect("reserved aliased target came from ranked candidates"); // Selection+reservation is atomic per slot; removing the endpoint // prevents duplicate workers even when other turns interleave slots. - healthy.retain(|candidate| candidate != &target); - spillover.retain(|candidate| candidate != &target); - selected.push((target, reservation)); + healthy.retain(|(candidate, _)| candidate != &target); + spillover.retain(|(candidate, _)| candidate != &target); + selected.push((target, alias, reservation)); } selected } @@ -78,61 +115,85 @@ pub(super) async fn self_fill_from_extra_instances( let Some(existing) = models.first().cloned() else { return; }; - let name = &existing.name; + let base = canonical_base_name(&existing.name); + let mut aliases = node.models_being_served().await; + if let Some(targets) = targets { + aliases.extend(targets.targets.keys().cloned()); + } + aliases.retain(|alias| canonical_base_name(alias) == base); + aliases.push(existing.name.clone()); + aliases.sort_by(|a, b| { + (b == &existing.name) + .cmp(&(a == &existing.name)) + .then_with(|| a.len().cmp(&b.len())) + .then_with(|| a.cmp(b)) + }); + aliases.dedup(); let mut candidates = Vec::new(); - if let Some(local) = targets - .and_then(|targets| targets.targets.get(name)) - .and_then(|targets| { - targets - .iter() - .find(|t| matches!(t, InferenceTarget::Local(_))) - }) - { - candidates.push(local.clone()); + for alias in &aliases { + if let Some(local) = targets + .and_then(|targets| targets.targets.get(alias)) + .and_then(|targets| { + targets + .iter() + .find(|t| matches!(t, InferenceTarget::Local(_))) + }) + { + candidates.push((local.clone(), alias.clone())); + } + candidates.extend( + node.hosts_for_model(alias) + .await + .into_iter() + .map(|peer_id| (InferenceTarget::Remote(peer_id), alias.clone())), + ); } - candidates.extend( - node.hosts_for_model(name) - .await - .into_iter() - .map(InferenceTarget::Remote), - ); + let mut physical = Vec::new(); + candidates.retain(|(target, _)| { + if physical.contains(target) { + false + } else { + physical.push(target.clone()); + true + } + }); if candidates.len() < 2 { return; } - let selected = select_clones(node, name, required_tokens, candidates, affinity).await; + let selected = select_aliased_clones(node, &base, required_tokens, candidates, affinity).await; if selected.len() < 2 { return; // Context filtering must not fabricate a second worker. } - *backends = selected - .into_iter() - .map(|(target, reservation)| { - let inner: Arc = match target { - InferenceTarget::Local(port) => Arc::new(LocalModelBackend { - port, - http: http.clone(), - }), - // No failover onto a sibling slot: every worker is a distinct sample. - InferenceTarget::Remote(peer_id) => Arc::new(RemoteModelBackend { - node: node.clone(), - peer_ids: vec![peer_id], - }), - InferenceTarget::None => unreachable!("self-fill only collects physical endpoints"), - }; - match reservation { - Some(reservation) => Arc::new(ReservedModelBackend { - inner, - _reservation: reservation, - }) as Arc, - None => inner, - } - }) - .collect(); - *models = (0..backends.len()) - .map(|backend_index| moa::ModelEntry { + let mut filled_backends = Vec::with_capacity(selected.len()); + let mut filled_models = Vec::with_capacity(selected.len()); + for (backend_index, (target, name, reservation)) in selected.into_iter().enumerate() { + let inner: Arc = match target { + InferenceTarget::Local(port) => Arc::new(LocalModelBackend { + port, + http: http.clone(), + }), + // No failover onto a sibling slot: every worker is a distinct sample. + InferenceTarget::Remote(peer_id) => Arc::new(RemoteModelBackend { + node: node.clone(), + peer_ids: vec![peer_id], + }), + InferenceTarget::None => unreachable!("self-fill only collects physical endpoints"), + }; + filled_backends.push(match reservation { + Some(reservation) => Arc::new(ReservedModelBackend { + inner, + _reservation: reservation, + }) as Arc, + None => inner, + }); + filled_models.push(moa::ModelEntry { + name, backend_index, ..existing.clone() - }) - .collect(); + }); + } + *backends = filled_backends; + *models = filled_models; } #[cfg(test)] diff --git a/crates/mesh-llm-host-runtime/src/network/openai/moa_gateway/self_fill/tests.rs b/crates/mesh-llm-host-runtime/src/network/openai/moa_gateway/self_fill/tests.rs index 01203d60f8..6bf4b76e37 100644 --- a/crates/mesh-llm-host-runtime/src/network/openai/moa_gateway/self_fill/tests.rs +++ b/crates/mesh-llm-host-runtime/src/network/openai/moa_gateway/self_fill/tests.rs @@ -17,6 +17,54 @@ async fn fleet(count: u32) -> (mesh::Node, Vec) { (node, candidates) } +fn rename_peer_model(peer: &mut mesh::PeerInfo, name: &str) { + peer.models = vec![name.to_string()]; + peer.serving_models = vec![name.to_string()]; + peer.hosted_models = vec![name.to_string()]; + for descriptor in &mut peer.served_model_descriptors { + descriptor.identity.model_name = name.to_string(); + } + for runtime in &mut peer.served_model_runtime { + runtime.model_name = name.to_string(); + } +} + +#[tokio::test] +async fn self_fill_preserves_each_physical_workers_routable_alias() { + let node = mesh::Node::new_for_tests(mesh::NodeRole::Client) + .await + .unwrap(); + let model = BIG_MODELS[1]; + let short = model.name; + let long = "unsloth/Qwen3-32B-GGUF:Q4_K_M"; + assert_eq!( + super::super::pool::canonical_base_name(short), + super::super::pool::canonical_base_name(long) + ); + let first = fleet_peer_with_health(1, model, None, Some(100_000)); + let mut second = fleet_peer_with_health(2, model, None, Some(100_000)); + rename_peer_model(&mut second, long); + let expected_aliases = first + .http_routable_models() + .into_iter() + .chain(second.http_routable_models()) + .collect::>(); + node.insert_test_peer(first).await; + node.insert_test_peer(second).await; + + let (backends, models) = + assemble_worker_pool(&node, None, Some(13_000), &reqwest::Client::new(), None).await; + + assert_eq!(backends.len(), 2); + assert_eq!( + models + .iter() + .map(|model| model.name.clone()) + .collect::>(), + expected_aliases + ); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn concurrent_committees_spread_across_twenty_clones() { let (node, candidates) = fleet(20).await; diff --git a/crates/mesh-llm-host-runtime/src/network/openai/routing_rank.rs b/crates/mesh-llm-host-runtime/src/network/openai/routing_rank.rs index 952cf0603d..93c3b00cfe 100644 --- a/crates/mesh-llm-host-runtime/src/network/openai/routing_rank.rs +++ b/crates/mesh-llm-host-runtime/src/network/openai/routing_rank.rs @@ -309,6 +309,31 @@ pub(super) async fn rank_targets_by_context( rank_candidates_by_context_and_throughput(&candidates, required_tokens) } +pub(super) async fn rank_aliased_targets_by_context( + node: &mesh::Node, + required_tokens: Option, + targets: &[(election::InferenceTarget, String)], +) -> RankedCandidates<(election::InferenceTarget, String)> { + let mut candidates = Vec::with_capacity(targets.len()); + for (target, model) in targets { + let context_length = match target { + election::InferenceTarget::Local(_) => node.local_model_context_length(model).await, + election::InferenceTarget::Remote(peer_id) => { + node.peer_model_context_length(*peer_id, model).await + } + election::InferenceTarget::None => None, + }; + let throughput = match target { + election::InferenceTarget::Remote(peer_id) => { + remote_target_throughput_rank(node, model, *peer_id).await + } + _ => local_target_throughput_rank(node, model, target), + }; + candidates.push(((target.clone(), model.clone()), context_length, throughput)); + } + rank_candidates_by_context_and_throughput(&candidates, required_tokens) +} + #[cfg(test)] pub(super) async fn order_targets_by_context( node: &mesh::Node, diff --git a/crates/mesh-llm-host-runtime/src/runtime/config_state.rs b/crates/mesh-llm-host-runtime/src/runtime/config_state.rs index 631c79aac6..651c4e06c4 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/config_state.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/config_state.rs @@ -1,6 +1,6 @@ use anyhow::Result; use mesh_llm_config::{ - ConfigDiagnostic, ConfigDiagnosticSeverity, ConfigPath, LoggingConfig, + ConfigDiagnostic, ConfigDiagnosticSeverity, ConfigPath, KvDiskTierConfig, LoggingConfig, built_in_config_schema_descriptor, legacy_validation_error_text, }; @@ -85,6 +85,8 @@ pub(crate) struct PendingConfigApply { logging_requires_restart: bool, dynamic_logging_only: bool, old_logging: LoggingConfig, + kv_disk_requires_restart: bool, + dynamic_kv_disk_changed: bool, } impl PendingConfigApply { @@ -120,8 +122,21 @@ impl PendingConfigApply { let _ = crate::apply_live_logging_limits(&self.old_logging); } + pub(crate) fn apply_live_kv_disk(&self) -> Option { + if !self.dynamic_kv_disk_changed { + return None; + } + match super::kv_disk_config::apply_live_kv_disk_limits(&self.config) { + Ok(rollback) => Some(rollback), + Err(error) => { + tracing::warn!(%error, "disk prompt-cache runtime unavailable; retaining limits as staged configuration"); + None + } + } + } + fn applied_result(&self, apply_mode: ConfigApplyMode) -> ApplyResult { - if self.logging_requires_restart { + if self.logging_requires_restart || self.kv_disk_requires_restart { ApplyResult::AppliedWithRestartRequired { revision: self.revision, hash: self.hash, @@ -309,6 +324,14 @@ fn logging_dynamic_limits_changed(old: &LoggingConfig, new: &LoggingConfig) -> b old.retention_ttl_secs != new.retention_ttl_secs || old.replay_capacity != new.replay_capacity } +fn kv_disk_changes_require_restart(old: &KvDiskTierConfig, new: &KvDiskTierConfig) -> bool { + old.mode != new.mode || old.directory != new.directory +} + +fn kv_disk_dynamic_limits_changed(old: &KvDiskTierConfig, new: &KvDiskTierConfig) -> bool { + old.budget_mib != new.budget_mib || old.minimum_free_mib != new.minimum_free_mib +} + impl Default for ConfigState { fn default() -> Self { let config = crate::plugin::MeshConfig::default(); @@ -424,6 +447,10 @@ impl ConfigState { let dynamic_logging_only = logging_dynamic_limits_changed(&old_logging, &new_config.logging) && !logging_requires_restart; + let old_kv_disk = &self.config.runtime.kv_cache.disk; + let new_kv_disk = &new_config.runtime.kv_cache.disk; + let kv_disk_requires_restart = kv_disk_changes_require_restart(old_kv_disk, new_kv_disk); + let dynamic_kv_disk_changed = kv_disk_dynamic_limits_changed(old_kv_disk, new_kv_disk); ConfigApplyPreparation::Pending(Box::new(PendingConfigApply { config: new_config, @@ -435,6 +462,8 @@ impl ConfigState { logging_requires_restart, dynamic_logging_only, old_logging, + kv_disk_requires_restart, + dynamic_kv_disk_changed, })) } @@ -515,6 +544,8 @@ where ConfigApplyPreparation::Pending(pending) => *pending, }; let live_logging_applied = pending.apply_live_logging(); + let live_kv_disk_rollback = pending.apply_live_kv_disk(); + let live_kv_disk_applied = live_kv_disk_rollback.is_some(); let persistence = persist(&pending); if live_logging_applied && matches!(&persistence, ConfigPersistence::PersistError(_)) { // A persistence failure after a live apply is rare, but restore the @@ -522,6 +553,11 @@ where // configuration and service do not diverge. pending.restore_live_logging(); } + if matches!(&persistence, ConfigPersistence::PersistError(_)) + && let Some(rollback) = live_kv_disk_rollback + { + super::kv_disk_config::restore_live_kv_disk_limits(rollback); + } let result = finish(pending, persistence); match result { @@ -530,7 +566,7 @@ where hash, diagnostics, .. - } if live_logging_applied => ApplyResult::Applied { + } if live_logging_applied || live_kv_disk_applied => ApplyResult::Applied { revision, hash, apply_mode: ConfigApplyMode::Live, diff --git a/crates/mesh-llm-host-runtime/src/runtime/config_state_tests.rs b/crates/mesh-llm-host-runtime/src/runtime/config_state_tests.rs index d1c5019e2c..b22a45c0e5 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/config_state_tests.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/config_state_tests.rs @@ -15,6 +15,8 @@ mod support; use support::*; #[path = "config_state_tests/diagnostics.rs"] mod diagnostics; +#[path = "config_state_tests/kv_disk.rs"] +mod kv_disk; #[path = "config_state_tests/logging.rs"] mod logging; #[path = "config_state_tests/persistence.rs"] diff --git a/crates/mesh-llm-host-runtime/src/runtime/config_state_tests/kv_disk.rs b/crates/mesh-llm-host-runtime/src/runtime/config_state_tests/kv_disk.rs new file mode 100644 index 0000000000..fbd43a50af --- /dev/null +++ b/crates/mesh-llm-host-runtime/src/runtime/config_state_tests/kv_disk.rs @@ -0,0 +1,48 @@ +use super::*; +use mesh_llm_config::{KvDiskTierConfig, KvDiskTierMode}; +use std::path::PathBuf; + +fn disk_config() -> KvDiskTierConfig { + KvDiskTierConfig { + mode: Some(KvDiskTierMode::Fixed), + directory: Some(PathBuf::from("/var/lib/mesh-llm/kv-cache")), + budget_mib: Some(32 * 1024), + minimum_free_mib: Some(16 * 1024), + } +} + +#[test] +fn disk_mode_and_directory_changes_require_restart() { + let old = disk_config(); + + let mut mode = old.clone(); + mode.mode = Some(KvDiskTierMode::Auto); + assert!(kv_disk_changes_require_restart(&old, &mode)); + + let mut directory = old.clone(); + directory.directory = Some(PathBuf::from("/var/lib/mesh-llm/other-cache")); + assert!(kv_disk_changes_require_restart(&old, &directory)); +} + +#[test] +fn disk_budget_and_reserve_changes_are_dynamic() { + let old = disk_config(); + + let mut limits = old.clone(); + limits.budget_mib = Some(48 * 1024); + limits.minimum_free_mib = Some(20 * 1024); + + assert!(!kv_disk_changes_require_restart(&old, &limits)); + assert!(kv_disk_dynamic_limits_changed(&old, &limits)); +} + +#[test] +fn mixed_static_and_dynamic_disk_change_preserves_both_classifications() { + let old = disk_config(); + let mut new = old.clone(); + new.directory = Some(PathBuf::from("/var/lib/mesh-llm/other-cache")); + new.budget_mib = Some(48 * 1024); + + assert!(kv_disk_changes_require_restart(&old, &new)); + assert!(kv_disk_dynamic_limits_changed(&old, &new)); +} diff --git a/crates/mesh-llm-host-runtime/src/runtime/context_planning.rs b/crates/mesh-llm-host-runtime/src/runtime/context_planning.rs index b8ce7535be..d4353e0e25 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/context_planning.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/context_planning.rs @@ -99,6 +99,59 @@ impl RuntimeResourcePlanningProfile { pub(super) struct RuntimeResourcePlan { pub(super) context_length: u32, pub(super) slots: usize, + /// Structured breakdown of the inputs and intermediate results the planner + /// used, emitted once per model start so every later memory-planning change + /// is measurable in production logs. `None` only for plans built outside + /// [`plan_runtime_resources`]. + pub(super) breakdown: Option, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) enum RuntimeResourcePlanSource { + ExplicitOverride, + StaticEstimate, + MeasuredFootprint, +} + +impl RuntimeResourcePlanSource { + pub(super) const fn as_str(self) -> &'static str { + match self { + Self::ExplicitOverride => "explicit_override", + Self::StaticEstimate => "static_estimate", + Self::MeasuredFootprint => "measured_footprint", + } + } +} + +/// The accounting selected at plan time. Static plans carry metadata-derived +/// estimates; measured plans carry the prior compatible load's native KV rate +/// and lane-scaled compute charge. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) struct RuntimeResourcePlanBreakdown { + pub(super) vram_bytes: u64, + pub(super) model_bytes: u64, + /// KV budget selected by the active planner. Static planning applies the + /// 85% utilization tax; measured planning applies its utilization target + /// and then subtracts the lane-scaled measured compute charge. + pub(super) kv_budget_bytes: u64, + /// Planned KV allocation: context_length x kv_bytes_per_token. + pub(super) planned_kv_bytes: u64, + /// Per-token KV cost used by the plan (layer-fraction scaled). + pub(super) kv_bytes_per_token: u64, + /// Compute charge held outside `kv_budget_bytes` by the selected planner. + pub(super) compute_charge_bytes: u64, + pub(super) planning_source: RuntimeResourcePlanSource, + /// `Some(false)` means a valid measured footprint proved that even the + /// minimum context cannot fit. Static and explicit plans use `None`. + pub(super) measured_fit: Option, + pub(super) slots: usize, + pub(super) context_length: u32, + /// True when slots came from the flat auto default rather than an + /// explicit override. + pub(super) slots_auto: bool, + /// True when the context length came from planning rather than an + /// explicit override. + pub(super) context_auto: bool, } #[derive(Clone, Copy, Debug)] @@ -117,6 +170,29 @@ pub(super) struct RuntimeResourcePlanInput<'a> { /// `None` means the whole model is local (fraction = 1.0). pub(super) local_layer_fraction: Option, pub(super) planning_profile: RuntimeResourcePlanningProfile, + /// Measured native buffer footprint from a prior context init of the same + /// shape on this node (compute + KV buffer sizes and the context length + /// they were measured at). When present, the planner charges these + /// measured sizes instead of the 85% KV tax (budget-driven sizing); when + /// absent it falls back to the static tax ladder. + pub(super) measured_buffers: Option, +} + +/// Measured native buffer sizes from one context init, the ground truth the +/// budget-driven planner charges in place of the KV-scaled tax. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) struct MeasuredBufferFootprint { + /// Compute-graph buffer(s) at context init, bytes. + pub(super) compute_bytes: u64, + /// KV buffer at `context_length`, bytes. + pub(super) kv_bytes: u64, + /// The context length the KV buffer was measured at. + pub(super) context_length: u32, + /// Lane count the buffers were measured at. Compute buffers scale + /// ~linearly with lanes (measured 399/783/1551 MiB at 2/4/8 lanes on the + /// 5080 for granite), so a plan resolving a different lane count scales + /// the compute charge by the lane ratio. + pub(super) lane_count: u32, } /// Plan context length and parallel slots. @@ -128,16 +204,80 @@ pub(super) struct RuntimeResourcePlanInput<'a> { /// override via CLI flags), and an explicit `--ctx-size` override bypasses this /// entirely. pub(super) fn plan_runtime_resources(input: RuntimeResourcePlanInput<'_>) -> RuntimeResourcePlan { - let context_length = input - .ctx_size_override - .unwrap_or_else(|| planned_context_length(&input)); + let context_auto = input.ctx_size_override.is_none(); + let slots_auto = input.parallel_override.is_none(); let slots = input .parallel_override .unwrap_or_else(planned_parallel_slots); + let estimated_kv_bytes_per_token = input + .metadata + .and_then(|metadata| { + input + .kv_cache_quant + .kv_cache_bytes_per_token(metadata) + .map(|bytes| scale_by_layer_fraction(bytes, &input)) + }) + .unwrap_or(0); + let estimated_kv_budget = usable_kv_cache_budget(input.vram_bytes, input.model_bytes); + let estimated_compute_charge = input + .vram_bytes + .saturating_sub(input.model_bytes) + .saturating_sub(estimated_kv_budget); + + let ( + context_length, + kv_budget_bytes, + kv_bytes_per_token, + compute_charge_bytes, + planning_source, + measured_fit, + ) = if let Some(context_length) = input.ctx_size_override { + ( + context_length, + estimated_kv_budget, + estimated_kv_bytes_per_token, + estimated_compute_charge, + RuntimeResourcePlanSource::ExplicitOverride, + None, + ) + } else if let Some(measured) = measured_context_plan(&input, slots) { + ( + measured.context_length, + measured.kv_budget_bytes, + measured.kv_bytes_per_token, + measured.compute_charge_bytes, + RuntimeResourcePlanSource::MeasuredFootprint, + Some(measured.fits), + ) + } else { + ( + planned_context_length(&input), + estimated_kv_budget, + estimated_kv_bytes_per_token, + estimated_compute_charge, + RuntimeResourcePlanSource::StaticEstimate, + None, + ) + }; + let planned_kv_bytes = kv_bytes_per_token.saturating_mul(u64::from(context_length)); RuntimeResourcePlan { context_length, slots, + breakdown: Some(RuntimeResourcePlanBreakdown { + vram_bytes: input.vram_bytes, + model_bytes: input.model_bytes, + kv_budget_bytes, + planned_kv_bytes, + kv_bytes_per_token, + compute_charge_bytes, + planning_source, + measured_fit, + slots, + context_length, + slots_auto, + context_auto, + }), } } @@ -248,6 +388,58 @@ fn usable_kv_cache_budget(vram_bytes: u64, model_bytes: u64) -> u64 { budget.min(u128::from(u64::MAX)) as u64 } +/// Measured-vs-charged reconciliation, produced once the native context +/// exists (after model open) and the `sched_reserve` buffer lines have been +/// parsed. +/// +/// `charged_compute_reserve_bytes` is the compute charge the selected planner +/// actually held outside its KV budget. `measured_*` are what llama.cpp +/// allocated, and the residual is what a re-plan with actual free memory would +/// see. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) struct MemoryPlanReconciliation { + /// Compute charge held outside the KV budget by the selected planner. + pub(super) charged_compute_reserve_bytes: u64, + /// Measured compute buffer(s) at context init, if any were observed. + pub(super) measured_compute_bytes: Option, + /// Measured KV buffer(s) at context init, if any were observed. + pub(super) measured_kv_bytes: Option, + /// `vram - model - measured_compute - measured_kv` when both + /// measurements exist; the headroom a re-plan starts from. + pub(super) residual_free_bytes: Option, +} + +pub(super) fn reconcile_memory_plan_with_measurements( + breakdown: &RuntimeResourcePlanBreakdown, + measured: Option, +) -> MemoryPlanReconciliation { + let charged_compute_reserve_bytes = breakdown.compute_charge_bytes; + let host_memory_observed = measured.is_some_and(|measurement| measurement.host_memory_observed); + let mib_to_bytes = |mib: Option| mib.map(|mib| (mib * 1024.0 * 1024.0).round() as u64); + let measured_compute_bytes = mib_to_bytes(measured.and_then(|m| m.compute_mib)); + let measured_kv_bytes = mib_to_bytes(measured.and_then(|m| m.kv_mib)); + let residual_free_bytes = match ( + host_memory_observed, + measured_compute_bytes, + measured_kv_bytes, + ) { + (false, Some(compute), Some(kv)) => Some( + breakdown + .vram_bytes + .saturating_sub(breakdown.model_bytes) + .saturating_sub(compute) + .saturating_sub(kv), + ), + _ => None, + }; + MemoryPlanReconciliation { + charged_compute_reserve_bytes, + measured_compute_bytes, + measured_kv_bytes, + residual_free_bytes, + } +} + fn fallback_context_length(input: &RuntimeResourcePlanInput<'_>) -> u32 { let free_bytes = input.vram_bytes.saturating_sub(input.model_bytes); if free_bytes >= FALLBACK_CONTEXT_64K_FREE_BYTES { @@ -273,6 +465,91 @@ fn snap_context_length_down(value: u32) -> u32 { .unwrap_or(value) } +/// Fraction of usable memory the budget-driven planner targets (Step 2). +/// +/// vLLM reserves ~8% of free memory after weights (`gpu_memory_utilization` +/// 0.92), SGLang ~10% (`mem-fraction-static` 0.9). Ours is deliberately a +/// little more conservative: mesh nodes can co-host other stages, and Metal +/// unified-memory nodes share the pool with the OS (where an over-commit +/// page-out stalls decode rather than failing loudly like a CUDA OOM). +const DEFAULT_UTILIZATION_TARGET_NUMERATOR: u64 = 88; +const DEFAULT_UTILIZATION_TARGET_DENOMINATOR: u64 = 100; + +/// Budget-driven context planning (Step 2) over the measured footprint from a +/// prior context init. +/// +/// Model: `budget = (vram - model) × utilization - measured_compute`, then +/// solve for the deepest context whose *scaled* KV cost fits the budget. KV +/// scales linearly with context (unified pool of `n_ctx` cells), so the +/// measured KV bytes at `measured_ctx` give `kv_bytes_per_token_measured = +/// kv_bytes / measured_ctx`, and the deepest affordable context is +/// `budget / kv_bytes_per_token_measured`. +/// +/// Returns `None` only when the measurement is structurally unusable (zero +/// context, KV, or lanes), letting the static tax ladder answer instead. A +/// valid measurement that cannot fit the minimum context returns an explicit +/// `fits = false` plan so the caller fails closed. Compute buffers do not scale +/// linearly with context (ubatch/graph shape dominates), so the measured value +/// is charged as-is apart from the measured-to-requested lane ratio. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +struct MeasuredContextPlan { + context_length: u32, + kv_budget_bytes: u64, + kv_bytes_per_token: u64, + compute_charge_bytes: u64, + fits: bool, +} + +fn measured_context_plan( + input: &RuntimeResourcePlanInput<'_>, + resolved_lanes: usize, +) -> Option { + let measured = input.measured_buffers?; + if measured.context_length == 0 || measured.kv_bytes == 0 || measured.lane_count == 0 { + return None; + } + let metadata = input.metadata?; + let native_context = metadata.context_length; + if native_context == 0 { + return None; + } + let native_context = native_context.min(MAX_AUTO_CONTEXT_LENGTH); + + let post_weight = input.vram_bytes.saturating_sub(input.model_bytes); + let utilised = u128::from(post_weight) * u128::from(DEFAULT_UTILIZATION_TARGET_NUMERATOR) + / u128::from(DEFAULT_UTILIZATION_TARGET_DENOMINATOR); + // Scale the measured compute charge by the lane ratio when the plan's + // resolved lane count differs from the one the footprint was measured at: + // compute buffers scale ~linearly with lanes (CUDA_Host + device graphs), + // while the unified KV pool is lane-invariant. Linear-through-origin + // slightly overestimates when scaling up (~2-3% at 4→8 on measured data), + // which is the conservative direction; scaling down is symmetric. + let lane_ratio_num = u128::from(resolved_lanes.max(1) as u64); + let lane_ratio_den = u128::from(measured.lane_count); + let compute_charge = u128::from(measured.compute_bytes) * lane_ratio_num / lane_ratio_den; + let budget = utilised.saturating_sub(compute_charge); + let measured_context = u128::from(measured.context_length); + let kv_per_token = u128::from(measured.kv_bytes).div_ceil(measured_context); + if kv_per_token == 0 { + return None; + } + let max_affordable = (budget / kv_per_token).min(u128::from(u32::MAX)) as u32; + let minimum = MIN_AUTO_CONTEXT_LENGTH.min(native_context); + let fits = max_affordable >= minimum; + let context_length = if fits { + snap_context_length_down(max_affordable.min(native_context)).max(minimum) + } else { + minimum + }; + Some(MeasuredContextPlan { + context_length, + kv_budget_bytes: budget.min(u128::from(u64::MAX)) as u64, + kv_bytes_per_token: kv_per_token.min(u128::from(u64::MAX)) as u64, + compute_charge_bytes: compute_charge.min(u128::from(u64::MAX)) as u64, + fits, + }) +} + #[cfg(test)] mod tests { use super::*; @@ -301,6 +578,7 @@ mod tests { kv_cache_quant: GgufKvCacheQuant::Q8_0, local_layer_fraction: None, planning_profile: RuntimeResourcePlanningProfile::DedicatedLocal, + measured_buffers: None, }); assert_eq!(plan.context_length, 16_384); @@ -319,6 +597,7 @@ mod tests { kv_cache_quant: GgufKvCacheQuant::Q8_0, local_layer_fraction: None, planning_profile: RuntimeResourcePlanningProfile::DedicatedLocal, + measured_buffers: None, }); assert_eq!( @@ -343,6 +622,7 @@ mod tests { kv_cache_quant: GgufKvCacheQuant::F16, local_layer_fraction: None, planning_profile: RuntimeResourcePlanningProfile::DedicatedLocal, + measured_buffers: None, }); let q8_plan = plan_runtime_resources(RuntimeResourcePlanInput { ctx_size_override: None, @@ -353,6 +633,7 @@ mod tests { kv_cache_quant: GgufKvCacheQuant::Q8_0, local_layer_fraction: None, planning_profile: RuntimeResourcePlanningProfile::DedicatedLocal, + measured_buffers: None, }); assert!( @@ -374,6 +655,7 @@ mod tests { kv_cache_quant: GgufKvCacheQuant::Q8_0, local_layer_fraction: None, planning_profile: RuntimeResourcePlanningProfile::DedicatedLocal, + measured_buffers: None, }); assert_eq!(plan.context_length, 16_384); @@ -397,6 +679,7 @@ mod tests { kv_cache_quant: GgufKvCacheQuant::Q8_0, local_layer_fraction: None, planning_profile: profile, + measured_buffers: None, }; let dedicated_plan = plan_runtime_resources(input(RuntimeResourcePlanningProfile::DedicatedLocal)); @@ -424,6 +707,7 @@ mod tests { kv_cache_quant: GgufKvCacheQuant::Q8_0, local_layer_fraction: None, planning_profile: RuntimeResourcePlanningProfile::SharedMesh, + measured_buffers: None, }); assert_eq!( @@ -456,6 +740,7 @@ mod tests { kv_cache_quant: GgufKvCacheQuant::Q8_0, local_layer_fraction: None, planning_profile: RuntimeResourcePlanningProfile::SharedMesh, + measured_buffers: None, }); assert_eq!( @@ -482,6 +767,7 @@ mod tests { kv_cache_quant: GgufKvCacheQuant::Q8_0, local_layer_fraction: None, planning_profile: RuntimeResourcePlanningProfile::DedicatedLocal, + measured_buffers: None, }); assert_eq!(plan.context_length, 32_768); @@ -514,6 +800,7 @@ mod tests { kv_cache_quant: GgufKvCacheQuant::Q8_0, local_layer_fraction: None, planning_profile: RuntimeResourcePlanningProfile::DedicatedLocal, + measured_buffers: None, }); assert_eq!(plan.context_length, 32_768); @@ -538,6 +825,7 @@ mod tests { kv_cache_quant: GgufKvCacheQuant::Q8_0, local_layer_fraction: None, planning_profile: RuntimeResourcePlanningProfile::DedicatedLocal, + measured_buffers: None, }); assert_eq!(plan.slots, 8); @@ -569,6 +857,7 @@ mod tests { kv_cache_quant: GgufKvCacheQuant::Q8_0, local_layer_fraction: None, planning_profile: RuntimeResourcePlanningProfile::DedicatedLocal, + measured_buffers: None, }); // With split awareness: local model ~174 GB, local KV fraction 0.66 @@ -581,6 +870,7 @@ mod tests { kv_cache_quant: GgufKvCacheQuant::Q8_0, local_layer_fraction: Some(local_fraction), planning_profile: RuntimeResourcePlanningProfile::DedicatedLocal, + measured_buffers: None, }); assert!( @@ -615,6 +905,7 @@ mod tests { kv_cache_quant: quant, local_layer_fraction: None, planning_profile: RuntimeResourcePlanningProfile::DedicatedLocal, + measured_buffers: None, }) }; let q8_plan = plan_with(GgufKvCacheQuant::Q8_0); @@ -627,4 +918,258 @@ mod tests { q4_plan.slots, q8_plan.slots ); } + + #[test] + fn budget_driven_context_uses_measured_kv_and_compute() { + // Roomy node: 16 GiB VRAM, 3 GiB weights. Here the measured KV + // per-token cost (2 GiB / 16_384 = 131_072 B) is twice the q8 + // estimate (~69_632 B) the ladder assumes, so the budget-driven plan + // correctly lands at 65_536 while the estimate-based ladder would + // clamp at native 131_072. Measured reality cuts both ways: when the + // real KV cost is higher than estimated, the correct plan is + // shallower, not deeper. (Conversely, charging measured compute under + // an 88% utilization target only buys depth over the 85% tax when + // compute is under ~3% of free memory — the estimate is what binds + // on roomy nodes, not the tax.) + let metadata = gqa_metadata(131_072); + let plan = plan_runtime_resources(RuntimeResourcePlanInput { + ctx_size_override: None, + parallel_override: None, + model_bytes: 3 * 1024 * 1024 * 1024, + vram_bytes: 16 * 1024 * 1024 * 1024, + metadata: Some(&metadata), + kv_cache_quant: GgufKvCacheQuant::Q8_0, + local_layer_fraction: None, + planning_profile: RuntimeResourcePlanningProfile::DedicatedLocal, + measured_buffers: Some(MeasuredBufferFootprint { + compute_bytes: 512 * 1024 * 1024, + kv_bytes: 2 * 1024 * 1024 * 1024, + context_length: 16_384, + lane_count: 4, + }), + }); + assert_eq!(plan.context_length, 65_536); + let measured_breakdown = plan.breakdown.expect("measured planner breakdown"); + assert_eq!( + measured_breakdown.planning_source, + RuntimeResourcePlanSource::MeasuredFootprint + ); + assert_eq!(measured_breakdown.kv_bytes_per_token, 131_072); + assert_eq!(measured_breakdown.compute_charge_bytes, 512 * 1024 * 1024); + assert_eq!(measured_breakdown.measured_fit, Some(true)); + assert_eq!(measured_breakdown.planned_kv_bytes, 65_536 * 131_072); + + // Tight node where the ladder's estimate binds: 5 GiB free, the q8 + // estimate (65,536 B/tok for these dims) caps the ladder at 65_536, + // while a measured KV cost half the estimate (32,768 B/tok measured + // at 16_384) lets the budget-driven plan reach the full native + // window. + let tight = plan_runtime_resources(RuntimeResourcePlanInput { + ctx_size_override: None, + parallel_override: None, + model_bytes: 1024 * 1024 * 1024, + vram_bytes: 6 * 1024 * 1024 * 1024, + metadata: Some(&metadata), + kv_cache_quant: GgufKvCacheQuant::Q8_0, + local_layer_fraction: None, + planning_profile: RuntimeResourcePlanningProfile::DedicatedLocal, + measured_buffers: Some(MeasuredBufferFootprint { + compute_bytes: 128 * 1024 * 1024, + kv_bytes: 512 * 1024 * 1024, + context_length: 16_384, + lane_count: 4, + }), + }); + assert_eq!(tight.context_length, 131_072); + + let tight_ladder = plan_runtime_resources(RuntimeResourcePlanInput { + ctx_size_override: None, + parallel_override: None, + model_bytes: 1024 * 1024 * 1024, + vram_bytes: 6 * 1024 * 1024 * 1024, + metadata: Some(&metadata), + kv_cache_quant: GgufKvCacheQuant::Q8_0, + local_layer_fraction: None, + planning_profile: RuntimeResourcePlanningProfile::DedicatedLocal, + measured_buffers: None, + }); + assert_eq!(tight_ladder.context_length, 65_536); + } + + #[test] + fn budget_driven_context_degrades_to_ladder_when_measurement_is_unusable() { + let metadata = gqa_metadata(131_072); + let base = RuntimeResourcePlanInput { + ctx_size_override: None, + parallel_override: None, + model_bytes: 5_000_000_000, + vram_bytes: 24_000_000_000, + metadata: Some(&metadata), + kv_cache_quant: GgufKvCacheQuant::Q8_0, + local_layer_fraction: None, + planning_profile: RuntimeResourcePlanningProfile::DedicatedLocal, + measured_buffers: None, + }; + + // Zero-context, zero-KV, or zero-lane measurements cannot drive the + // budget model; the plan must equal the static-ladder answer. + let ladder = plan_runtime_resources(base).context_length; + for unusable in [ + MeasuredBufferFootprint { + compute_bytes: 512 * 1024 * 1024, + kv_bytes: 0, + context_length: 16_384, + lane_count: 4, + }, + MeasuredBufferFootprint { + compute_bytes: 512 * 1024 * 1024, + kv_bytes: 2 * 1024 * 1024 * 1024, + context_length: 0, + lane_count: 4, + }, + MeasuredBufferFootprint { + compute_bytes: 512 * 1024 * 1024, + kv_bytes: 2 * 1024 * 1024 * 1024, + context_length: 16_384, + lane_count: 0, + }, + ] { + let degraded = plan_runtime_resources(RuntimeResourcePlanInput { + measured_buffers: Some(unusable), + ..base + }); + assert_eq!(degraded.context_length, ladder); + } + } + + #[test] + fn budget_driven_context_reports_exhausted_measurement_without_fallback() { + let metadata = gqa_metadata(131_072); + let plan = plan_runtime_resources(RuntimeResourcePlanInput { + ctx_size_override: None, + parallel_override: None, + model_bytes: 5 * 1024 * 1024 * 1024, + vram_bytes: 6 * 1024 * 1024 * 1024, + metadata: Some(&metadata), + kv_cache_quant: GgufKvCacheQuant::Q8_0, + local_layer_fraction: None, + planning_profile: RuntimeResourcePlanningProfile::DedicatedLocal, + measured_buffers: Some(MeasuredBufferFootprint { + compute_bytes: 1024 * 1024 * 1024, + kv_bytes: 2 * 1024 * 1024 * 1024, + context_length: 16_384, + lane_count: 4, + }), + }); + let breakdown = plan.breakdown.expect("planner breakdown"); + assert_eq!( + breakdown.planning_source, + RuntimeResourcePlanSource::MeasuredFootprint + ); + assert_eq!(breakdown.measured_fit, Some(false)); + assert_eq!(plan.context_length, MIN_AUTO_CONTEXT_LENGTH); + assert!(breakdown.planned_kv_bytes > breakdown.kv_budget_bytes); + } + + #[test] + fn budget_driven_compute_charge_scales_with_lane_count() { + // Compute buffers scale ~linearly with lanes (measured 399/783/1551 + // MiB at 2/4/8 on the 5080). A footprint measured at 4 lanes must be + // charged at 2x when the plan resolves 8 lanes, and at 0.5x when it + // resolves 2 — and the context depth must follow the charge. + let metadata = gqa_metadata(131_072); + let footprint = MeasuredBufferFootprint { + compute_bytes: 512 * 1024 * 1024, + kv_bytes: 2 * 1024 * 1024 * 1024, + context_length: 16_384, + lane_count: 4, + }; + let plan_at = |parallel_override: Option| { + plan_runtime_resources(RuntimeResourcePlanInput { + ctx_size_override: None, + parallel_override, + model_bytes: 3 * 1024 * 1024 * 1024, + vram_bytes: 16 * 1024 * 1024 * 1024, + metadata: Some(&metadata), + kv_cache_quant: GgufKvCacheQuant::Q8_0, + local_layer_fraction: None, + planning_profile: RuntimeResourcePlanningProfile::DedicatedLocal, + measured_buffers: Some(footprint), + }) + .context_length + }; + + let at2 = plan_at(Some(2)); + let at4 = plan_at(Some(4)); + let at8 = plan_at(Some(8)); + // Scaling the charge DOWN (fewer lanes than measured) frees budget and + // can only deepen (or hold) the plan; scaling UP can only shallow it. + assert!(at2 >= at4); + assert!(at8 <= at4); + // At the measured lane count the charge is exactly the measured value, + // so this matches the roomy-node case of + // budget_driven_context_uses_measured_kv_and_compute. + assert_eq!(at4, 65_536); + } + + #[test] + fn reconciliation_exposes_overcharge_and_residual() { + let breakdown = RuntimeResourcePlanBreakdown { + vram_bytes: 16 * 1024 * 1024 * 1024, + model_bytes: 3 * 1024 * 1024 * 1024, + kv_budget_bytes: (13 * 1024 * 1024 * 1024) * 85 / 100, + planned_kv_bytes: 2 * 1024 * 1024 * 1024, + kv_bytes_per_token: 131_072, + compute_charge_bytes: (13 * 1024 * 1024 * 1024) - (13 * 1024 * 1024 * 1024) * 85 / 100, + planning_source: RuntimeResourcePlanSource::StaticEstimate, + measured_fit: None, + slots: 4, + context_length: 16_384, + slots_auto: true, + context_auto: true, + }; + + // Without measurements the reconciliation still reports the charged + // proxy so the log line carries the plan side of the story. The proxy + // is what the 85% budget floor leaves behind, so derive it the same + // way rather than as a separate 15% floor (85+15 != 100 in integer + // math). + let post_weight_bytes = 13 * 1024 * 1024 * 1024; + let unmeasured = reconcile_memory_plan_with_measurements(&breakdown, None); + assert_eq!( + unmeasured.charged_compute_reserve_bytes, + post_weight_bytes - post_weight_bytes * 85 / 100 + ); + assert_eq!(unmeasured.measured_compute_bytes, None); + assert_eq!(unmeasured.residual_free_bytes, None); + + // With measurements: compute 0.5 GiB, KV 2.0 GiB over a 16 GiB node + // holding 3 GiB of weights leaves 10.5 GiB residual — far above the + // ~2 GiB proxy tax the planner charged. + let measured = skippy_runtime::MeasuredNativeBuffers { + compute_mib: Some(512.0), + kv_mib: Some(2048.0), + host_memory_observed: false, + }; + let reconciled = reconcile_memory_plan_with_measurements(&breakdown, Some(measured)); + assert_eq!(reconciled.measured_compute_bytes, Some(512 * 1024 * 1024)); + assert_eq!(reconciled.measured_kv_bytes, Some(2048 * 1024 * 1024)); + assert_eq!( + reconciled.residual_free_bytes, + Some(10 * 1024 * 1024 * 1024 + 512 * 1024 * 1024) + ); + assert_eq!( + reconciled.charged_compute_reserve_bytes, + unmeasured.charged_compute_reserve_bytes + ); + + let host_offloaded = reconcile_memory_plan_with_measurements( + &breakdown, + Some(skippy_runtime::MeasuredNativeBuffers { + host_memory_observed: true, + ..measured + }), + ); + assert_eq!(host_offloaded.residual_free_bytes, None); + } } diff --git a/crates/mesh-llm-host-runtime/src/runtime/kv_disk_config.rs b/crates/mesh-llm-host-runtime/src/runtime/kv_disk_config.rs new file mode 100644 index 0000000000..727c31a5ba --- /dev/null +++ b/crates/mesh-llm-host-runtime/src/runtime/kv_disk_config.rs @@ -0,0 +1,611 @@ +use std::ffi::OsString; +use std::path::{Path, PathBuf}; +use std::sync::{OnceLock, RwLock}; + +use anyhow::{Context, Result, bail}; +use mesh_llm_config::{ + DEFAULT_KV_DISK_MINIMUM_FREE_MIB, KvDiskTierMode, MIN_KV_DISK_MINIMUM_FREE_MIB, MeshConfig, + parse_iec_size, +}; +use serde::Serialize; +use skippy_cache::{L3CacheManager, StoreLimits}; + +use super::RuntimeOptions; + +const MIB: u64 = 1024 * 1024; +const LEGACY_DEFAULT_BUDGET_BYTES: u64 = 32 * 1024 * 1024 * 1024; +const AUTO_MAX_BUDGET_BYTES: u64 = 64 * 1024 * 1024 * 1024; + +static NODE_KV_DISK_CACHE: OnceLock> = OnceLock::new(); + +#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum KvDiskConfigSource { + Default, + Config, + Environment, + Cli, + LegacyEnvironment, +} + +#[derive(Clone, Debug, Eq, PartialEq, Serialize)] +pub(crate) struct KvDiskConfigSources { + pub(crate) mode: KvDiskConfigSource, + pub(crate) directory: KvDiskConfigSource, + pub(crate) budget: KvDiskConfigSource, + pub(crate) minimum_free: KvDiskConfigSource, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) struct ResolvedKvDiskConfig { + pub(crate) mode: KvDiskTierMode, + pub(crate) directory: PathBuf, + /// Fixed-mode cap. Auto mode resolves this from live filesystem facts + /// immediately before the node manager opens the root. + pub(crate) budget_bytes: Option, + pub(crate) minimum_free_bytes: u64, + pub(crate) sources: KvDiskConfigSources, + pub(crate) warnings: Vec, +} + +#[derive(Clone, Debug)] +pub(crate) struct NodeKvDiskCache { + pub(crate) configured: ResolvedKvDiskConfig, + pub(crate) manager: Option, + runtime_options: RuntimeOptions, +} + +pub(crate) struct KvDiskLiveRollback { + configured: ResolvedKvDiskConfig, + limits: Option, +} + +impl ResolvedKvDiskConfig { + pub(crate) fn enabled(&self) -> bool { + self.mode != KvDiskTierMode::Off + } +} + +pub(crate) fn resolve_kv_disk_config( + config: &MeshConfig, + options: &RuntimeOptions, +) -> Result { + resolve_kv_disk_config_with_env(config, options, |name| std::env::var_os(name)) +} + +/// Resolve the public/legacy configuration once and acquire the sole cache +/// manager for this node process. Store availability is fail-open for +/// inference: an unusable cache is reported as a warning and cold prefill +/// remains available. +pub(crate) fn configure_node_kv_disk_cache( + config: &MeshConfig, + options: &RuntimeOptions, +) -> Result { + if let Some(cache) = NODE_KV_DISK_CACHE.get() { + let snapshot = cache + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); + return Ok(snapshot); + } + let mut configured = resolve_kv_disk_config(config, options)?; + let mut manager = None; + if configured.enabled() { + match acquire_manager(&mut configured) { + Ok(acquired) => manager = acquired, + Err(error) => configured.warnings.push(format!( + "disk prompt cache is unavailable; inference will use cold prefill: {error:#}" + )), + } + } + let snapshot = NodeKvDiskCache { + configured, + manager, + runtime_options: options.clone(), + }; + let _ = NODE_KV_DISK_CACHE.set(RwLock::new(snapshot.clone())); + Ok(snapshot) +} + +pub(crate) fn node_kv_disk_manager() -> Option { + NODE_KV_DISK_CACHE.get().and_then(|cache| { + cache + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .manager + .clone() + }) +} + +pub(crate) fn node_kv_disk_cache() -> Option { + NODE_KV_DISK_CACHE.get().map(|cache| { + cache + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + }) +} + +pub(crate) fn apply_live_kv_disk_limits(config: &MeshConfig) -> Result { + let cache = NODE_KV_DISK_CACHE + .get() + .context("node disk prompt-cache runtime is not initialized")?; + let mut cache = cache + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut resolved = resolve_kv_disk_config(config, &cache.runtime_options)?; + let previous = cache.configured.clone(); + let previous_limits = cache.manager.as_ref().map(L3CacheManager::limits); + preserve_restart_fields(&previous, &mut resolved); + if let Some(manager) = cache.manager.as_ref() { + let current = manager.limits(); + let budget = match resolved.mode { + KvDiskTierMode::Fixed => resolved.budget_bytes.unwrap_or(current.budget_bytes), + KvDiskTierMode::Auto => auto_budget_bytes(manager.root(), resolved.minimum_free_bytes)? + .min(current.budget_bytes), + KvDiskTierMode::Off => current.budget_bytes, + }; + resolved.budget_bytes = Some(budget); + manager.update_limits(StoreLimits::new(budget, resolved.minimum_free_bytes))?; + } + cache.configured = resolved; + Ok(KvDiskLiveRollback { + configured: previous, + limits: previous_limits, + }) +} + +fn preserve_restart_fields(previous: &ResolvedKvDiskConfig, next: &mut ResolvedKvDiskConfig) { + if next.mode != previous.mode { + next.mode = previous.mode; + next.sources.mode = previous.sources.mode; + next.budget_bytes = previous.budget_bytes; + next.sources.budget = previous.sources.budget; + } + if next.directory != previous.directory { + next.directory = previous.directory.clone(); + next.sources.directory = previous.sources.directory; + } +} + +pub(crate) fn restore_live_kv_disk_limits(rollback: KvDiskLiveRollback) { + let Some(cache) = NODE_KV_DISK_CACHE.get() else { + return; + }; + let mut cache = cache + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let (Some(manager), Some(limits)) = (cache.manager.as_ref(), rollback.limits) { + let _ = manager.update_limits(limits); + } + cache.configured = rollback.configured; +} + +fn acquire_manager(config: &mut ResolvedKvDiskConfig) -> Result> { + std::fs::create_dir_all(&config.directory).with_context(|| { + format!( + "create disk prompt-cache root {}", + config.directory.display() + ) + })?; + let budget_bytes = match config.mode { + KvDiskTierMode::Off => return Ok(None), + KvDiskTierMode::Fixed => config + .budget_bytes + .context("fixed disk prompt-cache mode has no budget")?, + KvDiskTierMode::Auto => { + let budget = auto_budget_bytes(&config.directory, config.minimum_free_bytes)?; + config.budget_bytes = Some(budget); + if budget == 0 { + config.warnings.push( + "automatic disk prompt-cache budget is zero after the minimum-free reserve; disk cache remains disabled" + .to_string(), + ); + return Ok(None); + } + budget + } + }; + Ok(Some(L3CacheManager::acquire( + &config.directory, + StoreLimits::new(budget_bytes, config.minimum_free_bytes), + )?)) +} + +fn auto_budget_bytes(root: &Path, minimum_free_bytes: u64) -> Result { + let available = skippy_cache::fsinfo::available_bytes(root)?; + let managed = managed_root_bytes(root)?; + Ok(auto_budget_from_space( + available, + managed, + minimum_free_bytes, + )) +} + +fn auto_budget_from_space(available: u64, managed: u64, minimum_free: u64) -> u64 { + let capacity_basis = available.saturating_add(managed); + let twenty_percent = capacity_basis / 5; + let allocatable = available + .saturating_sub(minimum_free) + .saturating_add(managed); + twenty_percent.min(allocatable).min(AUTO_MAX_BUDGET_BYTES) +} + +fn managed_root_bytes(root: &Path) -> Result { + let mut total = 0_u64; + let mut pending = vec![root.to_path_buf()]; + while let Some(directory) = pending.pop() { + for entry in std::fs::read_dir(&directory) + .with_context(|| format!("read cache directory {}", directory.display()))? + { + let entry = entry?; + let metadata = std::fs::symlink_metadata(entry.path())?; + if metadata.file_type().is_symlink() { + bail!( + "disk prompt-cache root contains a symlink: {}", + entry.path().display() + ); + } + if metadata.is_dir() { + pending.push(entry.path()); + } else { + total = total.saturating_add(metadata.len()); + } + } + } + Ok(total) +} + +fn resolve_kv_disk_config_with_env( + config: &MeshConfig, + options: &RuntimeOptions, + env: impl Fn(&str) -> Option, +) -> Result { + let disk = &config.runtime.kv_cache.disk; + let mut warnings = Vec::new(); + + let mesh_home = match env("MESH_LLM_HOME") { + Some(value) => absolute_path(value, "MESH_LLM_HOME")?, + None => dirs::home_dir() + .context("cannot determine home directory for the disk prompt cache")? + .join(".mesh-llm"), + }; + let mut directory = mesh_home.join("kv-cache"); + let mut directory_source = KvDiskConfigSource::Default; + if let Some(value) = disk.directory.as_ref() { + directory = value.clone(); + directory_source = KvDiskConfigSource::Config; + } + + let mut mode = disk.mode.unwrap_or_default(); + let mut mode_source = if disk.mode.is_some() { + KvDiskConfigSource::Config + } else { + KvDiskConfigSource::Default + }; + let mut budget_bytes = disk.budget_mib.map(|mib| mib.saturating_mul(MIB)); + let mut budget_source = if disk.budget_mib.is_some() { + KvDiskConfigSource::Config + } else { + KvDiskConfigSource::Default + }; + let mut minimum_free_bytes = disk + .minimum_free_mib + .unwrap_or(DEFAULT_KV_DISK_MINIMUM_FREE_MIB) + .saturating_mul(MIB); + let mut minimum_free_source = if disk.minimum_free_mib.is_some() { + KvDiskConfigSource::Config + } else { + KvDiskConfigSource::Default + }; + + let public_env_disk = env_utf8(&env, "MESH_LLM_KV_CACHE_DISK")?; + let public_env_directory = env_path(&env, "MESH_LLM_KV_CACHE_DISK_DIR")?; + let public_env_minimum = env_utf8(&env, "MESH_LLM_KV_CACHE_MIN_FREE")?; + let public_mode = + disk.mode.is_some() || public_env_disk.is_some() || options.kv_cache_disk.is_some(); + let public_directory = disk.directory.is_some() + || public_env_directory.is_some() + || options.kv_cache_disk_dir.is_some(); + let public_budget = disk.budget_mib.is_some() || public_mode; + + if let Some(legacy_directory) = env_path(&env, "SKIPPY_L3_DIR")? { + let mut used_legacy = false; + if !public_directory { + directory = require_absolute(legacy_directory, "SKIPPY_L3_DIR")?; + directory_source = KvDiskConfigSource::LegacyEnvironment; + used_legacy = true; + } + if !public_mode { + mode = KvDiskTierMode::Fixed; + mode_source = KvDiskConfigSource::LegacyEnvironment; + used_legacy = true; + } + if !public_budget { + budget_source = KvDiskConfigSource::LegacyEnvironment; + budget_bytes = match env_utf8(&env, "SKIPPY_L3_BUDGET_BYTES")? { + Some(value) => match value.parse::() { + Ok(0) => { + warnings.push( + "SKIPPY_L3_BUDGET_BYTES=0 no longer means unbounded; using the legacy 32GiB default" + .to_string(), + ); + Some(LEGACY_DEFAULT_BUDGET_BYTES) + } + Ok(value) => Some(value), + Err(_) => bail!("SKIPPY_L3_BUDGET_BYTES must be a positive byte count"), + }, + None => Some(LEGACY_DEFAULT_BUDGET_BYTES), + }; + used_legacy = true; + } + if used_legacy { + warnings.push( + "SKIPPY_L3_DIR and SKIPPY_L3_BUDGET_BYTES are deprecated; use runtime.kv_cache.disk or MESH_LLM_KV_CACHE_*" + .to_string(), + ); + } + } else if env("SKIPPY_L3_BUDGET_BYTES").is_some() { + warnings.push( + "ignoring deprecated SKIPPY_L3_BUDGET_BYTES because SKIPPY_L3_DIR is not set" + .to_string(), + ); + } + + if let Some(value) = public_env_disk.as_deref() { + (mode, budget_bytes) = parse_mode_or_size(value, "MESH_LLM_KV_CACHE_DISK")?; + mode_source = KvDiskConfigSource::Environment; + budget_source = KvDiskConfigSource::Environment; + } + if let Some(value) = public_env_directory { + directory = require_absolute(value, "MESH_LLM_KV_CACHE_DISK_DIR")?; + directory_source = KvDiskConfigSource::Environment; + } + if let Some(value) = public_env_minimum.as_deref() { + minimum_free_bytes = parse_minimum_free(value, "MESH_LLM_KV_CACHE_MIN_FREE")?; + minimum_free_source = KvDiskConfigSource::Environment; + } + + if let Some(value) = options.kv_cache_disk.as_deref() { + (mode, budget_bytes) = parse_mode_or_size(value, "--kv-cache-disk")?; + mode_source = KvDiskConfigSource::Cli; + budget_source = KvDiskConfigSource::Cli; + } + if let Some(value) = options.kv_cache_disk_dir.as_ref() { + directory = require_absolute(value.clone(), "--kv-cache-disk-dir")?; + directory_source = KvDiskConfigSource::Cli; + } + if let Some(value) = options.kv_cache_min_free.as_deref() { + minimum_free_bytes = parse_minimum_free(value, "--kv-cache-min-free")?; + minimum_free_source = KvDiskConfigSource::Cli; + } + + match mode { + KvDiskTierMode::Fixed if budget_bytes.is_none() => { + bail!("fixed disk prompt-cache mode requires a positive budget") + } + KvDiskTierMode::Off | KvDiskTierMode::Auto => budget_bytes = None, + KvDiskTierMode::Fixed => {} + } + + Ok(ResolvedKvDiskConfig { + mode, + directory, + budget_bytes, + minimum_free_bytes, + sources: KvDiskConfigSources { + mode: mode_source, + directory: directory_source, + budget: budget_source, + minimum_free: minimum_free_source, + }, + warnings, + }) +} + +fn parse_mode_or_size(value: &str, name: &str) -> Result<(KvDiskTierMode, Option)> { + match value.trim() { + "off" => Ok((KvDiskTierMode::Off, None)), + "auto" => Ok((KvDiskTierMode::Auto, None)), + size => parse_iec_size(size) + .map(|bytes| (KvDiskTierMode::Fixed, Some(bytes))) + .with_context(|| format!("invalid {name} value {value:?}")), + } +} + +fn parse_minimum_free(value: &str, name: &str) -> Result { + let bytes = parse_iec_size(value).with_context(|| format!("invalid {name} value {value:?}"))?; + let minimum = MIN_KV_DISK_MINIMUM_FREE_MIB.saturating_mul(MIB); + if bytes < minimum { + bail!("{name} must preserve at least {MIN_KV_DISK_MINIMUM_FREE_MIB}MiB"); + } + Ok(bytes) +} + +fn env_utf8(env: &impl Fn(&str) -> Option, name: &str) -> Result> { + env(name) + .map(|value| { + value + .into_string() + .map_err(|_| anyhow::anyhow!("{name} must contain valid UTF-8")) + }) + .transpose() +} + +fn env_path(env: &impl Fn(&str) -> Option, name: &str) -> Result> { + Ok(env(name).map(PathBuf::from)) +} + +fn absolute_path(value: OsString, name: &str) -> Result { + require_absolute(PathBuf::from(value), name) +} + +fn require_absolute(path: PathBuf, name: &str) -> Result { + if !is_absolute_on_supported_host(&path) { + bail!("{name} must be an absolute path: {}", path.display()); + } + Ok(path) +} + +fn is_absolute_on_supported_host(path: &Path) -> bool { + if path.is_absolute() { + return true; + } + if !cfg!(windows) { + // `C:\cache` is a relative path on this host; treating it as absolute + // would create the cache root under the working directory. + return false; + } + let rendered = path.to_string_lossy(); + rendered.as_bytes().get(1) == Some(&b':') + && rendered + .as_bytes() + .get(2) + .is_some_and(|byte| matches!(byte, b'/' | b'\\')) +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use super::*; + + fn resolve( + config: &str, + options: RuntimeOptions, + env: &[(&str, &str)], + ) -> Result { + let config = mesh_llm_config::parse_config_toml(config)?; + let env = env + .iter() + .map(|(key, value)| ((*key).to_string(), OsString::from(value))) + .collect::>(); + resolve_kv_disk_config_with_env(&config, &options, |name| env.get(name).cloned()) + } + + #[test] + fn precedence_is_field_by_field_and_reports_each_source() { + let options = RuntimeOptions { + kv_cache_disk: Some("48GiB".to_string()), + kv_cache_min_free: Some("20GiB".to_string()), + ..RuntimeOptions::default() + }; + let resolved = resolve( + "[runtime.kv_cache.disk]\nmode='fixed'\ndirectory='/config/cache'\nbudget_mib=32768\nminimum_free_mib=12288\n", + options, + &[ + ("MESH_LLM_KV_CACHE_DISK", "auto"), + ("MESH_LLM_KV_CACHE_DISK_DIR", "/env/cache"), + ], + ) + .unwrap(); + + assert_eq!(resolved.mode, KvDiskTierMode::Fixed); + assert_eq!(resolved.budget_bytes, Some(48 * 1024_u64.pow(3))); + assert_eq!(resolved.directory, PathBuf::from("/env/cache")); + assert_eq!(resolved.minimum_free_bytes, 20 * 1024_u64.pow(3)); + assert_eq!(resolved.sources.mode, KvDiskConfigSource::Cli); + assert_eq!(resolved.sources.budget, KvDiskConfigSource::Cli); + assert_eq!(resolved.sources.directory, KvDiskConfigSource::Environment); + assert_eq!(resolved.sources.minimum_free, KvDiskConfigSource::Cli); + } + + #[test] + fn legacy_environment_is_fallback_only_and_zero_is_bounded() { + let resolved = resolve( + "", + RuntimeOptions::default(), + &[ + ("SKIPPY_L3_DIR", "/legacy/cache"), + ("SKIPPY_L3_BUDGET_BYTES", "0"), + ], + ) + .unwrap(); + assert_eq!(resolved.mode, KvDiskTierMode::Fixed); + assert_eq!(resolved.budget_bytes, Some(LEGACY_DEFAULT_BUDGET_BYTES)); + assert_eq!(resolved.sources.mode, KvDiskConfigSource::LegacyEnvironment); + assert_eq!(resolved.warnings.len(), 2); + + let public = resolve( + "[runtime.kv_cache.disk]\nmode='off'\n", + RuntimeOptions::default(), + &[("SKIPPY_L3_DIR", "/legacy/cache")], + ) + .unwrap(); + assert_eq!(public.mode, KvDiskTierMode::Off); + assert_eq!(public.sources.mode, KvDiskConfigSource::Config); + assert_eq!( + public.directory, + PathBuf::from("/legacy/cache"), + "legacy directory remains a field-level fallback" + ); + + let minimum_only = resolve( + "[runtime.kv_cache.disk]\nminimum_free_mib=20480\n", + RuntimeOptions::default(), + &[("SKIPPY_L3_DIR", "/legacy/cache")], + ) + .unwrap(); + assert_eq!(minimum_only.mode, KvDiskTierMode::Fixed); + assert_eq!(minimum_only.budget_bytes, Some(LEGACY_DEFAULT_BUDGET_BYTES)); + assert_eq!( + minimum_only.sources.minimum_free, + KvDiskConfigSource::Config + ); + } + + #[test] + fn ambiguous_sizes_and_relative_paths_fail_closed() { + let options = RuntimeOptions { + kv_cache_disk: Some("32".to_string()), + ..RuntimeOptions::default() + }; + assert!(resolve("", options, &[]).is_err()); + + let options = RuntimeOptions { + kv_cache_disk_dir: Some(PathBuf::from("relative/cache")), + ..RuntimeOptions::default() + }; + assert!(resolve("", options, &[]).is_err()); + + let options = RuntimeOptions { + kv_cache_min_free: Some("512MiB".to_string()), + ..RuntimeOptions::default() + }; + assert!(resolve("", options, &[]).is_err()); + } + + #[test] + fn live_apply_preserves_restart_only_mode_directory_and_coupled_budget() { + let previous = resolve( + "[runtime.kv_cache.disk]\nmode='fixed'\ndirectory='/old/cache'\nbudget_mib=32768\nminimum_free_mib=16384\n", + RuntimeOptions::default(), + &[], + ) + .unwrap(); + let mut next = resolve( + "[runtime.kv_cache.disk]\nmode='auto'\ndirectory='/new/cache'\nminimum_free_mib=20480\n", + RuntimeOptions::default(), + &[], + ) + .unwrap(); + + preserve_restart_fields(&previous, &mut next); + + assert_eq!(next.mode, KvDiskTierMode::Fixed); + assert_eq!(next.directory, PathBuf::from("/old/cache")); + assert_eq!(next.budget_bytes, Some(32 * 1024_u64.pow(3))); + assert_eq!(next.minimum_free_bytes, 20 * 1024_u64.pow(3)); + } + + #[test] + fn auto_budget_uses_available_plus_managed_as_stable_capacity_basis() { + let gib = 1024_u64.pow(3); + assert_eq!(auto_budget_from_space(100 * gib, 0, 16 * gib), 20 * gib); + assert_eq!( + auto_budget_from_space(84 * gib, 16 * gib, 16 * gib), + 20 * gib + ); + assert_eq!(auto_budget_from_space(8 * gib, 0, 16 * gib), 0); + } +} diff --git a/crates/mesh-llm-host-runtime/src/runtime/local.rs b/crates/mesh-llm-host-runtime/src/runtime/local.rs index 5169c70397..848424ea9f 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local.rs @@ -3,6 +3,10 @@ use super::context_planning::{ RuntimeResourcePlan, RuntimeResourcePlanInput, RuntimeResourcePlanningProfile, plan_runtime_resources, }; +use super::local_memory_plan::{ + MemoryPlanMeasurementKey, MemoryPlanStartPath, emit_measured_memory_reconciliation, + emit_memory_plan_resolved, measured_buffers_footprint, +}; use super::split_planning::format_gb; use crate::api; use crate::inference::{election, skippy}; @@ -704,6 +708,17 @@ pub(super) async fn start_local_openai_model( effective_cache_type_v, ) .unwrap_or(models::gguf::GgufKvCacheQuant::Q8_0); + let measurement_key = MemoryPlanMeasurementKey::new(format!( + "model={runtime_model_name:?};path={:?};bytes={local_model_bytes};capacity={my_vram};config={:?};config_model={:?};device={:?};pinned_gpu={:?};cache_k={effective_cache_type_k:?};cache_v={effective_cache_type_v:?};batch={:?};ubatch={:?};flash={:?}", + spec.model_path, + spec.mesh_config, + spec.config_model_id, + spec.device_override, + spec.pinned_gpu, + spec.n_batch_override, + spec.n_ubatch_override, + spec.flash_attention_override, + )); let plan = plan_runtime_resources(RuntimeResourcePlanInput { ctx_size_override: spec.ctx_size_override, @@ -714,7 +729,15 @@ pub(super) async fn start_local_openai_model( kv_cache_quant, local_layer_fraction, planning_profile: spec.planning_profile, + measured_buffers: measured_buffers_footprint(&measurement_key), }); + anyhow::ensure!( + !plan + .breakdown + .as_ref() + .is_some_and(|breakdown| breakdown.measured_fit == Some(false)), + "measured native buffers leave no capacity for the minimum context under the current model configuration" + ); if let Some(package) = package { start_local_package_v2_model( @@ -723,6 +746,7 @@ pub(super) async fn start_local_openai_model( progress_ingress, package, plan, + measurement_key, compact_meta.as_ref(), ) .await @@ -732,6 +756,7 @@ pub(super) async fn start_local_openai_model( model_name, progress_ingress, plan, + measurement_key, compact_meta.as_ref(), ) .await @@ -743,6 +768,7 @@ async fn start_local_skippy_model( model_name: String, progress_ingress: Option, plan: RuntimeResourcePlan, + measurement_key: MemoryPlanMeasurementKey, compact_meta: Option<&models::gguf::GgufCompactMeta>, ) -> Result<( String, @@ -750,6 +776,11 @@ async fn start_local_skippy_model( tokio::sync::oneshot::Receiver<()>, )> { let context_length = plan.context_length; + emit_memory_plan_resolved( + &model_name, + plan.breakdown.as_ref(), + MemoryPlanStartPath::Direct, + ); let fallback_projector_path = mmproj_path_for_model(&model_name).filter(|path| path.exists()); let mut resolved = resolve_local_openai_skippy_config( &spec, @@ -813,6 +844,7 @@ async fn start_local_skippy_model( }) .await .context("join load skippy direct GGUF task")??; + emit_measured_memory_reconciliation(&model_name, &measurement_key, &plan); let _ = emit_event(OutputEvent::ModelLoaded { model: model_name.clone(), bytes: None, @@ -844,6 +876,7 @@ async fn start_local_package_v2_model( progress_ingress: Option, package: skippy::SkippyPackageIdentity, plan: RuntimeResourcePlan, + measurement_key: MemoryPlanMeasurementKey, compact_meta: Option<&models::gguf::GgufCompactMeta>, ) -> Result<( String, @@ -880,6 +913,11 @@ async fn start_local_package_v2_model( ) }; let context_length = plan.context_length; + emit_memory_plan_resolved( + &model_name, + plan.breakdown.as_ref(), + MemoryPlanStartPath::PackageV2, + ); let fallback_projector_path = package_projector_path .or_else(|| mmproj_path_for_model(&model_name).filter(|path| path.exists())); let mut resolved = resolve_local_openai_skippy_config( @@ -979,6 +1017,7 @@ async fn start_local_package_v2_model( }) .await .context("join load skippy package-v2 task")??; + emit_measured_memory_reconciliation(&model_name, &measurement_key, &plan); let _ = emit_event(OutputEvent::ModelLoaded { model: model_ref, bytes: None, diff --git a/crates/mesh-llm-host-runtime/src/runtime/local_memory_plan.rs b/crates/mesh-llm-host-runtime/src/runtime/local_memory_plan.rs new file mode 100644 index 0000000000..8bc6b37588 --- /dev/null +++ b/crates/mesh-llm-host-runtime/src/runtime/local_memory_plan.rs @@ -0,0 +1,259 @@ +use std::sync::Mutex; + +use super::context_planning::{ + MeasuredBufferFootprint, RuntimeResourcePlan, RuntimeResourcePlanBreakdown, + reconcile_memory_plan_with_measurements, +}; + +#[derive(Clone, Copy)] +pub(super) enum MemoryPlanStartPath { + Direct, + PackageV2, +} + +/// Host-side tie between this process's measured native buffers and the plan +/// that produced them. The model key, context length, and lane count are +/// written together after model open so later planning cannot combine state +/// from two different starts. +#[derive(Clone, Debug, Eq, PartialEq)] +pub(super) struct MemoryPlanMeasurementKey(String); + +impl MemoryPlanMeasurementKey { + pub(super) fn new(value: String) -> Self { + Self(value) + } +} + +#[derive(Clone, Debug, Eq, PartialEq)] +struct MeasuredPlanSnapshot { + key: MemoryPlanMeasurementKey, + footprint: MeasuredBufferFootprint, +} + +static MEASURED_PLAN_SNAPSHOT: Mutex> = Mutex::new(None); + +fn completed_measurement_snapshot( + key: &MemoryPlanMeasurementKey, + breakdown: &RuntimeResourcePlanBreakdown, + measured: Option, +) -> Option { + let measured = measured?; + if measured.host_memory_observed { + return None; + } + let mib_to_bytes = |mib: f64| (mib * 1024.0 * 1024.0).round() as u64; + Some(MeasuredPlanSnapshot { + key: key.clone(), + footprint: MeasuredBufferFootprint { + compute_bytes: mib_to_bytes(measured.compute_mib?), + kv_bytes: mib_to_bytes(measured.kv_mib?), + context_length: breakdown.context_length, + lane_count: breakdown.slots as u32, + }, + }) +} + +/// Return a completed native buffer measurement only when its model identity, +/// capacity pool, and allocation-affecting configuration all match. +pub(super) fn measured_buffers_footprint( + key: &MemoryPlanMeasurementKey, +) -> Option { + let snapshot = MEASURED_PLAN_SNAPSHOT.lock().ok()?.clone()?; + if snapshot.key != *key || snapshot.footprint.context_length == 0 { + return None; + } + Some(snapshot.footprint) +} + +/// Emit the structured plan-time estimate while preserving the package-v2 +/// discriminator used by existing telemetry queries. +pub(super) fn emit_memory_plan_resolved( + model_name: &str, + breakdown: Option<&RuntimeResourcePlanBreakdown>, + start_path: MemoryPlanStartPath, +) { + let Some(breakdown) = breakdown else { + return; + }; + let slots_source = plan_value_source(breakdown.slots_auto); + let context_source = plan_value_source(breakdown.context_auto); + + macro_rules! emit { + ($($package_field:tt)*) => { + tracing::info!( + model = model_name, + $($package_field)* + memory_plan.vram_bytes = breakdown.vram_bytes, + memory_plan.model_bytes = breakdown.model_bytes, + memory_plan.kv_budget_bytes = breakdown.kv_budget_bytes, + memory_plan.planned_kv_bytes = breakdown.planned_kv_bytes, + memory_plan.kv_bytes_per_token = breakdown.kv_bytes_per_token, + memory_plan.compute_charge_bytes = breakdown.compute_charge_bytes, + memory_plan.planning_source = breakdown.planning_source.as_str(), + memory_plan.measured_fit = breakdown.measured_fit.unwrap_or(true), + memory_plan.measured_fit_available = breakdown.measured_fit.is_some(), + memory_plan.context_length = breakdown.context_length, + memory_plan.slots = breakdown.slots, + memory_plan.slots_source = slots_source, + memory_plan.context_source = context_source, + "memory plan resolved: charged estimates at plan time; compare with measured buffer_mib native events" + ) + }; + } + + match start_path { + MemoryPlanStartPath::Direct => emit!(), + MemoryPlanStartPath::PackageV2 => emit!(memory_plan.package = "v2",), + } +} + +fn plan_value_source(automatic: bool) -> &'static str { + if automatic { "auto" } else { "override" } +} + +/// Reconcile a resolved plan against the native buffers captured during open. +pub(super) fn emit_measured_memory_reconciliation( + model_name: &str, + measurement_key: &MemoryPlanMeasurementKey, + plan: &RuntimeResourcePlan, +) { + let Some(breakdown) = plan.breakdown.as_ref() else { + return; + }; + let measured = skippy_runtime::measured_native_buffers(); + let reconciliation = reconcile_memory_plan_with_measurements(breakdown, measured); + if let Ok(mut snapshot) = MEASURED_PLAN_SNAPSHOT.lock() { + *snapshot = completed_measurement_snapshot(measurement_key, breakdown, measured); + } + let memory_plan_measured = + measured.is_some_and(|m| m.compute_mib.is_some() || m.kv_mib.is_some()); + let measurement_reusable = measured.is_some_and(|measurement| { + !measurement.host_memory_observed + && measurement.compute_mib.is_some() + && measurement.kv_mib.is_some() + }); + tracing::info!( + model = model_name, + memory_plan.measured_available = memory_plan_measured, + memory_plan.charged_compute_reserve_bytes = reconciliation.charged_compute_reserve_bytes, + memory_plan.measured_compute_bytes = reconciliation.measured_compute_bytes.unwrap_or(0), + memory_plan.measured_kv_bytes = reconciliation.measured_kv_bytes.unwrap_or(0), + memory_plan.residual_free_bytes = reconciliation.residual_free_bytes.unwrap_or(0), + memory_plan.measured_residual_available = reconciliation.residual_free_bytes.is_some(), + memory_plan.measurement_reusable = measurement_reusable, + "memory plan reconciled with measured native buffers" + ); +} + +#[cfg(test)] +mod tests { + use super::{ + MEASURED_PLAN_SNAPSHOT, MeasuredPlanSnapshot, MemoryPlanMeasurementKey, + completed_measurement_snapshot, measured_buffers_footprint, + }; + use crate::runtime::context_planning::{ + MeasuredBufferFootprint, RuntimeResourcePlanBreakdown, RuntimeResourcePlanSource, + }; + + #[test] + fn measured_footprint_reads_one_coherent_plan_snapshot() { + let mut snapshot = MEASURED_PLAN_SNAPSHOT.lock().unwrap(); + *snapshot = None; + drop(snapshot); + let first_key = MemoryPlanMeasurementKey::new("first".to_string()); + let second_key = MemoryPlanMeasurementKey::new("second".to_string()); + assert!(measured_buffers_footprint(&first_key).is_none()); + + snapshot = MEASURED_PLAN_SNAPSHOT.lock().unwrap(); + *snapshot = Some(MeasuredPlanSnapshot { + key: second_key, + footprint: MeasuredBufferFootprint { + compute_bytes: 10, + kv_bytes: 20, + context_length: 32768, + lane_count: 4, + }, + }); + drop(snapshot); + assert!(measured_buffers_footprint(&first_key).is_none()); + + snapshot = MEASURED_PLAN_SNAPSHOT.lock().unwrap(); + *snapshot = Some(MeasuredPlanSnapshot { + key: first_key.clone(), + footprint: MeasuredBufferFootprint { + compute_bytes: 30, + kv_bytes: 40, + context_length: 8192, + lane_count: 2, + }, + }); + drop(snapshot); + let footprint = measured_buffers_footprint(&first_key).expect("matching snapshot"); + assert_eq!(footprint.compute_bytes, 30); + assert_eq!(footprint.kv_bytes, 40); + assert_eq!(footprint.context_length, 8192); + assert_eq!(footprint.lane_count, 2); + + snapshot = MEASURED_PLAN_SNAPSHOT.lock().unwrap(); + *snapshot = Some(MeasuredPlanSnapshot { + key: first_key.clone(), + footprint: MeasuredBufferFootprint { + compute_bytes: 30, + kv_bytes: 40, + context_length: 0, + lane_count: 4, + }, + }); + drop(snapshot); + assert!(measured_buffers_footprint(&first_key).is_none()); + + *MEASURED_PLAN_SNAPSHOT.lock().unwrap() = None; + } + + #[test] + fn completed_snapshot_rejects_host_offload_and_incomplete_measurements() { + let key = MemoryPlanMeasurementKey::new("config".to_string()); + let breakdown = RuntimeResourcePlanBreakdown { + vram_bytes: 10_000, + model_bytes: 2_000, + kv_budget_bytes: 6_000, + planned_kv_bytes: 4_000, + kv_bytes_per_token: 4, + compute_charge_bytes: 2_000, + planning_source: RuntimeResourcePlanSource::StaticEstimate, + measured_fit: None, + slots: 4, + context_length: 1_000, + slots_auto: true, + context_auto: true, + }; + let measured = skippy_runtime::MeasuredNativeBuffers { + compute_mib: Some(1.0), + kv_mib: Some(2.0), + host_memory_observed: false, + }; + assert!(completed_measurement_snapshot(&key, &breakdown, Some(measured)).is_some()); + assert!( + completed_measurement_snapshot( + &key, + &breakdown, + Some(skippy_runtime::MeasuredNativeBuffers { + host_memory_observed: true, + ..measured + }) + ) + .is_none() + ); + assert!( + completed_measurement_snapshot( + &key, + &breakdown, + Some(skippy_runtime::MeasuredNativeBuffers { + kv_mib: None, + ..measured + }) + ) + .is_none() + ); + } +} diff --git a/crates/mesh-llm-host-runtime/src/runtime/mod.rs b/crates/mesh-llm-host-runtime/src/runtime/mod.rs index 357e669610..25dd6589bd 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/mod.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/mod.rs @@ -10,7 +10,9 @@ mod discovery; pub mod instance; mod instance_lifecycle; mod interactive; +pub(crate) mod kv_disk_config; mod local; +mod local_memory_plan; mod local_model_only; mod local_package; mod local_split; diff --git a/crates/mesh-llm-host-runtime/src/runtime/options.rs b/crates/mesh-llm-host-runtime/src/runtime/options.rs index ffa9cc1bd7..3a74b45bb1 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/options.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/options.rs @@ -87,6 +87,9 @@ pub struct RuntimeOptions { pub nostr_relay: Vec, pub no_console: bool, pub config: Option, + pub kv_cache_disk: Option, + pub kv_cache_disk_dir: Option, + pub kv_cache_min_free: Option, pub owner_key: Option, pub control_bind: Option, pub control_advertise_addr: Option, @@ -166,6 +169,9 @@ impl Default for RuntimeOptions { nostr_relay: Vec::new(), no_console: false, config: None, + kv_cache_disk: None, + kv_cache_disk_dir: None, + kv_cache_min_free: None, owner_key: None, control_bind: None, control_advertise_addr: None, diff --git a/crates/mesh-llm-host-runtime/src/runtime/proxy/tests/mod.rs b/crates/mesh-llm-host-runtime/src/runtime/proxy/tests/mod.rs index b38afcf387..1212a16eab 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/proxy/tests/mod.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/proxy/tests/mod.rs @@ -378,7 +378,8 @@ async fn spawn_held_upstream( tokio::spawn(async move { let _raw = read_raw_http_request(&mut stream).await; accepted.fetch_add(1, std::sync::atomic::Ordering::SeqCst); - let _permit = release.acquire().await.expect("release semaphore"); + let permit = release.acquire().await.expect("release semaphore"); + permit.forget(); let reply = format!( "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", response.len(), diff --git a/crates/mesh-llm-host-runtime/src/runtime/run_auto.rs b/crates/mesh-llm-host-runtime/src/runtime/run_auto.rs index bcb46395f2..fd7cf8a206 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/run_auto.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/run_auto.rs @@ -14,10 +14,10 @@ use super::{ StartupReadyReporter, bridge_skippy_native_logs, build_serving_list, cli_has_explicit_models, configure_skippy_native_logging, emit_configuration_ui_read_only_hint, initialize_embedded_runtime_entrypoint, initialize_runtime_entrypoint, - maybe_discover_join_candidates, next_runtime_instance_id, nostr_rediscovery, nostr_relays, - openai_guardrail_policy_handle, owner_runtime_config, prepare_runtime_startup, - publish_initial_openai_guardrails_status, record_first_joined_mesh_ts, - record_runtime_operational_event, resolve_runtime_owner_key_path, + kv_disk_config::configure_node_kv_disk_cache, maybe_discover_join_candidates, + next_runtime_instance_id, nostr_rediscovery, nostr_relays, openai_guardrail_policy_handle, + owner_runtime_config, prepare_runtime_startup, publish_initial_openai_guardrails_status, + record_first_joined_mesh_ts, record_runtime_operational_event, resolve_runtime_owner_key_path, resolve_startup_mesh_creation_state, run_auto_join_mesh_phase, run_auto_model_identity, run_auto_model_path_or_shutdown, run_auto_runtime_loop_and_shutdown, run_local_model_only, runtime_data_producer_for_console, runtime_startup_requirements, setup_run_auto_console_state, @@ -308,6 +308,14 @@ pub(super) async fn run_runtime_cli( )?; apply_runtime_config_options(&mut options, &config); + let disk_cache = configure_node_kv_disk_cache(&config, &options)?; + for warning in &disk_cache.configured.warnings { + let _ = emit_event(OutputEvent::Warning { + message: warning.clone(), + context: None, + }); + } + initialize_audit_logging_for_options(&options)?; let startup_mesh_creation_state = resolve_startup_mesh_creation_state(&options, &config)?; diff --git a/crates/mesh-llm-protocol/proto/node.proto b/crates/mesh-llm-protocol/proto/node.proto index d15c77e653..09c0ab925d 100644 --- a/crates/mesh-llm-protocol/proto/node.proto +++ b/crates/mesh-llm-protocol/proto/node.proto @@ -541,6 +541,7 @@ message OwnerControlRequest { OwnerControlUnloadModelRequest unload_model = 7; OwnerControlEnsureModelRequest ensure_model = 8; OwnerControlDrainModelRequest drain_model = 9; + OwnerControlKvCacheRequest kv_cache = 10; } message OwnerControlResponse { @@ -553,6 +554,7 @@ message OwnerControlResponse { OwnerControlUnloadModelResponse unload_model = 7; OwnerControlEnsureModelResponse ensure_model = 8; OwnerControlDrainModelResponse drain_model = 9; + OwnerControlKvCacheResponse kv_cache = 10; } message OwnerControlError { @@ -654,6 +656,21 @@ message OwnerControlDrainModelRequest { optional uint64 drain_timeout_secs = 4; } +message OwnerControlKvCacheRequest { + bytes requester_node_id = 1; // exactly 32 bytes + bytes target_node_id = 2; // exactly 32 bytes + OwnerControlKvCacheOperation operation = 3; + optional uint64 target_bytes = 4; // prune only; default is 85% of budget + optional string model_identity = 5; // exact internal numerical identity +} + +enum OwnerControlKvCacheOperation { + OWNER_CONTROL_KV_CACHE_OPERATION_UNSPECIFIED = 0; + OWNER_CONTROL_KV_CACHE_OPERATION_STATUS = 1; + OWNER_CONTROL_KV_CACHE_OPERATION_PRUNE = 2; + OWNER_CONTROL_KV_CACHE_OPERATION_CLEAR = 3; +} + message OwnerControlLoadModelResponse { string intent_id = 1; string accepted_state = 2; @@ -678,6 +695,11 @@ message OwnerControlDrainModelResponse { OwnerControlModelRef target = 3; } +message OwnerControlKvCacheResponse { + bytes status_json = 1; // versioned KvCacheStatusPayload JSON + optional uint64 freed_bytes = 2; // present for prune and clear +} + message OwnerControlRefreshInventoryResponse { OwnerControlConfigSnapshot snapshot = 1; OwnerControlRefreshInventory inventory = 2; diff --git a/crates/mesh-llm-protocol/src/proto/node.rs b/crates/mesh-llm-protocol/src/proto/node.rs index 36b009cfeb..bd1be8e3cb 100644 --- a/crates/mesh-llm-protocol/src/proto/node.rs +++ b/crates/mesh-llm-protocol/src/proto/node.rs @@ -648,6 +648,8 @@ pub struct OwnerControlRequest { pub ensure_model: ::core::option::Option, #[prost(message, optional, tag = "9")] pub drain_model: ::core::option::Option, + #[prost(message, optional, tag = "10")] + pub kv_cache: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct OwnerControlResponse { @@ -669,6 +671,8 @@ pub struct OwnerControlResponse { pub ensure_model: ::core::option::Option, #[prost(message, optional, tag = "9")] pub drain_model: ::core::option::Option, + #[prost(message, optional, tag = "10")] + pub kv_cache: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct OwnerControlError { @@ -834,6 +838,23 @@ pub struct OwnerControlDrainModelRequest { pub drain_timeout_secs: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct OwnerControlKvCacheRequest { + /// exactly 32 bytes + #[prost(bytes = "vec", tag = "1")] + pub requester_node_id: ::prost::alloc::vec::Vec, + /// exactly 32 bytes + #[prost(bytes = "vec", tag = "2")] + pub target_node_id: ::prost::alloc::vec::Vec, + #[prost(enumeration = "OwnerControlKvCacheOperation", tag = "3")] + pub operation: i32, + /// prune only; default is 85% of budget + #[prost(uint64, optional, tag = "4")] + pub target_bytes: ::core::option::Option, + /// exact internal numerical identity + #[prost(string, optional, tag = "5")] + pub model_identity: ::core::option::Option<::prost::alloc::string::String>, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct OwnerControlLoadModelResponse { #[prost(string, tag = "1")] pub intent_id: ::prost::alloc::string::String, @@ -870,6 +891,15 @@ pub struct OwnerControlDrainModelResponse { pub target: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct OwnerControlKvCacheResponse { + /// versioned KvCacheStatusPayload JSON + #[prost(bytes = "vec", tag = "1")] + pub status_json: ::prost::alloc::vec::Vec, + /// present for prune and clear + #[prost(uint64, optional, tag = "2")] + pub freed_bytes: ::core::option::Option, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct OwnerControlRefreshInventory { #[prost(message, repeated, tag = "1")] pub entries: ::prost::alloc::vec::Vec, @@ -1414,6 +1444,35 @@ impl OwnerControlErrorCode { } #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] +pub enum OwnerControlKvCacheOperation { + Unspecified = 0, + Status = 1, + Prune = 2, + Clear = 3, +} +impl OwnerControlKvCacheOperation { + /// String value of the enum field names used in the ProtoBuf definition. + pub fn as_str_name(&self) -> &'static str { + match self { + Self::Unspecified => "OWNER_CONTROL_KV_CACHE_OPERATION_UNSPECIFIED", + Self::Status => "OWNER_CONTROL_KV_CACHE_OPERATION_STATUS", + Self::Prune => "OWNER_CONTROL_KV_CACHE_OPERATION_PRUNE", + Self::Clear => "OWNER_CONTROL_KV_CACHE_OPERATION_CLEAR", + } + } + + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "OWNER_CONTROL_KV_CACHE_OPERATION_UNSPECIFIED" => Some(Self::Unspecified), + "OWNER_CONTROL_KV_CACHE_OPERATION_STATUS" => Some(Self::Status), + "OWNER_CONTROL_KV_CACHE_OPERATION_PRUNE" => Some(Self::Prune), + "OWNER_CONTROL_KV_CACHE_OPERATION_CLEAR" => Some(Self::Clear), + _ => None, + } + } +} +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] pub enum OwnerControlRefreshInventoryDisposition { Unspecified = 0, Executed = 1, diff --git a/crates/mesh-llm-protocol/src/protocol/mod.rs b/crates/mesh-llm-protocol/src/protocol/mod.rs index 3ddcb420fe..b0c61e68af 100644 --- a/crates/mesh-llm-protocol/src/protocol/mod.rs +++ b/crates/mesh-llm-protocol/src/protocol/mod.rs @@ -71,6 +71,8 @@ pub enum ControlFrameError { MissingModelRef, InvalidModelRefCombination, InvalidInventoryOrder, + InvalidKvCacheOperation { got: i32 }, + MissingKvCacheStatus, DecodeError(String), WrongStreamType { expected: u8, got: u8 }, ForgedSender, @@ -169,6 +171,12 @@ impl std::fmt::Display for ControlFrameError { "inventory entries must be strictly sorted by canonical model ref" ) } + ControlFrameError::InvalidKvCacheOperation { got } => { + write!(f, "invalid owner-control kv-cache operation: {got}") + } + ControlFrameError::MissingKvCacheStatus => { + write!(f, "owner-control kv-cache response missing status payload") + } ControlFrameError::DecodeError(msg) => write!(f, "protobuf decode error: {}", msg), ControlFrameError::WrongStreamType { expected, got } => write!( f, @@ -363,6 +371,7 @@ impl ValidateControlFrame for crate::proto::node::OwnerControlRequest { self.unload_model.is_some(), self.ensure_model.is_some(), self.drain_model.is_some(), + self.kv_cache.is_some(), ]; if commands.into_iter().filter(|present| *present).count() != 1 { return Err(ControlFrameError::MissingControlCommand); @@ -391,6 +400,9 @@ impl ValidateControlFrame for crate::proto::node::OwnerControlRequest { if let Some(request) = &self.drain_model { request.validate_frame()?; } + if let Some(request) = &self.kv_cache { + request.validate_frame()?; + } Ok(()) } } @@ -409,6 +421,7 @@ impl ValidateControlFrame for crate::proto::node::OwnerControlResponse { self.unload_model.is_some(), self.ensure_model.is_some(), self.drain_model.is_some(), + self.kv_cache.is_some(), ]; if results.into_iter().filter(|present| *present).count() != 1 { return Err(ControlFrameError::MissingControlResult); @@ -437,6 +450,9 @@ impl ValidateControlFrame for crate::proto::node::OwnerControlResponse { if let Some(response) = &self.drain_model { response.validate_frame()?; } + if let Some(response) = &self.kv_cache { + response.validate_frame()?; + } Ok(()) } } @@ -597,6 +613,23 @@ impl ValidateControlFrame for crate::proto::node::OwnerControlDrainModelRequest } } +impl ValidateControlFrame for crate::proto::node::OwnerControlKvCacheRequest { + fn validate_frame(&self) -> Result<(), ControlFrameError> { + validate_endpoint_id_length(self.requester_node_id.len())?; + validate_endpoint_id_length(self.target_node_id.len())?; + match crate::proto::node::OwnerControlKvCacheOperation::try_from(self.operation) { + Ok(crate::proto::node::OwnerControlKvCacheOperation::Status) + | Ok(crate::proto::node::OwnerControlKvCacheOperation::Prune) + | Ok(crate::proto::node::OwnerControlKvCacheOperation::Clear) => Ok(()), + Ok(crate::proto::node::OwnerControlKvCacheOperation::Unspecified) | Err(_) => { + Err(ControlFrameError::InvalidKvCacheOperation { + got: self.operation, + }) + } + } + } +} + impl ValidateControlFrame for crate::proto::node::OwnerControlLoadModelResponse { fn validate_frame(&self) -> Result<(), ControlFrameError> { validate_owner_control_model_for_load_or_ensure( @@ -637,6 +670,15 @@ impl ValidateControlFrame for crate::proto::node::OwnerControlDrainModelResponse } } +impl ValidateControlFrame for crate::proto::node::OwnerControlKvCacheResponse { + fn validate_frame(&self) -> Result<(), ControlFrameError> { + if self.status_json.is_empty() { + return Err(ControlFrameError::MissingKvCacheStatus); + } + Ok(()) + } +} + impl ValidateControlFrame for crate::proto::node::OwnerControlRefreshInventory { fn validate_frame(&self) -> Result<(), ControlFrameError> { use crate::proto::node::OwnerControlRefreshInventoryDisposition; @@ -917,6 +959,7 @@ mod tests { OwnerControlConfigSnapshot, OwnerControlConfigUpdate, OwnerControlEnvelope, OwnerControlError, OwnerControlErrorCode, OwnerControlGetConfigRequest, OwnerControlGetConfigResponse, OwnerControlHandshake, OwnerControlInventoryEntry, + OwnerControlKvCacheOperation, OwnerControlKvCacheRequest, OwnerControlKvCacheResponse, OwnerControlRefreshInventory, OwnerControlRefreshInventoryDisposition, OwnerControlRefreshInventoryRequest, OwnerControlRefreshInventoryResponse, OwnerControlRequest, OwnerControlResponse, OwnerControlWatchAccepted, @@ -1012,6 +1055,7 @@ mod tests { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), response: None, error: None, @@ -1040,6 +1084,7 @@ mod tests { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), error: None, }; @@ -1064,6 +1109,7 @@ mod tests { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), response: None, error: None, @@ -1092,6 +1138,7 @@ mod tests { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), error: None, }; @@ -1114,6 +1161,7 @@ mod tests { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), response: None, error: None, @@ -1138,6 +1186,7 @@ mod tests { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), error: None, }; @@ -1160,6 +1209,7 @@ mod tests { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), error: None, }; @@ -1189,11 +1239,50 @@ mod tests { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), error: None, }; decode_owner_control_envelope(&encode_owner_control_envelope(&update_response)) .expect("watch update response must decode"); + + let cache_request = OwnerControlEnvelope { + r#gen: NODE_PROTOCOL_GENERATION, + handshake: None, + request: Some(OwnerControlRequest { + request_id: 16, + kv_cache: Some(OwnerControlKvCacheRequest { + requester_node_id: vec![0x61; 32], + target_node_id: vec![0x62; 32], + operation: OwnerControlKvCacheOperation::Prune as i32, + target_bytes: Some(1024), + model_identity: Some("blake3:model".to_string()), + }), + ..Default::default() + }), + response: None, + error: None, + }; + let decoded = decode_owner_control_envelope(&encode_owner_control_envelope(&cache_request)) + .expect("kv-cache request must decode"); + assert_eq!(decoded.request.unwrap().request_id, 16); + + let cache_response = OwnerControlEnvelope { + r#gen: NODE_PROTOCOL_GENERATION, + handshake: None, + request: None, + response: Some(OwnerControlResponse { + request_id: 16, + kv_cache: Some(OwnerControlKvCacheResponse { + status_json: br#"{"version":1}"#.to_vec(), + freed_bytes: Some(512), + }), + ..Default::default() + }), + error: None, + }; + decode_owner_control_envelope(&encode_owner_control_envelope(&cache_response)) + .expect("kv-cache response must decode"); } #[test] @@ -1211,6 +1300,7 @@ mod tests { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), response: None, error: None, diff --git a/crates/mesh-llm-routing/src/lib.rs b/crates/mesh-llm-routing/src/lib.rs index cc3c340d1f..8700736eea 100644 --- a/crates/mesh-llm-routing/src/lib.rs +++ b/crates/mesh-llm-routing/src/lib.rs @@ -26,7 +26,112 @@ pub fn total_model_bytes(model: &Path) -> u64 { return total; } } - std::fs::metadata(model).map(|m| m.len()).unwrap_or(0) + match std::fs::metadata(model) { + Ok(metadata) if metadata.is_dir() => { + // A SafeTensors checkpoint directory (model.safetensors + + // tokenizer/config siblings) reports the directory inode's size + // (~4 KiB) as its length. Sum the checkpoint files instead so + // memory planning charges the real weight bytes; a directory of + // plain GGUFs is not a loadable single model, but summing its + // files is still the least-wrong size estimate for routing. + dir_file_bytes(model) + } + Ok(metadata) => metadata.len(), + Err(_) => 0, + } +} + +/// Sum the regular-file sizes directly inside `dir` (non-recursive; model +/// checkpoint directories are flat). +fn dir_file_bytes(dir: &Path) -> u64 { + std::fs::read_dir(dir) + .map(|entries| { + entries + .filter_map(|entry| entry.ok()) + // Follow symlinks: HuggingFace-style snapshot entries are + // symlinks into a blobs/ store, and DirEntry::metadata is + // lstat — it would report the link itself (is_file() false, + // len 0), zeroing the total. fs::metadata follows the link + // and reports the target file. + .filter_map(|entry| std::fs::metadata(entry.path()).ok()) + .filter(|metadata| metadata.is_file()) + .map(|metadata| metadata.len()) + .sum() + }) + .unwrap_or(0) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn directory_model_reports_summed_file_bytes() { + // A SafeTensors checkpoint dir must report the summed weight-file + // bytes, not the directory inode's ~4 KiB st_size — memory planning + // charges this quantity against the KV budget. + let dir = + std::env::temp_dir().join(format!("mesh-routing-total-bytes-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + std::fs::write(dir.join("model.safetensors"), vec![0u8; 4096]).unwrap(); + std::fs::write(dir.join("config.json"), vec![0u8; 128]).unwrap(); + std::fs::create_dir(dir.join("nested")).unwrap(); + std::fs::write(dir.join("nested").join("ignored.bin"), vec![0u8; 999_999]).unwrap(); + let total = total_model_bytes(&dir); + let _ = std::fs::remove_dir_all(&dir); + assert_eq!(total, 4096 + 128); + } + + #[test] + fn file_model_reports_its_own_bytes() { + let file = std::env::temp_dir().join(format!( + "mesh-routing-total-bytes-file-{}", + std::process::id() + )); + std::fs::write(&file, vec![0u8; 2048]).unwrap(); + let total = total_model_bytes(&file); + let _ = std::fs::remove_file(&file); + assert_eq!(total, 2048); + } + + #[test] + fn symlinked_snapshot_entries_report_target_bytes() { + // HuggingFace-style snapshot dirs symlink weight files into a blobs/ + // store. DirEntry::metadata is lstat: it sees the link itself + // (is_file() false, len 0) and would total zero — worse than the + // directory-inode bug this replaced. The reader must follow links. + let root = std::env::temp_dir().join(format!( + "mesh-routing-total-bytes-symlink-{}", + std::process::id() + )); + let blobs = root.join("blobs"); + let snapshot = root.join("snapshots").join("abc123"); + std::fs::create_dir_all(&blobs).unwrap(); + std::fs::create_dir_all(&snapshot).unwrap(); + std::fs::write(blobs.join("blob-weights"), vec![0u8; 8192]).unwrap(); + std::fs::write(blobs.join("blob-config"), vec![0u8; 256]).unwrap(); + #[cfg(unix)] + std::os::unix::fs::symlink( + "../../blobs/blob-weights", + snapshot.join("model.safetensors"), + ) + .unwrap(); + #[cfg(unix)] + std::os::unix::fs::symlink("../../blobs/blob-config", snapshot.join("config.json")) + .unwrap(); + let total = total_model_bytes(&snapshot); + let _ = std::fs::remove_dir_all(&root); + #[cfg(unix)] + assert_eq!(total, 8192 + 256); + } + + #[test] + fn missing_model_reports_zero() { + assert_eq!( + total_model_bytes(Path::new("/nonexistent/mesh-routing-missing-model")), + 0 + ); + } } /// The current inference target selected by runtime planning. diff --git a/crates/mesh-llm-system/src/hardware/tests.rs b/crates/mesh-llm-system/src/hardware/tests.rs index 42bd454454..d2ebae7c6d 100644 --- a/crates/mesh-llm-system/src/hardware/tests.rs +++ b/crates/mesh-llm-system/src/hardware/tests.rs @@ -1,6 +1,55 @@ use super::*; use serial_test::serial; +#[cfg(all( + target_os = "linux", + any( + not(feature = "skippy-devices"), + feature = "dynamic-native-runtime", + test + ) +))] +struct TestDirectory(std::path::PathBuf); + +#[cfg(all( + target_os = "linux", + any( + not(feature = "skippy-devices"), + feature = "dynamic-native-runtime", + test + ) +))] +impl TestDirectory { + fn new(label: &str) -> Self { + static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + let id = NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .expect("test clock follows Unix epoch") + .as_nanos(); + let path = std::env::temp_dir().join(format!( + "mesh-llm-{label}-{}-{timestamp}-{id}", + std::process::id() + )); + std::fs::create_dir(&path).expect("create collision-resistant test directory"); + Self(path) + } +} + +#[cfg(all( + target_os = "linux", + any( + not(feature = "skippy-devices"), + feature = "dynamic-native-runtime", + test + ) +))] +impl Drop for TestDirectory { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.0); + } +} + fn synthetic_gpu(index: usize, stable_id: Option<&str>) -> GpuFacts { GpuFacts { index, @@ -414,11 +463,12 @@ fn test_tegra_collector_gpu_name_absent_leaves_source_none() { // model file present, which made the old `TegraCollector.collect` form flip // on such a host). With the model file absent, both the name and its source // must stay absent — never a guessed source for a name that was never read. - let missing = std::path::Path::new("/nonexistent/mesh-llm/tegra/devicetree/base/model"); + let directory = TestDirectory::new("tegra-model-absent"); + let missing = directory.0.join("missing-model"); assert!(!missing.exists()); let mut survey = HardwareSurvey::default(); - tegra_gpu_name_from_model_path(&mut survey, missing); + tegra_gpu_name_from_model_path(&mut survey, &missing); assert_eq!(survey.gpu_name, None); assert_eq!(survey.gpu_name_source, None); @@ -439,7 +489,8 @@ fn test_tegra_collector_gpu_name_absent_leaves_source_none() { fn test_tegra_collector_gpu_name_present_tags_sysfs_source() { use std::io::Write as _; - let path = std::env::temp_dir().join("mesh_llm_test_tegra_model_present"); + let directory = TestDirectory::new("tegra-model-present"); + let path = directory.0.join("model"); let mut f = std::fs::File::create(&path).expect("create temp model file"); write!(f, "NVIDIA Jetson AGX Orin Developer Kit\0").expect("write model file"); drop(f); @@ -447,8 +498,6 @@ fn test_tegra_collector_gpu_name_present_tags_sysfs_source() { let mut survey = HardwareSurvey::default(); tegra_gpu_name_from_model_path(&mut survey, &path); - let _ = std::fs::remove_file(&path); - assert_eq!(survey.gpu_name.as_deref(), Some("Jetson AGX Orin")); assert_eq!(survey.gpu_name_source, Some(GpuNameSource::Sysfs)); } diff --git a/crates/mesh-llm-ui/src/features/app-tabs/configuration-defaults-runtime.ts b/crates/mesh-llm-ui/src/features/app-tabs/configuration-defaults-runtime.ts index b4f1bfdac1..0dfe146543 100644 --- a/crates/mesh-llm-ui/src/features/app-tabs/configuration-defaults-runtime.ts +++ b/crates/mesh-llm-ui/src/features/app-tabs/configuration-defaults-runtime.ts @@ -411,7 +411,8 @@ export const CONFIGURATION_DEFAULT_RUNTIME_SETTINGS = [ categoryId: 'memory', icon: 'layers', label: 'Micro-batch size', - description: 'Set the default decode micro-batch size.', + description: + 'Set the default micro-batch (physical prefill chunk) size. Values at or below 128 keep the CUDA SSM sequential-scan fallback; larger values enable the SSD chunked kernel for recurrent models.', inheritedLabel: 'Applied when a placement does not override micro-batch size', visibility: 'advanced', tomlSection: MODEL_FIT_TOML_SECTION, diff --git a/crates/mesh-llm/src/commands/mod.rs b/crates/mesh-llm/src/commands/mod.rs index bc85b8a161..055c9a9968 100644 --- a/crates/mesh-llm/src/commands/mod.rs +++ b/crates/mesh-llm/src/commands/mod.rs @@ -43,6 +43,13 @@ pub async fn dispatch(cli: &Cli) -> Result { async fn dispatch_command(cli: &Cli, cmd: &Command) -> Result<()> { match cmd { Command::Auth { command } => mesh_llm_commands::auth::run_auth_command(command), + Command::KvCache { command } => { + mesh_llm_commands::kv_cache::dispatch_kv_cache_command(command).await + } + Command::Runtime { command } => { + dispatch_runtime_command(command.as_ref(), cli.config.as_deref(), cli.llama_flavor) + .await + } Command::ModelPrepare { .. } => dispatch_model_prepare(cmd).await, _ => dispatch_general_command(cli, cmd).await, } @@ -67,9 +74,8 @@ async fn dispatch_general_command(cli: &Cli, cmd: &Command) -> Result<()> { )?; Ok(()) } - Command::Runtime { command } => { - dispatch_runtime_command(command.as_ref(), cli.config.as_deref(), cli.llama_flavor) - .await + Command::Runtime { .. } | Command::KvCache { .. } => { + unreachable!("runtime and kv-cache commands are dispatched before general commands") } Command::Setup { .. } => { dispatch_setup_command(cmd, cli.config.as_deref(), cli.llama_flavor).await diff --git a/crates/mesh-llm/src/lib.rs b/crates/mesh-llm/src/lib.rs index 3e727b8a0b..b8a1e0f3e4 100644 --- a/crates/mesh-llm/src/lib.rs +++ b/crates/mesh-llm/src/lib.rs @@ -479,6 +479,9 @@ fn runtime_options_from_cli(cli: mesh_llm_cli::Cli) -> mesh_llm_host_runtime::Ru nostr_relay: cli.nostr_relay, no_console: cli.no_console, config: cli.config, + kv_cache_disk: cli.kv_cache_disk, + kv_cache_disk_dir: cli.kv_cache_disk_dir, + kv_cache_min_free: cli.kv_cache_min_free, owner_key: cli.owner_key, control_bind: cli.control_bind, control_advertise_addr: cli.control_advertise_addr, diff --git a/crates/mesh-llm/tests/protocol_convert_matrix.rs b/crates/mesh-llm/tests/protocol_convert_matrix.rs index b557446946..890d78d9d7 100644 --- a/crates/mesh-llm/tests/protocol_convert_matrix.rs +++ b/crates/mesh-llm/tests/protocol_convert_matrix.rs @@ -277,6 +277,7 @@ fn owner_control_get_config_roundtrip_works() { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), response: None, error: None, @@ -301,6 +302,7 @@ fn owner_control_unknown_command_maps_to_structured_error() { unload_model: None, ensure_model: None, drain_model: None, + kv_cache: None, }), response: None, error: None, diff --git a/crates/skippy-bench/Cargo.toml b/crates/skippy-bench/Cargo.toml index 8604d172aa..d845ecf5f5 100644 --- a/crates/skippy-bench/Cargo.toml +++ b/crates/skippy-bench/Cargo.toml @@ -10,6 +10,7 @@ clap.workspace = true csv = "1.4.0" dirs = "6.0.0" libc = "0.2" +skippy-cache = { path = "../skippy-cache" } skippy-protocol = { path = "../skippy-protocol" } skippy-runtime = { path = "../skippy-runtime" } skippy-topology = { path = "../skippy-topology" } diff --git a/crates/skippy-bench/src/cli.rs b/crates/skippy-bench/src/cli.rs index 8b68c33065..9ebca354ba 100644 --- a/crates/skippy-bench/src/cli.rs +++ b/crates/skippy-bench/src/cli.rs @@ -35,6 +35,8 @@ pub enum CommandKind { LocalSplitChainBinary(LocalSplitChainBinaryArgs), #[command(name = "verify-window-local")] VerifyWindowLocal(VerifyWindowLocalArgs), + #[command(name = "l2-tier")] + L2Tier(L2TierArgs), #[command(name = "chat-corpus")] ChatCorpus(ChatCorpusArgs), #[command(name = "token-lengths")] @@ -45,6 +47,35 @@ pub enum CommandKind { Run(RunArgs), } +#[derive(Parser)] +pub struct L2TierArgs { + /// Working directory for the L3 store. Must not exist: the bench + /// creates, sentinel-marks, and (unless `--keep-store`) removes a + /// directory it owns, and refuses any pre-existing path instead of + /// deleting it. + #[arg(long, default_value = "/tmp/skippy-l2-tier-bench")] + pub store_root: PathBuf, + /// Number of timed L3-cold / L2-warm matched pairs after warmup. + #[arg(long, default_value_t = 50)] + pub pairs: usize, + /// Recorded prefix length in tokens (the synthetic conversation length). + #[arg(long, default_value_t = 1_893)] + pub tokens: usize, + /// Bytes of KV payload per token — sized to mimic a real dense model's + /// per-token KV footprint at the target dtype. + #[arg(long, default_value_t = 512)] + pub kv_bytes_per_token: usize, + /// L2 budget in MiB. Defaults to four times one entry. + #[arg(long)] + pub l2_budget_mib: Option, + /// Keep the L3 store directory after the run for inspection. + #[arg(long, default_value_t = false)] + pub keep_store: bool, + /// Model identity stamped into the L3 tier and L2 keys. + #[arg(long, default_value = "bench-model")] + pub model_identity: String, +} + #[derive(Parser)] pub struct EvalArgs { #[command(subcommand)] diff --git a/crates/skippy-bench/src/l2_tier.rs b/crates/skippy-bench/src/l2_tier.rs new file mode 100644 index 0000000000..4a7b6bac04 --- /dev/null +++ b/crates/skippy-bench/src/l2_tier.rs @@ -0,0 +1,309 @@ +//! `l2-tier` benchmark: cold L3 fill versus warm L2 lookup on identical +//! packed entries (#1651). +//! +//! Builds a temporary L3 store, spills a synthetic multi-turn prompt at a +//! recorded prefix length, then measures two restore paths in-process: +//! +//! - **L3 cold fill**: `L3Tier::fill_longest` — index probe + segment +//! assembly + digest verification from disk. +//! - **L2 warm lookup**: `L2Tier::get` + `to_payload` + materialization — +//! the entry was admitted from an identical verified L3 fill, so the +//! lookup is a digest-keyed handle assembly (no re-hash: admission +//! verified the wire once; segments are immutable afterward). +//! +//! Both arms are timed through the same boundary: the moment their bytes +//! are usable (`full_state_bytes_timed`). A multi-segment L2 entry still +//! materializes its wire on read, so stopping the L2 timer at handle +//! creation would understate the real cost; the equality gate compares the +//! materialized bytes of both arms before the pair is counted. The +//! handle-only lookup time is reported separately as +//! `l2_handle_lookup_ns`. +//! +//! Admission hashing (the one-time wire BLAKE3) is measured separately +//! and reported as its own metric, never inside the timed lookup. +//! +//! The store root is owned by the run: it must not exist beforehand (the +//! bench refuses existing paths instead of deleting user data) and it is +//! marked with an ownership sentinel so cleanup never touches a directory +//! the bench did not create. +use std::path::{Path, PathBuf}; +use std::time::Instant; + +use anyhow::{Context, Result}; + +use crate::cli::L2TierArgs; +use skippy_cache::{ + ExactStatePayload, ExactStatePayloadMirror, L2Origin, L2Tier, l2_cache_key, l3_prefix_key, +}; + +/// Marker file proving the bench created the store root itself; cleanup +/// refuses to `remove_dir_all` a directory without it. +const OWNERSHIP_SENTINEL: &str = ".skippy-l2-tier-bench-owned"; + +/// Create a fresh, bench-owned store root. Existing paths are refused — +/// the bench must never delete a user-supplied directory it did not +/// create. +fn prepare_store_root(requested: &Path) -> Result { + if requested.symlink_metadata().is_ok() { + anyhow::bail!( + "refusing to use store root {}: the path already exists; the bench only runs \ + in a root it created itself", + requested.display() + ); + } + std::fs::create_dir_all(requested) + .with_context(|| format!("failed to create bench store root {}", requested.display()))?; + std::fs::write( + requested.join(OWNERSHIP_SENTINEL), + b"skippy-bench l2-tier store\n", + ) + .with_context(|| format!("failed to mark {} as bench-owned", requested.display()))?; + Ok(requested.to_path_buf()) +} + +/// Remove a bench-owned store root. Refuses paths without the ownership +/// sentinel so `remove_dir_all` can never hit arbitrary input. +fn remove_owned_store_root(root: &Path) -> Result<()> { + if !root.join(OWNERSHIP_SENTINEL).is_file() { + anyhow::bail!( + "refusing to remove store root {}: missing bench ownership sentinel", + root.display() + ); + } + std::fs::remove_dir_all(root) + .with_context(|| format!("failed to remove bench store root {}", root.display())) +} + +fn percentile(samples_ns: &mut [u128], pct: f64) -> f64 { + samples_ns.sort_unstable(); + let index = ((pct / 100.0) * (samples_ns.len() as f64 - 1.0)).round() as usize; + samples_ns[index.min(samples_ns.len() - 1)] as f64 +} + +pub fn l2_tier(args: L2TierArgs) -> Result<()> { + // Validation: reject degenerate configurations up front instead of + // dividing by zero or allocating nothing below. + if args.pairs == 0 { + anyhow::bail!("--pairs must be at least 1"); + } + if args.tokens == 0 { + anyhow::bail!("--tokens must be at least 1"); + } + if args.kv_bytes_per_token == 0 { + anyhow::bail!("--kv-bytes-per-token must be at least 1"); + } + + // Owned store root: refuse existing paths rather than deleting them, + // and mark the created directory so cleanup stays bounded to it. + let store_root = prepare_store_root(&args.store_root)?; + + let namespace = "bench-namespace"; + let state_identity = args.model_identity.clone(); + let token_ids: Vec = (0..args.tokens).map(|i| (i % 128_000) as i32).collect(); + + // Deterministic synthetic KV payload: content matters only for digests, + // size matters for timing. + let payload_len = args + .tokens + .checked_mul(args.kv_bytes_per_token) + .context("--tokens * --kv-bytes-per-token overflows")?; + let payload_bytes: Vec = (0..payload_len).map(|i| (i % 251) as u8).collect(); + let payload = ExactStatePayload::full_state(payload_bytes); + + let tier = skippy_cache::L3Tier::open( + store_root.clone(), + (payload_len as u64) * 8, + state_identity.clone(), + 64 * 1024, + ) + .context("failed to open bench L3 tier")?; + + // Spill once: this is the population path, not the measured path. + let manifest_key = tier + .spill(namespace, &token_ids, &payload, None, None) + .context("bench spill failed")?; + let _ = manifest_key; + + // Locate once to learn the recorded prefix key/digest used by both paths. + let location = tier + .locate_longest(namespace, &token_ids, 8) + .context("bench locate failed")? + .context("bench spill was not locatable")?; + let manifest = tier.store().load_manifest(&location.manifest_key)?; + let payload_digest = manifest.payload_digest.clone(); + let recorded_tokens = manifest.token_count; + + let l2_budget_bytes = args + .l2_budget_mib + .map(|mib| { + mib.checked_mul(1024 * 1024) + .context("--l2-budget-mib overflows") + }) + .transpose()? + .unwrap_or(payload_len as u64 * 4); + let l2 = L2Tier::new(l2_budget_bytes); + + let cache_key = l2_cache_key(&args.model_identity, &state_identity, namespace, &token_ids); + + // Warmup: one L3 fill, then admit it into L2 from the verified wire. + // The wire check inside `admit` is the one-time admission hash; it is + // timed separately below. + let warm_fill = tier + .fill_longest(namespace, &token_ids, 8) + .context("bench warmup L3 fill failed")? + .context("bench warmup L3 fill missed")?; + let (warm_wire, _) = warm_fill.payload.full_state_bytes_timed().context("wire")?; + let admission_started = Instant::now(); + l2.admit( + cache_key.clone(), + warm_fill.token_count, + payload_digest.clone(), + &warm_wire, + ExactStatePayloadMirror::from_manifest(&manifest) + .map_err(|refusal| anyhow::anyhow!(refusal.reason()))?, + L2Origin::FromL3, + ) + .map_err(|refusal| anyhow::anyhow!("bench warmup L2 admit refused: {}", refusal.reason()))?; + let admission_hash_ns = admission_started.elapsed().as_nanos(); + + let mut l3_samples: Vec = Vec::with_capacity(args.pairs); + let mut l2_samples: Vec = Vec::with_capacity(args.pairs); + let mut l2_handle_samples: Vec = Vec::with_capacity(args.pairs); + + for pair in 0..args.pairs { + // Cold-ish L3 fill: the OS page cache will help after warmup, which + // matches the production comparison — both paths run on the same + // machine state, the delta is the tier delta. + let start = Instant::now(); + let fill = tier + .fill_longest(namespace, &token_ids, 8) + .context("bench L3 fill failed")? + .context("bench L3 fill missed")?; + let l3_ns = start.elapsed().as_nanos(); + + // Timed through the same boundary as the L3 arm, starting before + // the lookup: the L3 timer covers index probe + assembly + + // verification, so the L2 timer covers lookup + handle assembly + + // materialization — both arms measure "nothing to usable bytes". + // The handle-only lookup time (this same `get`, inner timer) is + // reported separately as `l2_handle_lookup_ns`. + let start = Instant::now(); + let handle_start = Instant::now(); + let hit = l2.get(&cache_key); + let l2_handle_ns = handle_start.elapsed().as_nanos(); + let l2_payload = hit.as_ref().map(|hit| hit.to_payload()); + let (l2_bytes, _) = l2_payload + .as_ref() + .context("bench L2 payload missing")? + .full_state_bytes_timed() + .context("bench L2 bytes")?; + let l2_ns = start.elapsed().as_nanos(); + + // Correctness gate, outside the timer: L2 must return + // byte-identical state to the L3 fill, or the speedup is + // meaningless. + let hit = hit.context("bench L2 lookup missed")?; + anyhow::ensure!(hit.token_count == fill.token_count); + anyhow::ensure!(hit.payload_digest == payload_digest); + let (l3_bytes, _) = fill.payload.full_state_bytes_timed().context("l3 bytes")?; + anyhow::ensure!( + l3_bytes.as_ref() == l2_bytes.as_ref(), + "pair {pair}: L2 payload diverged from L3 fill" + ); + + l3_samples.push(l3_ns); + l2_samples.push(l2_ns); + l2_handle_samples.push(l2_handle_ns); + } + + let stats = l2.stats(); + let mut l3_sorted = l3_samples.clone(); + let mut l2_sorted = l2_samples.clone(); + let mut l2_handle_sorted = l2_handle_samples.clone(); + let summary = serde_json::json!({ + "bench": "l2-tier", + "pairs": args.pairs, + "tokens": recorded_tokens, + "payload_bytes": payload_len, + "l2_budget_bytes": l2_budget_bytes, + "model_identity": args.model_identity, + "l3_fill_ns": { + "p50": percentile(&mut l3_sorted, 50.0), + "p99": percentile(&mut l3_sorted, 99.0), + }, + "l2_lookup_to_usable_bytes_ns": { + "p50": percentile(&mut l2_sorted, 50.0), + "p99": percentile(&mut l2_sorted, 99.0), + }, + "l2_handle_lookup_ns": { + "p50": percentile(&mut l2_handle_sorted, 50.0), + "p99": percentile(&mut l2_handle_sorted, 99.0), + }, + "l2_admission_hash_ns_one_time": admission_hash_ns, + "speedup_p50": percentile(&mut l3_sorted, 50.0) / percentile(&mut l2_sorted, 50.0).max(1.0), + "l2_stats": { + "hits": stats.hits, + "misses": stats.misses, + "evictions": stats.evictions, + "bytes": stats.bytes, + "segments": stats.segments, + "shared_bytes_admitted": stats.shared_bytes_admitted, + }, + "l3_prefix_key": l3_prefix_key(namespace, &token_ids), + }); + println!("{summary}"); + + if !args.keep_store { + remove_owned_store_root(&store_root)?; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn store_root_refuses_existing_paths() { + let dir = + std::env::temp_dir().join(format!("skippy-l2-bench-refuse-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).expect("create pre-existing dir"); + std::fs::write(dir.join("precious.txt"), b"user data").expect("seed user data"); + + let err = prepare_store_root(&dir).expect_err("existing path must be refused"); + assert!( + err.to_string().contains("refusing to use store root"), + "unexpected error: {err}" + ); + assert!( + dir.join("precious.txt").is_file(), + "pre-existing contents must survive the refusal" + ); + + // A root the bench created is removable; a lookalike without the + // sentinel is not. (Creation itself stays refused for any + // pre-existing path, owned or not.) + let _ = std::fs::remove_dir_all(&dir); + + let owned = + std::env::temp_dir().join(format!("skippy-l2-bench-owned-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&owned); + prepare_store_root(&owned).expect("fresh root is created and owned"); + assert!(owned.join(OWNERSHIP_SENTINEL).is_file()); + remove_owned_store_root(&owned).expect("owned root is removable"); + assert!(!owned.exists()); + + // Unmarked directory: cleanup must refuse. + let unowned = + std::env::temp_dir().join(format!("skippy-l2-bench-unowned-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&unowned); + std::fs::create_dir_all(&unowned).expect("create unowned dir"); + std::fs::write(unowned.join("keep.txt"), b"user data").expect("seed"); + let err = + remove_owned_store_root(&unowned).expect_err("unowned root removal must be refused"); + assert!(err.to_string().contains("sentinel"), "unexpected: {err}"); + assert!(unowned.join("keep.txt").is_file(), "contents survive"); + let _ = std::fs::remove_dir_all(&unowned); + } +} diff --git a/crates/skippy-bench/src/main.rs b/crates/skippy-bench/src/main.rs index 1f0de777b0..39743c8bd7 100644 --- a/crates/skippy-bench/src/main.rs +++ b/crates/skippy-bench/src/main.rs @@ -3,6 +3,7 @@ mod cli; mod direct_return_listener; mod distributed; mod evals; +mod l2_tier; mod local_single; mod local_split; mod model_identity; @@ -54,6 +55,7 @@ fn main() -> Result<()> { CommandKind::LocalSplitCompare(args) => local_split_compare(args), CommandKind::LocalSplitChainBinary(args) => local_split_chain_binary(args), CommandKind::VerifyWindowLocal(args) => verify_window_local(args), + CommandKind::L2Tier(args) => l2_tier::l2_tier(args), CommandKind::ChatCorpus(args) => chat_corpus(args), CommandKind::TokenLengths(args) => token_lengths(args), CommandKind::FocusedRuntime(args) => focused_runtime(args), diff --git a/crates/skippy-cache/Cargo.toml b/crates/skippy-cache/Cargo.toml index aba7e04908..8664f1f5ef 100644 --- a/crates/skippy-cache/Cargo.toml +++ b/crates/skippy-cache/Cargo.toml @@ -14,4 +14,28 @@ path = "src/lib.rs" [dependencies] anyhow.workspace = true blake3.workspace = true +cubecl = { version = "=0.10.0", optional = true } +fs2 = "0.4" +libc = "0.2.183" +serde.workspace = true +serde_json.workspace = true skippy-protocol = { path = "../skippy-protocol", version = "0.76.1" } + +[target.'cfg(windows)'.dependencies] +windows-sys = { version = "0.61", features = [ + "Win32_Foundation", + "Win32_Security", + "Win32_Security_Authorization", + "Win32_Storage_FileSystem", + "Win32_System_Threading", + "Win32_System_WindowsProgramming", +] } + +[features] +# Evidence-only CubeCL spike (#1652): nothing in the library links CubeCL +# unless this feature is on. The example is excluded from default builds. +cachegen-spike = ["dep:cubecl", "cubecl/cpu", "cubecl/wgpu"] + +[[example]] +name = "cachegen_cubecl_spike" +required-features = ["cachegen-spike"] diff --git a/crates/skippy-cache/README.md b/crates/skippy-cache/README.md index 360b2bf756..0a5cfc9c64 100644 --- a/crates/skippy-cache/README.md +++ b/crates/skippy-cache/README.md @@ -122,6 +122,16 @@ continuation state for those models. That is where the largest wins come from: llama-server warm slots still have to reprocess the recurrent prefix in many request shapes, while Skippy can restore the exact compact state and decode. +Durable L3 manifests identify runtime-native KV segments as +`native-kv-page/1`. Those segments are the bytes exported by the active runtime +and are stored and restored verbatim; F32, F16, Q8_0, and Q4_0 remain in their +active runtime representation when that backend supports the type. Recurrent +and other auxiliary continuation state stays `raw/1`, with a hard segment +boundary between the representations. Before reading native segment bytes, +the server checks the manifest's runtime page descriptor against the located +prefix and encoded KV length. The exact-state identity already binds the +runtime ABI, platform, model, layer range, and KV configuration. + ```mermaid flowchart LR Payload["KV + recurrent payload"] --> Split["1 MiB chunks"] diff --git a/crates/skippy-cache/examples/cachegen_cubecl_spike.rs b/crates/skippy-cache/examples/cachegen_cubecl_spike.rs new file mode 100644 index 0000000000..8b166b7a5c --- /dev/null +++ b/crates/skippy-cache/examples/cachegen_cubecl_spike.rs @@ -0,0 +1,387 @@ +//! CubeCL feasibility spike for the CacheGen hot kernels (#1652). +//! +//! Evidence, not commitment: this example is feature-gated behind +//! `cachegen-spike`; nothing in the library links CubeCL. Per the agreed +//! gate it reports six numbers separately on CPU and on the best available +//! wgpu backend (Metal on Apple Silicon): +//! +//! 1. cold JIT/compile time (first launch of a kernel specialization) +//! 2. warm dispatch time (steady-state launches) +//! 3. host→device bytes copied +//! 4. device→host bytes copied +//! 5. peak temporary device memory +//! 6. output equality against the CPU reference (bitwise) +//! +//! Measurement honesty, per the review of PR #1752: every timed stage +//! ends with a `client.sync()` *inside* the timer, so the number covers +//! real completion — not unsynchronized enqueue. "Cold JIT" is the first +//! launch of a kernel specialization in this process (compile + pipeline +//! creation + execution); in cubecl 0.10.0 the only on-disk kernel cache +//! is the Vulkan-only SPIR-V cache, which this build does not enable, so +//! per-process first launch is the true cold path for CPU and Metal +//! alike. Warm numbers are the average over synchronized steady-state +//! launches. Equality is exact `==` on symbols and f32 values: dequantized +//! values are exact products of small integers and dyadic floats, so any +//! non-bit-equal output is a real divergence, not rounding noise. +//! +//! Kernel decomposition, stated honestly: quantization is embarrassingly +//! parallel; the token-axis delta is a scan along rows. KV tiles are +//! tall-thin (thousands of token rows, `dims` of one head layout), so the +//! spike assigns one unit per column and walks rows sequentially inside +//! the unit — the same arithmetic as the CPU reference, ordered so the +//! result is bit-exact. A parallel-scan kernel is the performance +//! follow-up; the rANS entropy stage remains CPU in this reference and is +//! measured here only as an encoded-size ratio. +//! +//! Run: +//! cargo run -p skippy-cache --example cachegen_cubecl_spike \ +//! --features cachegen-spike -- --rows 4096 --dims 128 [--iterations 20] +//! +//! `--iterations` only widens the warm-dispatch sample; the cold-JIT +//! number is a single first launch by definition. + +use cubecl::prelude::*; +use skippy_cache::cachegen::reference; +use skippy_protocol::binary::{f16_bits_to_f32, f32_to_f16_bits}; + +/// One unit per column; walks token rows sequentially so the delta ring +/// matches the CPU reference exactly. +#[cube(launch)] +fn quantize_delta_columns( + values: &Array, + symbols: &mut Array, + calib: &Array, + #[comptime] rows: usize, + #[comptime] dims: usize, +) { + let column = UNIT_POS_X as usize; + if column < dims { + let min = calib[0]; + let scale = calib[1]; + let mut prev: u32 = 0; + for row in 0..rows { + let index = row * dims + column; + let scaled = ((values[index] - min) / scale).round(); + let plain = u32::cast_from(scaled.clamp(0.0, 15.0)); + symbols[index] = (plain + 16 - prev) % 16; + prev = plain; + } + } +} + +/// Inverse scan: one unit per column re-accumulates the reconstructed +/// symbol sequence, then dequantizes. Bit-exact against the CPU reference. +#[cube(launch)] +fn undelta_dequantize_columns( + symbols: &Array, + values: &mut Array, + calib: &Array, + #[comptime] rows: usize, + #[comptime] dims: usize, +) { + let column = UNIT_POS_X as usize; + if column < dims { + let min = calib[0]; + let scale = calib[1]; + let mut prev: u32 = 0; + for row in 0..rows { + let index = row * dims + column; + let plain = (symbols[index] + prev) % 16; + values[index] = f32::cast_from(plain) * scale + min; + prev = plain; + } + } +} + +struct StageTiming { + /// First synchronized launch of this kernel specialization in this + /// process: JIT compile + pipeline creation + execution to completion. + cold_compile_ms: u128, + /// Average over synchronized steady-state launches (completion, not + /// enqueue). + warm_dispatch_us: u128, +} + +/// Blocks until the client's stream drains. Called inside every timed +/// region: without it the timer measures launch enqueue only and the +/// actual work lands after the clock stops (the PR #1752 review bug). +fn synchronize(client: &ComputeClient) { + cubecl::future::block_on(client.sync()) + .unwrap_or_else(|error| panic!("device sync failed: {error}")); +} + +fn timed_stage( + client: &ComputeClient, + iterations: u32, + mut launch: impl FnMut(), +) -> StageTiming { + synchronize::(client); + let cold = std::time::Instant::now(); + launch(); + synchronize::(client); + let cold_compile_ms = cold.elapsed().as_millis(); + // Steady state: the first launch populated the pipeline cache, so + // every launch here is the warm path. Each launch is individually + // synchronized inside the timed region; the reported number is total + // elapsed / iterations, i.e. the steady-state cost per launch + // including completion. + let warm_start = std::time::Instant::now(); + for _ in 0..iterations { + launch(); + synchronize::(client); + } + StageTiming { + cold_compile_ms, + warm_dispatch_us: warm_start.elapsed().as_micros() / u128::from(iterations), + } +} + +fn run_backend( + backend: &'static str, + tile: &[u8], + rows: usize, + dims: usize, + expected_symbols: &[u8], + expected_values: &[f32], + iterations: u32, +) -> Result<(), String> { + if dims > 1024 { + return Err(format!( + "{backend}: spike shape needs dims <= 1024 units, got {dims}" + )); + } + + let client = R::client(&R::Device::default()); + let count = rows * dims; + let values: Vec = tile + .as_chunks::<2>() + .0 + .iter() + .map(|bytes| f16_bits_to_f32(u16::from_le_bytes(*bytes))) + .collect(); + let calibration = reference::calibrate(&values).map_err(|error| error.to_string())?; + + // Timed H2D: the honest copy path is every byte the kernel work needs + // (the f32 tile plus the 8-byte calibration vector), measured to + // completion. + let h2d_start = std::time::Instant::now(); + let values_handle = client.create_from_slice(f32::as_bytes(&values)); + let calib_handle = client.create_from_slice(f32::as_bytes(&[ + f32::from_bits(calibration.min_bits), + f32::from_bits(calibration.scale_bits), + ])); + synchronize::(&client); + let h2d_us = h2d_start.elapsed().as_micros(); + let h2d_bytes = values.len() * core::mem::size_of::() + 2 * core::mem::size_of::(); + // Working-set allocations (outputs) are not timed; they are part of + // the live peak, not the transfer. + let symbols_handle = client.empty(count * core::mem::size_of::()); + let rebuilt_handle = client.empty(count * core::mem::size_of::()); + // Actual live allocation peak while the kernels run: every buffer the + // harness holds at once — inputs (tile + calibration) and both + // outputs. Reported rather than derived, per the review. + let peak_temporary_bytes = + values_handle.size() + calib_handle.size() + symbols_handle.size() + rebuilt_handle.size(); + if calibration.scale_bits == 0.0f32.to_bits() { + return Err( + "flat tile (scale == 0): not exercised by the spike; the CPU reference covers it" + .to_string(), + ); + } + + let encode = || unsafe { + quantize_delta_columns::launch::( + &client, + CubeCount::Static(1, 1, 1), + CubeDim::new_1d(dims as u32), + ArrayArg::from_raw_parts(values_handle.clone(), count), + ArrayArg::from_raw_parts(symbols_handle.clone(), count), + ArrayArg::from_raw_parts(calib_handle.clone(), 2), + rows, + dims, + ) + }; + let encode_timing = timed_stage::(&client, iterations, encode); + + let decode = || unsafe { + undelta_dequantize_columns::launch::( + &client, + CubeCount::Static(1, 1, 1), + CubeDim::new_1d(dims as u32), + ArrayArg::from_raw_parts(symbols_handle.clone(), count), + ArrayArg::from_raw_parts(rebuilt_handle.clone(), count), + ArrayArg::from_raw_parts(calib_handle.clone(), 2), + rows, + dims, + ) + }; + let decode_timing = timed_stage::(&client, iterations, decode); + + // Timed D2H: both returns measured to completion, like the uploads. + let d2h_start = std::time::Instant::now(); + let symbols_bytes = client + .read_one(symbols_handle.clone()) + .map_err(|error| error.to_string())?; + let rebuilt_bytes = client + .read_one(rebuilt_handle.clone()) + .map_err(|error| error.to_string())?; + synchronize::(&client); + let d2h_us = d2h_start.elapsed().as_micros(); + let d2h_bytes = symbols_bytes.len() + rebuilt_bytes.len(); + let device_symbols_u32 = u32::from_bytes(&symbols_bytes); + // Range-check before the narrowing cast: `as u8` would silently alias + // an out-of-alphabet u32 (e.g. 256 -> 0) into a false parity pass. + let all_in_alphabet = device_symbols_u32.iter().all(|&symbol| symbol < 16); + let device_symbols: Vec = device_symbols_u32.iter().map(|&s| s as u8).collect(); + let mut histogram = vec![0u32; reference::TOKEN_COUNT]; + for &symbol in &device_symbols { + histogram[usize::from(symbol)] += 1; + } + let freqs = + reference::histogram_to_freqs(&histogram, count).map_err(|error| error.to_string())?; + let table = skippy_cache::cachegen::rans::SymbolTable::from_freqs(&freqs) + .ok_or("rANS table construction failed")?; + let mut encoder = skippy_cache::cachegen::rans::RansEncoder::new(); + for &symbol in device_symbols.iter().rev() { + encoder.put(&table, usize::from(symbol)); + } + let stream = encoder.finish(); + let ratio = stream.len() as f64 / tile.len() as f64; + + let rebuilt = f32::from_bytes(&rebuilt_bytes); + let symbols_match = all_in_alphabet && device_symbols == expected_symbols; + // Exact bitwise comparison: dequantized values are `symbol * scale + + // min` where symbol is a small integer and scale/min are identical f32 + // bits on both sides, so the f32 words must match exactly. A tolerance + // here is what let an enqueue-only measurement pass for parity. + let values_match = rebuilt.len() == expected_values.len() + && rebuilt + .iter() + .zip(expected_values.iter()) + .all(|(device, reference_value)| device.to_bits() == reference_value.to_bits()); + let symbol_mismatches = device_symbols + .iter() + .zip(expected_symbols.iter()) + .filter(|(device, expected)| device != expected) + .count(); + let value_mismatches = rebuilt + .iter() + .zip(expected_values.iter()) + .filter(|(device, expected)| device.to_bits() != expected.to_bits()) + .count(); + + println!("=== {backend} ==="); + println!( + "quantize+delta: cold {} ms, warm {} us | undelta+dequantize: cold {} ms, warm {} us", + encode_timing.cold_compile_ms, + encode_timing.warm_dispatch_us, + decode_timing.cold_compile_ms, + decode_timing.warm_dispatch_us, + ); + println!( + "copies: H2D {} bytes in {} us (f32 tile + 8 B calibration, to completion) | D2H {} bytes in {} us (symbols + rebuilt) | live device buffer peak {} bytes (tile + calibration + both outputs)", + h2d_bytes, h2d_us, d2h_bytes, d2h_us, peak_temporary_bytes, + ); + println!( + "encoded-size ratio: rANS {} bytes / raw {} bytes = {:.3}", + stream.len(), + tile.len(), + ratio + ); + println!("equality vs CPU reference (bitwise): symbols={symbols_match}, values={values_match}"); + println!( + "mismatch counts: symbols {symbol_mismatches}/{}, values {value_mismatches}/{}", + device_symbols.len(), + rebuilt.len() + ); + if symbols_match && values_match { + Ok(()) + } else { + Err(format!( + "{backend}: device output diverges from the CPU reference (symbols={symbols_match}, values={values_match})" + )) + } +} + +fn main() { + let args: Vec = std::env::args().collect(); + let parse = |name: &str, default: usize| -> usize { + let position = args.iter().position(|argument| argument == name); + match position { + Some(index) => args + .get(index + 1) + .and_then(|value| value.parse().ok()) + .unwrap_or(default), + None => default, + } + }; + let rows = parse("--rows", 4096); + let dims = parse("--dims", 128); + let iterations = u32::try_from(parse("--iterations", 20)).unwrap_or(20); + + // Smooth KV-like fixture, deterministic, the shape CacheGen gains come + // from. + let count = rows * dims; + let mut tile = Vec::with_capacity(count * 2); + for row in 0..rows { + for column in 0..dims { + let phase = (row * dims + column) as f32; + let value = (phase * 0.000_5).sin() * 0.4 + 0.5; + tile.extend_from_slice(&f32_to_f16_bits(value).to_le_bytes()); + } + } + + // CPU reference outputs for the equality gate. + let values: Vec = tile + .as_chunks::<2>() + .0 + .iter() + .map(|bytes| f16_bits_to_f32(u16::from_le_bytes(*bytes))) + .collect(); + let calibration = reference::calibrate(&values).expect("calibrate"); + let mut symbols = reference::quantize(&calibration, &values).expect("quantize"); + reference::delta_encode(&mut symbols, dims).expect("delta"); + // The device decode path undeltas before dequantizing; the expected + // values must follow the same order of operations. + let mut undeltaed = symbols.clone(); + reference::delta_decode(&mut undeltaed, dims).expect("undelta"); + let expected_values = reference::dequantize(&calibration, &undeltaed); + + println!( + "tile: {rows} rows x {dims} dims = {count} f16 values ({} raw bytes)", + tile.len() + ); + + let mut failures = Vec::new(); + if let Err(error) = run_backend::( + "cubecl-cpu", + &tile, + rows, + dims, + &symbols, + &expected_values, + iterations, + ) { + failures.push(error); + } + if let Err(error) = run_backend::( + "wgpu(Metal)", + &tile, + rows, + dims, + &symbols, + &expected_values, + iterations, + ) { + failures.push(error); + } + + if failures.is_empty() { + println!("SPIKE PASS: device outputs match the CPU reference on every available backend"); + } else { + for failure in &failures { + eprintln!("SPIKE FAILURE: {failure}"); + } + std::process::exit(1); + } +} diff --git a/crates/skippy-cache/src/cachegen/archive.rs b/crates/skippy-cache/src/cachegen/archive.rs new file mode 100644 index 0000000000..47c7bfb8d8 --- /dev/null +++ b/crates/skippy-cache/src/cachegen/archive.rs @@ -0,0 +1,1133 @@ +//! Bounded CacheGen archives for complete native KV pages. +//! +//! A page archive covers every byte in the decoded runtime payload exactly +//! once. K/V records carry LMCache-compatible segments while auxiliary +//! indexer state remains exact. The archive is portable; native backends may +//! decode its validated records directly into resident KV tensors. + +use std::collections::BTreeMap; + +use anyhow::{Context, Result, anyhow, bail}; +use skippy_protocol::binary::{f16_bits_to_f32, f32_to_f16_bits}; + +use super::lmcache::{ + MAX_TOKENS_PER_CHUNK, bins_for_layer, decode_f16_segment, encode_f16_segment_packed, + validate_f16_segment, +}; + +pub const ARCHIVE_MAGIC: [u8; 4] = *b"CKG1"; +pub const ARCHIVE_HEADER_BYTES: usize = 16; +pub const RECORD_HEADER_BYTES: usize = 52; + +type ByteRange = (usize, usize); +type TransposedRegion = (usize, usize, usize, usize); +type TransposedCoverage = BTreeMap>; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(u8)] +pub enum RecordKind { + CacheGen = 0, + Exact = 1, + CacheGenTransposed = 2, + CacheGenF32 = 3, + CacheGenF32Transposed = 4, + CacheGenQ8_0 = 5, + CacheGenQ4_0 = 6, +} + +impl TryFrom for RecordKind { + type Error = anyhow::Error; + + fn try_from(value: u8) -> Result { + match value { + 0 => Ok(Self::CacheGen), + 1 => Ok(Self::Exact), + 2 => Ok(Self::CacheGenTransposed), + 3 => Ok(Self::CacheGenF32), + 4 => Ok(Self::CacheGenF32Transposed), + 5 => Ok(Self::CacheGenQ8_0), + 6 => Ok(Self::CacheGenQ4_0), + _ => bail!("unknown CacheGen archive record kind {value}"), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ValueType { + F32, + F16, + Q8_0, + Q4_0, +} + +impl ValueType { + fn record_kind(self, transposed: bool) -> Result { + match (self, transposed) { + (Self::F16, false) => Ok(RecordKind::CacheGen), + (Self::F16, true) => Ok(RecordKind::CacheGenTransposed), + (Self::F32, false) => Ok(RecordKind::CacheGenF32), + (Self::F32, true) => Ok(RecordKind::CacheGenF32Transposed), + (Self::Q8_0, false) => Ok(RecordKind::CacheGenQ8_0), + (Self::Q4_0, false) => Ok(RecordKind::CacheGenQ4_0), + (Self::Q8_0 | Self::Q4_0, true) => { + bail!("CacheGen does not accept transposed block-quantized rows") + } + } + } + + fn element_bytes(self) -> u8 { + match self { + Self::F32 => 4, + Self::F16 => 2, + Self::Q8_0 => 34, + Self::Q4_0 => 18, + } + } + + fn channels_for_row(self, row_bytes: usize) -> Result { + let (block_bytes, block_values) = match self { + Self::F32 => (4, 1), + Self::F16 => (2, 1), + Self::Q8_0 => (34, 32), + Self::Q4_0 => (18, 32), + }; + if row_bytes == 0 || !row_bytes.is_multiple_of(block_bytes) { + bail!("CacheGen row size is invalid for {self:?}"); + } + Ok(row_bytes / block_bytes * block_values) + } +} + +fn record_value_type(kind: RecordKind) -> Option<(ValueType, bool)> { + match kind { + RecordKind::CacheGen => Some((ValueType::F16, false)), + RecordKind::CacheGenTransposed => Some((ValueType::F16, true)), + RecordKind::CacheGenF32 => Some((ValueType::F32, false)), + RecordKind::CacheGenF32Transposed => Some((ValueType::F32, true)), + RecordKind::CacheGenQ8_0 => Some((ValueType::Q8_0, false)), + RecordKind::CacheGenQ4_0 => Some((ValueType::Q4_0, false)), + RecordKind::Exact => None, + } +} + +/// Native-page geometry for one independent base or sliding-window cache. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ComponentLayout { + pub token_count: u64, + pub layer_count: u32, + pub k_type: ValueType, + pub v_type: ValueType, + pub k_row_bytes: u32, + pub v_row_bytes: u32, + pub v_element_bytes: u32, + pub k_idx_row_bytes: u32, + pub payload_offset: u64, + pub payload_bytes: u64, + pub v_transposed: bool, +} + +/// Complete native-page geometry. Components must cover `payload_bytes` +/// exactly and appear in output order. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PageLayout { + pub payload_bytes: u64, + pub components: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CacheGenArchive { + pub bytes: Vec, + pub tile_count: usize, + pub estimated_peak_codec_working_bytes: usize, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Record<'a> { + pub kind: RecordKind, + pub element_bytes: usize, + pub output_offset: usize, + pub decoded_len: usize, + pub token_count: usize, + pub token_start: usize, + pub total_tokens: usize, + pub payload: &'a [u8], +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ValidatedArchive<'a> { + pub raw_len: usize, + pub records: Vec>, +} + +struct OwnedRecord { + kind: RecordKind, + element_bytes: u8, + output_offset: usize, + decoded_len: usize, + token_count: u64, + token_start: u64, + total_tokens: u64, + payload: Vec, +} + +pub fn encode_page(layout: &PageLayout, raw: &[u8]) -> Result { + let raw_len = usize::try_from(layout.payload_bytes).context("page size exceeds usize")?; + if raw.len() != raw_len { + bail!("CacheGen page length disagrees with its layout"); + } + validate_components(layout, raw_len)?; + + let mut records = Vec::new(); + for component in &layout.components { + encode_component(*component, raw, &mut records)?; + } + records.sort_by_key(|record| (record.output_offset, record.token_start)); + validate_owned_record_coverage(&records, raw_len)?; + + let record_count = u32::try_from(records.len()).context("too many CacheGen archive records")?; + let payload_bytes = + records + .iter() + .try_fold(ARCHIVE_HEADER_BYTES, |total, record| -> Result { + total + .checked_add(RECORD_HEADER_BYTES) + .and_then(|value| value.checked_add(record.payload.len())) + .ok_or_else(|| anyhow!("CacheGen archive length overflow")) + })?; + let mut bytes = Vec::with_capacity(payload_bytes); + bytes.extend_from_slice(&ARCHIVE_MAGIC); + bytes.extend_from_slice(&layout.payload_bytes.to_le_bytes()); + bytes.extend_from_slice(&record_count.to_le_bytes()); + + let mut largest_working_set = 0usize; + let mut tile_count = 0usize; + for record in records { + bytes.push(record.kind as u8); + bytes.push(record.element_bytes); + bytes.extend_from_slice(&[0, 0]); + bytes.extend_from_slice(&(record.output_offset as u64).to_le_bytes()); + bytes.extend_from_slice(&(record.decoded_len as u64).to_le_bytes()); + bytes.extend_from_slice(&record.token_count.to_le_bytes()); + bytes.extend_from_slice(&record.token_start.to_le_bytes()); + bytes.extend_from_slice(&record.total_tokens.to_le_bytes()); + bytes.extend_from_slice(&(record.payload.len() as u64).to_le_bytes()); + bytes.extend_from_slice(&record.payload); + if record.kind != RecordKind::Exact { + tile_count += 1; + let segment = validate_f16_segment(&record.payload)?; + let values = segment.rows.saturating_mul(segment.channels); + let codec_working = record + .decoded_len + .saturating_add(values.saturating_mul(5)) + .saturating_add(record.payload.len()); + largest_working_set = largest_working_set.max(codec_working); + } + } + + validate_archive(&bytes, raw_len).context("validate encoded CacheGen page")?; + + Ok(CacheGenArchive { + estimated_peak_codec_working_bytes: bytes + .len() + .saturating_add(raw.len()) + .saturating_add(largest_working_set), + bytes, + tile_count, + }) +} + +pub fn validate_archive(archive: &[u8], expected_raw_len: usize) -> Result> { + if archive.len() < ARCHIVE_HEADER_BYTES || archive[..4] != ARCHIVE_MAGIC { + bail!("invalid CacheGen archive header"); + } + let raw_len = usize::try_from(read_u64(&archive[4..12])?) + .context("CacheGen archive raw length exceeds usize")?; + if raw_len != expected_raw_len { + bail!("CacheGen archive raw length disagrees with KV descriptor"); + } + let record_count = read_u32(&archive[12..16])? as usize; + let minimum_headers = record_count + .checked_mul(RECORD_HEADER_BYTES) + .and_then(|value| value.checked_add(ARCHIVE_HEADER_BYTES)) + .ok_or_else(|| anyhow!("CacheGen archive record table overflows"))?; + if minimum_headers > archive.len() { + bail!("truncated CacheGen archive record table"); + } + + let mut cursor = ARCHIVE_HEADER_BYTES; + let mut records = Vec::with_capacity(record_count); + let mut coverage = Vec::with_capacity(record_count); + let mut transposed_coverage = TransposedCoverage::new(); + let mut decoded_total = 0usize; + for _ in 0..record_count { + let header = archive + .get(cursor..cursor + RECORD_HEADER_BYTES) + .ok_or_else(|| anyhow!("truncated CacheGen archive record header"))?; + cursor += RECORD_HEADER_BYTES; + let kind = RecordKind::try_from(header[0])?; + let element_bytes = header[1] as usize; + if header[2..4] != [0, 0] { + bail!("CacheGen archive record reserved bytes are non-zero"); + } + let output_offset = usize::try_from(read_u64(&header[4..12])?) + .context("record output offset exceeds usize")?; + let decoded_len = usize::try_from(read_u64(&header[12..20])?) + .context("record decoded length exceeds usize")?; + let token_count = usize::try_from(read_u64(&header[20..28])?) + .context("record token count exceeds usize")?; + let token_start = usize::try_from(read_u64(&header[28..36])?) + .context("record token start exceeds usize")?; + let total_tokens = usize::try_from(read_u64(&header[36..44])?) + .context("record total token count exceeds usize")?; + let payload_len = usize::try_from(read_u64(&header[44..52])?) + .context("record payload length exceeds usize")?; + let payload_end = cursor + .checked_add(payload_len) + .ok_or_else(|| anyhow!("record payload range overflow"))?; + let payload = archive + .get(cursor..payload_end) + .ok_or_else(|| anyhow!("truncated CacheGen archive record payload"))?; + cursor = payload_end; + decoded_total = decoded_total + .checked_add(decoded_len) + .context("decoded archive size overflow")?; + + validate_record_geometry(Record { + kind, + element_bytes, + output_offset, + decoded_len, + token_count, + token_start, + total_tokens, + payload, + })?; + record_coverage( + kind, + output_offset, + decoded_len, + token_start, + token_count, + total_tokens, + element_bytes, + &mut coverage, + &mut transposed_coverage, + )?; + records.push(Record { + kind, + element_bytes, + output_offset, + decoded_len, + token_count, + token_start, + total_tokens, + payload, + }); + } + if cursor != archive.len() { + bail!("CacheGen archive has trailing bytes"); + } + validate_decoded_coverage( + &mut coverage, + &mut transposed_coverage, + raw_len, + decoded_total, + )?; + Ok(ValidatedArchive { raw_len, records }) +} + +pub fn decode_page(archive: &[u8], expected_raw_len: usize) -> Result> { + let validated = validate_archive(archive, expected_raw_len)?; + let mut decoded = vec![0u8; validated.raw_len]; + for record in validated.records { + match record_value_type(record.kind) { + Some((value_type, false)) => { + let output_end = record + .output_offset + .checked_add(record.decoded_len) + .context("record output range overflow")?; + let output = decoded + .get_mut(record.output_offset..output_end) + .ok_or_else(|| anyhow!("record output range exceeds KV payload"))?; + let f16 = decode_f16_segment(record.payload)?; + let tile = encode_native_rows_from_f16(value_type, &f16, record.token_count)?; + output.copy_from_slice(&tile); + } + None => { + let output_end = record + .output_offset + .checked_add(record.decoded_len) + .context("record output range overflow")?; + decoded[record.output_offset..output_end].copy_from_slice(record.payload); + } + Some((value_type, true)) => { + let f16 = decode_f16_segment(record.payload)?; + let token_major = + encode_native_rows_from_f16(value_type, &f16, record.token_count)?; + let dims = record.decoded_len / record.token_count / record.element_bytes; + let layer_bytes = record.total_tokens * dims * record.element_bytes; + let output_end = record.output_offset + layer_bytes; + transpose_range_from_token_major( + &token_major, + &mut decoded[record.output_offset..output_end], + record.total_tokens, + record.token_start, + record.token_count, + dims, + record.element_bytes, + )?; + } + } + } + Ok(decoded) +} + +fn validate_components(layout: &PageLayout, raw_len: usize) -> Result<()> { + if layout.components.is_empty() { + bail!("CacheGen page has no components"); + } + let mut next = 0usize; + for component in &layout.components { + let offset = + usize::try_from(component.payload_offset).context("component offset exceeds usize")?; + let bytes = + usize::try_from(component.payload_bytes).context("component size exceeds usize")?; + if offset != next || component.token_count == 0 || component.layer_count == 0 { + bail!("CacheGen components are empty, reordered, or leave a gap"); + } + next = next + .checked_add(bytes) + .ok_or_else(|| anyhow!("component range overflow"))?; + } + if next != raw_len { + bail!("CacheGen components do not cover the native page"); + } + Ok(()) +} + +fn encode_component( + component: ComponentLayout, + raw: &[u8], + records: &mut Vec, +) -> Result<()> { + let token_count = + usize::try_from(component.token_count).context("token count exceeds usize")?; + let layer_count = component.layer_count as usize; + let k_row = component.k_row_bytes as usize; + let k_channels = component.k_type.channels_for_row(k_row)?; + let base = usize::try_from(component.payload_offset).context("payload offset exceeds usize")?; + let component_len = + usize::try_from(component.payload_bytes).context("payload size exceeds usize")?; + let k_layer_bytes = token_count + .checked_mul(k_row) + .context("K layer size overflow")?; + let k_bytes = layer_count + .checked_mul(k_layer_bytes) + .context("K size overflow")?; + let k_idx_bytes = layer_count + .checked_mul(token_count) + .and_then(|value| value.checked_mul(component.k_idx_row_bytes as usize)) + .context("K-index size overflow")?; + let v_bytes = component_len + .checked_sub(k_bytes) + .and_then(|value| value.checked_sub(k_idx_bytes)) + .ok_or_else(|| anyhow!("component payload is shorter than K and indexer state"))?; + if !v_bytes.is_multiple_of(layer_count) { + bail!("V payload is not uniform across layers"); + } + let v_layer_bytes = v_bytes / layer_count; + if !v_layer_bytes.is_multiple_of(token_count) { + bail!("V payload is not uniform across tokens"); + } + let v_row = v_layer_bytes / token_count; + let v_channels = component.v_type.channels_for_row(v_row)?; + if component.v_transposed { + if component.v_type.element_bytes() as u32 != component.v_element_bytes { + bail!("transposed V page element size disagrees with its value type"); + } + } else if component.v_row_bytes as usize != v_row { + bail!("non-transposed V payload disagrees with its row size"); + } + + let end = base + .checked_add(component_len) + .context("component range overflow")?; + let component_raw = raw + .get(base..end) + .ok_or_else(|| anyhow!("component range exceeds KV payload"))?; + + for layer in 0..layer_count { + let layer_offset = layer * k_layer_bytes; + for row_start in (0..token_count).step_by(MAX_TOKENS_PER_CHUNK) { + let rows = (token_count - row_start).min(MAX_TOKENS_PER_CHUNK); + let local_offset = layer_offset + row_start * k_row; + let tile = &component_raw[local_offset..local_offset + rows * k_row]; + records.push(OwnedRecord { + kind: component.k_type.record_kind(false)?, + element_bytes: component.k_type.element_bytes(), + output_offset: base + local_offset, + decoded_len: tile.len(), + token_count: rows as u64, + token_start: 0, + total_tokens: 0, + payload: encode_f16_segment_packed( + &decode_native_rows_to_f16(component.k_type, tile, rows, k_row)?, + k_channels, + bins_for_layer(layer, layer_count, true), + )?, + }); + } + } + + let v_base = k_bytes; + for layer in 0..layer_count { + let layer_offset = v_base + layer * v_layer_bytes; + let layer_tile = &component_raw[layer_offset..layer_offset + v_layer_bytes]; + for row_start in (0..token_count).step_by(MAX_TOKENS_PER_CHUNK) { + let rows = (token_count - row_start).min(MAX_TOKENS_PER_CHUNK); + let (kind, output_offset, encoded_input, token_start, total_tokens) = + if component.v_transposed { + let native = transpose_range_to_token_major( + layer_tile, + token_count, + row_start, + rows, + v_channels, + component.v_element_bytes as usize, + )?; + ( + component.v_type.record_kind(true)?, + base + layer_offset, + decode_native_rows_to_f16(component.v_type, &native, rows, v_row)?, + row_start as u64, + token_count as u64, + ) + } else { + let local_offset = layer_offset + row_start * v_row; + let native = &component_raw[local_offset..local_offset + rows * v_row]; + ( + component.v_type.record_kind(false)?, + base + local_offset, + decode_native_rows_to_f16(component.v_type, native, rows, v_row)?, + 0, + 0, + ) + }; + records.push(OwnedRecord { + kind, + element_bytes: component.v_type.element_bytes(), + output_offset, + decoded_len: rows * v_row, + token_count: rows as u64, + token_start, + total_tokens, + payload: encode_f16_segment_packed( + &encoded_input, + v_channels, + bins_for_layer(layer, layer_count, false), + )?, + }); + } + } + if k_idx_bytes > 0 { + let local_offset = k_bytes + v_bytes; + records.push(OwnedRecord { + kind: RecordKind::Exact, + element_bytes: 1, + output_offset: base + local_offset, + decoded_len: k_idx_bytes, + token_count: component.token_count, + token_start: 0, + total_tokens: 0, + payload: component_raw[local_offset..local_offset + k_idx_bytes].to_vec(), + }); + } + Ok(()) +} + +fn validate_record_geometry(record: Record<'_>) -> Result<()> { + if record.decoded_len == 0 || record.token_count == 0 { + bail!("CacheGen archive record has empty geometry"); + } + match record_value_type(record.kind) { + Some((value_type, transposed)) => { + if record.element_bytes != value_type.element_bytes() as usize { + bail!("CacheGen record has an invalid element size"); + } + let segment = validate_f16_segment(record.payload)?; + let expected_row = native_row_bytes(value_type, segment.channels)?; + let expected = segment + .rows + .checked_mul(expected_row) + .ok_or_else(|| anyhow!("CacheGen segment native decoded length overflow"))?; + if segment.rows != record.token_count || expected != record.decoded_len { + bail!("CacheGen segment geometry disagrees with its archive record"); + } + if transposed { + if record.total_tokens == 0 + || record + .token_start + .checked_add(record.token_count) + .is_none_or(|end| end > record.total_tokens) + { + bail!("invalid transposed CacheGen archive record geometry"); + } + } else if record.token_start != 0 || record.total_tokens != 0 { + bail!("non-transposed CacheGen record carries transpose geometry"); + } + } + None => { + if record.element_bytes != 1 + || record.payload.len() != record.decoded_len + || record.token_start != 0 + || record.total_tokens != 0 + { + bail!("exact CacheGen archive record has invalid geometry"); + } + } + } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +fn record_coverage( + kind: RecordKind, + output_offset: usize, + decoded_len: usize, + token_start: usize, + token_count: usize, + total_tokens: usize, + element_bytes: usize, + coverage: &mut Vec, + transposed: &mut TransposedCoverage, +) -> Result<()> { + if record_value_type(kind).is_some_and(|(_, transposed)| transposed) { + let dims = decoded_len + .checked_div(token_count) + .and_then(|value| value.checked_div(element_bytes)) + .filter(|dims| *dims > 0) + .ok_or_else(|| anyhow!("transposed record geometry overflow"))?; + transposed + .entry((output_offset, total_tokens, dims, element_bytes)) + .or_default() + .push((token_start, token_start + token_count)); + } else { + let end = output_offset + .checked_add(decoded_len) + .ok_or_else(|| anyhow!("record output range overflow"))?; + coverage.push((output_offset, end)); + } + Ok(()) +} + +fn validate_owned_record_coverage(records: &[OwnedRecord], raw_len: usize) -> Result<()> { + let decoded_total = records.iter().try_fold(0usize, |total, record| { + total + .checked_add(record.decoded_len) + .ok_or_else(|| anyhow!("archive decoded size overflow")) + })?; + if decoded_total != raw_len { + bail!("CacheGen archive records do not account for the KV payload"); + } + Ok(()) +} + +fn validate_decoded_coverage( + coverage: &mut Vec, + transposed: &mut TransposedCoverage, + raw_len: usize, + decoded_total: usize, +) -> Result<()> { + if decoded_total != raw_len { + bail!("CacheGen archive decoded bytes do not account for the KV payload"); + } + for (&(output_offset, total_tokens, dims, element_bytes), ranges) in transposed.iter_mut() { + ranges.sort_unstable(); + let mut next_token = 0usize; + for &(start, end) in ranges.iter() { + if start != next_token || end < start { + bail!("transposed CacheGen records overlap or leave a token gap"); + } + next_token = end; + } + if next_token != total_tokens { + bail!("transposed CacheGen records leave a token gap"); + } + let end = output_offset + .checked_add( + total_tokens + .checked_mul(dims) + .and_then(|value| value.checked_mul(element_bytes)) + .context("transposed coverage size overflow")?, + ) + .context("transposed coverage range overflow")?; + coverage.push((output_offset, end)); + } + coverage.sort_unstable(); + let mut next = 0usize; + for &(start, end) in coverage.iter() { + if start != next || end < start { + bail!("CacheGen archive records do not exactly cover the KV payload"); + } + next = end; + } + if next != raw_len { + bail!("CacheGen archive records leave a gap in the KV payload"); + } + Ok(()) +} + +fn native_row_bytes(value_type: ValueType, channels: usize) -> Result { + match value_type { + ValueType::F32 => channels.checked_mul(4).context("F32 row size overflow"), + ValueType::F16 => channels.checked_mul(2).context("F16 row size overflow"), + ValueType::Q8_0 => { + if !channels.is_multiple_of(32) { + bail!("Q8_0 rows require a multiple of 32 values"); + } + channels + .checked_div(32) + .and_then(|blocks| blocks.checked_mul(34)) + .context("Q8_0 row size overflow") + } + ValueType::Q4_0 => { + if !channels.is_multiple_of(32) { + bail!("Q4_0 rows require a multiple of 32 values"); + } + channels + .checked_div(32) + .and_then(|blocks| blocks.checked_mul(18)) + .context("Q4_0 row size overflow") + } + } +} + +fn decode_native_rows_to_f16( + value_type: ValueType, + native: &[u8], + rows: usize, + row_bytes: usize, +) -> Result> { + let expected = rows + .checked_mul(row_bytes) + .context("native tile size overflow")?; + if native.len() != expected { + bail!("native CacheGen tile length disagrees with its row geometry"); + } + let channels = value_type.channels_for_row(row_bytes)?; + let mut output = Vec::with_capacity( + rows.checked_mul(channels) + .and_then(|values| values.checked_mul(2)) + .context("F16 adapter output size overflow")?, + ); + match value_type { + ValueType::F16 => output.extend_from_slice(native), + ValueType::F32 => { + for value in native.as_chunks::<4>().0 { + let value = f32::from_le_bytes(*value); + output.extend_from_slice(&f32_to_f16_bits(value).to_le_bytes()); + } + } + ValueType::Q8_0 => { + for row in native.chunks_exact(row_bytes) { + for block in row.as_chunks::<34>().0 { + let scale = f16_bits_to_f32(u16::from_le_bytes([block[0], block[1]])); + for &quant in &block[2..] { + let value = f32::from(quant as i8) * scale; + output.extend_from_slice(&f32_to_f16_bits(value).to_le_bytes()); + } + } + } + } + ValueType::Q4_0 => { + for row in native.chunks_exact(row_bytes) { + for block in row.as_chunks::<18>().0 { + let scale = f16_bits_to_f32(u16::from_le_bytes([block[0], block[1]])); + for half in 0..2 { + for &packed in &block[2..] { + let quant = if half == 0 { + packed & 0x0f + } else { + packed >> 4 + }; + let value = (i32::from(quant) - 8) as f32 * scale; + output.extend_from_slice(&f32_to_f16_bits(value).to_le_bytes()); + } + } + } + } + } + } + Ok(output) +} + +fn encode_native_rows_from_f16(value_type: ValueType, f16: &[u8], rows: usize) -> Result> { + if rows == 0 || !f16.len().is_multiple_of(rows * 2) { + bail!("decoded CacheGen tile has invalid row geometry"); + } + let channels = f16.len() / rows / 2; + let row_bytes = native_row_bytes(value_type, channels)?; + let mut output = Vec::with_capacity( + rows.checked_mul(row_bytes) + .context("native adapter output size overflow")?, + ); + for row in f16.chunks_exact(channels * 2) { + match value_type { + ValueType::F16 => output.extend_from_slice(row), + ValueType::F32 => { + for value in row.as_chunks::<2>().0 { + let value = f16_bits_to_f32(u16::from_le_bytes([value[0], value[1]])); + output.extend_from_slice(&value.to_le_bytes()); + } + } + ValueType::Q8_0 => quantize_q8_0_row(row, &mut output), + ValueType::Q4_0 => quantize_q4_0_row(row, &mut output), + } + } + Ok(output) +} + +fn f16_row_values(row: &[u8]) -> impl Iterator + '_ { + row.as_chunks::<2>() + .0 + .iter() + .map(|value| f16_bits_to_f32(u16::from_le_bytes([value[0], value[1]]))) +} + +fn quantize_q8_0_row(row: &[u8], output: &mut Vec) { + for block in row.as_chunks::<64>().0 { + let values = f16_row_values(block).collect::>(); + let amax = values + .iter() + .fold(0.0_f32, |acc, value| acc.max(value.abs())); + let scale = amax / 127.0; + let inverse = if scale == 0.0 { 0.0 } else { scale.recip() }; + output.extend_from_slice(&f32_to_f16_bits(scale).to_le_bytes()); + output.extend( + values + .into_iter() + .map(|value| (value * inverse).round().clamp(-127.0, 127.0) as i8 as u8), + ); + } +} + +fn quantize_q4_0_row(row: &[u8], output: &mut Vec) { + for block in row.as_chunks::<64>().0 { + let values = f16_row_values(block).collect::>(); + let mut amax = 0.0_f32; + let mut signed_max = 0.0_f32; + for &value in &values { + if amax < value.abs() { + amax = value.abs(); + signed_max = value; + } + } + let scale = signed_max / -8.0; + let inverse = if scale == 0.0 { 0.0 } else { scale.recip() }; + output.extend_from_slice(&f32_to_f16_bits(scale).to_le_bytes()); + for index in 0..16 { + let low = (values[index] * inverse + 8.5).trunc().clamp(0.0, 15.0) as u8; + let high = (values[index + 16] * inverse + 8.5) + .trunc() + .clamp(0.0, 15.0) as u8; + output.push(low | (high << 4)); + } + } +} + +fn transpose_range_to_token_major( + source: &[u8], + total_tokens: usize, + token_start: usize, + token_count: usize, + dims: usize, + element_bytes: usize, +) -> Result> { + let expected = total_tokens + .checked_mul(dims) + .and_then(|value| value.checked_mul(element_bytes)) + .context("transpose size overflow")?; + if source.len() != expected { + bail!("transposed V tile length does not match its geometry"); + } + let output_len = token_count + .checked_mul(dims) + .and_then(|value| value.checked_mul(element_bytes)) + .context("transpose output size overflow")?; + if token_start + .checked_add(token_count) + .is_none_or(|end| end > total_tokens) + { + bail!("transpose token range exceeds source geometry"); + } + let mut output = vec![0u8; output_len]; + for local_token in 0..token_count { + let token = token_start + local_token; + for dim in 0..dims { + let source_offset = (dim * total_tokens + token) * element_bytes; + let output_offset = (local_token * dims + dim) * element_bytes; + output[output_offset..output_offset + element_bytes] + .copy_from_slice(&source[source_offset..source_offset + element_bytes]); + } + } + Ok(output) +} + +fn transpose_range_from_token_major( + source: &[u8], + output: &mut [u8], + total_tokens: usize, + token_start: usize, + token_count: usize, + dims: usize, + element_bytes: usize, +) -> Result<()> { + let expected_source = token_count + .checked_mul(dims) + .and_then(|value| value.checked_mul(element_bytes)) + .context("transpose source size overflow")?; + let expected_output = total_tokens + .checked_mul(dims) + .and_then(|value| value.checked_mul(element_bytes)) + .context("transpose output size overflow")?; + if source.len() != expected_source || output.len() != expected_output { + bail!("transposed CacheGen decode length mismatch"); + } + for local_token in 0..token_count { + let token = token_start + local_token; + for dim in 0..dims { + let source_offset = (local_token * dims + dim) * element_bytes; + let output_offset = (dim * total_tokens + token) * element_bytes; + output[output_offset..output_offset + element_bytes] + .copy_from_slice(&source[source_offset..source_offset + element_bytes]); + } + } + Ok(()) +} + +fn read_u64(bytes: &[u8]) -> Result { + Ok(u64::from_le_bytes( + bytes.try_into().context("u64 field length")?, + )) +} + +fn read_u32(bytes: &[u8]) -> Result { + Ok(u32::from_le_bytes( + bytes.try_into().context("u32 field length")?, + )) +} + +#[cfg(test)] +mod tests { + use super::*; + use skippy_protocol::binary::f32_to_f16_bits; + + fn f16_bytes(values: usize) -> Vec { + (0..values) + .flat_map(|index| f32_to_f16_bits(index as f32 / 17.0).to_le_bytes()) + .collect() + } + + fn layout(token_count: u64, v_transposed: bool) -> PageLayout { + let payload_bytes = token_count * 2 * 2 * 6; + PageLayout { + payload_bytes, + components: vec![ComponentLayout { + token_count, + layer_count: 2, + k_type: ValueType::F16, + v_type: ValueType::F16, + k_row_bytes: 6, + v_row_bytes: if v_transposed { 0 } else { 6 }, + v_element_bytes: if v_transposed { 2 } else { 0 }, + k_idx_row_bytes: 0, + payload_offset: 0, + payload_bytes, + v_transposed, + }], + } + } + + #[test] + fn page_roundtrip_preserves_geometry_and_length() { + let layout = layout(4, false); + let raw = f16_bytes(48); + let archive = encode_page(&layout, &raw).expect("encode"); + let decoded = decode_page(&archive.bytes, raw.len()).expect("decode"); + assert_eq!(decoded.len(), raw.len()); + assert_eq!(archive.tile_count, 4); + assert_ne!(decoded, raw, "fixture must exercise lossy quantization"); + assert_eq!( + &archive.bytes[ARCHIVE_HEADER_BYTES + RECORD_HEADER_BYTES..][..4], + b"LCG2" + ); + } + + #[test] + fn transposed_v_layout_uses_inferred_row_width() { + let layout = layout(4, true); + let raw = f16_bytes(48); + let archive = encode_page(&layout, &raw).expect("encode"); + let decoded = decode_page(&archive.bytes, raw.len()).expect("decode"); + assert_eq!(decoded.len(), raw.len()); + assert_eq!(archive.tile_count, 4); + } + + #[test] + fn validation_rejects_a_corrupt_nested_segment_before_decode() { + let layout = layout(4, false); + let raw = f16_bytes(48); + let mut archive = encode_page(&layout, &raw).expect("encode").bytes; + archive[ARCHIVE_HEADER_BYTES + RECORD_HEADER_BYTES] = b'X'; + assert!(validate_archive(&archive, raw.len()).is_err()); + } + + #[test] + fn long_transposed_pages_are_chunked_at_the_reference_limit() { + let token_count = MAX_TOKENS_PER_CHUNK as u64 + 4; + let layout = layout(token_count, true); + let raw = f16_bytes(layout.payload_bytes as usize / 2); + let archive = encode_page(&layout, &raw).expect("encode"); + let decoded = decode_page(&archive.bytes, raw.len()).expect("decode"); + assert_eq!(decoded.len(), raw.len()); + assert_eq!(archive.tile_count, 8); + } + + #[test] + fn archive_rejects_trailing_and_uncovered_bytes() { + let layout = layout(4, false); + let raw = f16_bytes(48); + let mut archive = encode_page(&layout, &raw).expect("encode").bytes; + archive.push(0); + assert!(validate_archive(&archive, raw.len()).is_err()); + + let mut uncovered = layout; + uncovered.payload_bytes += 2; + assert!(encode_page(&uncovered, &f16_bytes(49)).is_err()); + } + + fn typed_layout(k_type: ValueType, v_type: ValueType) -> PageLayout { + let token_count = 4; + let layer_count = 2; + let channels = 32; + let k_row_bytes = native_row_bytes(k_type, channels).expect("K row"); + let v_row_bytes = native_row_bytes(v_type, channels).expect("V row"); + let payload_bytes = token_count * layer_count * (k_row_bytes + v_row_bytes); + PageLayout { + payload_bytes: payload_bytes as u64, + components: vec![ComponentLayout { + token_count: token_count as u64, + layer_count: layer_count as u32, + k_type, + v_type, + k_row_bytes: k_row_bytes as u32, + v_row_bytes: v_row_bytes as u32, + v_element_bytes: 0, + k_idx_row_bytes: 0, + payload_offset: 0, + payload_bytes: payload_bytes as u64, + v_transposed: false, + }], + } + } + + fn typed_page(layout: &PageLayout) -> Vec { + let component = layout.components[0]; + let rows = component.token_count as usize * component.layer_count as usize; + let channels = 32; + let mut f16 = Vec::with_capacity(rows * channels * 2); + for index in 0..rows * channels { + let value = ((index % 37) as f32 - 18.0) / 11.0; + f16.extend_from_slice(&f32_to_f16_bits(value).to_le_bytes()); + } + let mut raw = encode_native_rows_from_f16(component.k_type, &f16, rows).expect("K"); + raw.extend(encode_native_rows_from_f16(component.v_type, &f16, rows).expect("V")); + raw + } + + #[test] + fn portable_archive_adapts_f32_q8_q4_and_mixed_pages() { + for (k_type, v_type, k_kind, v_kind) in [ + ( + ValueType::F32, + ValueType::F32, + RecordKind::CacheGenF32, + RecordKind::CacheGenF32, + ), + ( + ValueType::Q8_0, + ValueType::Q8_0, + RecordKind::CacheGenQ8_0, + RecordKind::CacheGenQ8_0, + ), + ( + ValueType::Q4_0, + ValueType::Q4_0, + RecordKind::CacheGenQ4_0, + RecordKind::CacheGenQ4_0, + ), + ( + ValueType::Q8_0, + ValueType::Q4_0, + RecordKind::CacheGenQ8_0, + RecordKind::CacheGenQ4_0, + ), + ] { + let layout = typed_layout(k_type, v_type); + let raw = typed_page(&layout); + let archive = encode_page(&layout, &raw).expect("encode typed page"); + let validated = validate_archive(&archive.bytes, raw.len()).expect("validate"); + assert_eq!(validated.records.len(), 4); + assert!( + validated.records[..2] + .iter() + .all(|record| record.kind == k_kind) + ); + assert!( + validated.records[2..] + .iter() + .all(|record| record.kind == v_kind) + ); + + let decoded = decode_page(&archive.bytes, raw.len()).expect("decode typed page"); + assert_eq!(decoded.len(), raw.len()); + assert_eq!(archive.tile_count, 4); + } + } + + #[test] + fn f32_transposed_v_uses_a_typed_record_without_changing_the_container() { + let token_count = 4_u64; + let row_bytes = 32 * 4; + let payload_bytes = token_count * 2 * row_bytes; + let layout = PageLayout { + payload_bytes, + components: vec![ComponentLayout { + token_count, + layer_count: 1, + k_type: ValueType::F32, + v_type: ValueType::F32, + k_row_bytes: row_bytes as u32, + v_row_bytes: 0, + v_element_bytes: 4, + k_idx_row_bytes: 0, + payload_offset: 0, + payload_bytes, + v_transposed: true, + }], + }; + let raw = (0..payload_bytes / 4) + .flat_map(|index| ((index as f32 - 50.0) / 13.0).to_le_bytes()) + .collect::>(); + let archive = encode_page(&layout, &raw).expect("encode F32 transposed page"); + let validated = validate_archive(&archive.bytes, raw.len()).expect("validate"); + assert_eq!(validated.records[0].kind, RecordKind::CacheGenF32); + assert_eq!(validated.records[1].kind, RecordKind::CacheGenF32Transposed); + assert_eq!( + decode_page(&archive.bytes, raw.len()) + .expect("decode") + .len(), + raw.len() + ); + } +} diff --git a/crates/skippy-cache/src/cachegen/container.rs b/crates/skippy-cache/src/cachegen/container.rs new file mode 100644 index 0000000000..8c144330ee --- /dev/null +++ b/crates/skippy-cache/src/cachegen/container.rs @@ -0,0 +1,742 @@ +//! The CacheGen v1 segment container and its deterministic CPU encoder and +//! decoder. +//! +//! One container holds exactly one self-contained segment tile: affine +//! 4-bit quantization, token-axis delta decorrelation, static byte-rANS. +//! The container carries the calibration values and the 16-entry symbol +//! histogram (the CDF metadata), so decoding is self-contained and the +//! rANS table is bit-identical on both sides by construction. +//! +//! The calibration digest binds the calibration and tile shape into the +//! segment identity: a lossy lookup only matches entries calibrated +//! identically, and can never satisfy an exact lookup (the +//! `CodecClass::Lossy` contract on the v4 per-segment identity). +//! +//! Determinism contract: encoding is a pure function of the input bytes. +//! No wall clock, no map iteration, no float reassociation. The GPU work +//! in the later CubeCL slices must produce byte-identical containers. + +use anyhow::{Result, anyhow, bail}; +use skippy_protocol::binary::{f16_bits_to_f32, f32_to_f16_bits}; + +use super::rans::{RansDecoder, RansEncoder, SymbolTable}; +use super::reference::{self, Calibration, TOKEN_COUNT}; +use crate::l3::{CodecClass, SegmentCodecIdentity}; + +/// Segment codec name stamped into the per-segment identity. +pub const CACHEGEN_CODEC_NAME: &str = "cachegen"; +/// Container format version. +pub const CACHEGEN_CODEC_VERSION: u32 = 1; +/// Container magic: CacheGen, version 1. +const MAGIC: [u8; 4] = *b"CGv1"; +/// Fixed prefix: 4 magic + 2 dims + 4 rows + 4 min bits + 4 scale bits + +/// 4 stream length + 2 reserved. +const FIXED_HEADER_LEN: usize = 24; +/// The histogram travels as `TOKEN_COUNT` little-endian `u32` counts. +/// Wide counts keep the format honest at real tile sizes: a 4096x128 tile +/// is 524,288 symbols and would overflow a `u16` count, and CGv1 is one +/// format for every tile the codec can encode. +const HISTOGRAM_LEN: usize = TOKEN_COUNT * 4; +/// Total bytes before the rANS stream. +const HEADER_LEN: usize = FIXED_HEADER_LEN + HISTOGRAM_LEN; +/// Hard format ceiling on the tile one container may declare, in values +/// (`2^24` = 16,777,216 values = 32 MiB decoded f16). Bounds the decode +/// working set — symbols + f32 materialization + f16 output is ~7 bytes +/// per value, so this caps the worst case near 112 MiB instead of the +/// multi-GiB bomb an unbounded product permits — and is enforced +/// symmetrically by [`encode_f16_segment`] and [`parse_container`]. The +/// measured 4096x128 tile is 524,288 values, so real tiles sit 32x under +/// the ceiling; raising it is a format-version decision, not a per-parse +/// judgment call. A later slice additionally binds decode to the +/// manifest's per-segment `decoded_len`. +pub const MAX_DECODED_VALUES: usize = 1 << 24; + +/// The per-segment identity a CacheGen segment carries: lossy, calibrated. +pub fn segment_identity(decoded_len: u64, calibration_digest: String) -> SegmentCodecIdentity { + SegmentCodecIdentity { + name: CACHEGEN_CODEC_NAME.to_string(), + version: CACHEGEN_CODEC_VERSION, + class: CodecClass::Lossy, + decoded_len, + calibration_digest: Some(calibration_digest), + } +} + +/// BLAKE3 digest binding the calibration parameters to the tile shape, in +/// f32 bit-exact form. Two entries decode against each other only if this +/// digest matches. +pub fn calibration_digest(calibration: &Calibration, dims: usize) -> Result { + let dims = u16::try_from(dims).map_err(|_| anyhow!("dims exceed container field width"))?; + let mut hasher = blake3::Hasher::new(); + hasher.update(b"skippy-cachegen-calibration-v1"); + hasher.update(&dims.to_le_bytes()); + hasher.update(&calibration.min_bits.to_le_bytes()); + hasher.update(&calibration.scale_bits.to_le_bytes()); + Ok(hasher.finalize().to_hex().to_string()) +} + +/// Encodes one segment of little-endian f16 KV values as a CacheGen v1 +/// container. `dims` is the number of values per token along the token +/// axis (the delta stride); the value count must be a non-zero multiple +/// of it and each of its parts must fit the container's fixed fields. +pub fn encode_f16_segment(raw_segment: &[u8], dims: usize) -> Result> { + if dims == 0 || raw_segment.is_empty() || !raw_segment.len().is_multiple_of(2) { + bail!("segment must be a non-empty run of f16 pairs"); + } + // Same ceiling the parser enforces, applied before any decode-sized + // allocation on either side: encode and decode agree on what the + // format bounds, so a tile that would be undecodable is rejected at + // entry — and before the f32 materialization below, not after it. + checked_tile_len(raw_segment.len() / 2)?; + let values: Vec = raw_segment + .as_chunks::<2>() + .0 + .iter() + .map(|bytes| f16_bits_to_f32(u16::from_le_bytes(*bytes))) + .collect(); + if !values.len().is_multiple_of(dims) { + bail!( + "segment shape mismatch: {} values are not rows of {dims}", + values.len() + ); + } + let rows = u32::try_from(values.len() / dims) + .map_err(|_| anyhow!("segment row count exceeds container field width"))?; + let dims16 = u16::try_from(dims).map_err(|_| anyhow!("dims exceed container field width"))?; + + let calibration = reference::calibrate(&values)?; + let mut symbols = reference::quantize(&calibration, &values)?; + reference::delta_encode(&mut symbols, dims)?; + + let mut histogram = vec![0u32; TOKEN_COUNT]; + for &symbol in &symbols { + histogram[usize::from(symbol)] += 1; + } + let freqs = reference::histogram_to_freqs(&histogram, symbols.len())?; + let table = SymbolTable::from_freqs(&freqs).expect("histogram_to_freqs yields a valid table"); + let mut encoder = RansEncoder::new(); + // rANS encodes in reverse stream order. + for &symbol in symbols.iter().rev() { + encoder.put(&table, usize::from(symbol)); + } + let stream = encoder.finish(); + + let mut out = Vec::with_capacity(HEADER_LEN + stream.len()); + out.extend_from_slice(&MAGIC); + out.extend_from_slice(&dims16.to_le_bytes()); + out.extend_from_slice(&rows.to_le_bytes()); + out.extend_from_slice(&calibration.min_bits.to_le_bytes()); + out.extend_from_slice(&calibration.scale_bits.to_le_bytes()); + let stream_len = u32::try_from(stream.len()) + .map_err(|_| anyhow!("encoded stream exceeds container field width"))?; + out.extend_from_slice(&stream_len.to_le_bytes()); + out.extend_from_slice(&[0u8; 2]); + // Counts are `u32` throughout the encoder; the write is direct. + for &count in &histogram { + out.extend_from_slice(&count.to_le_bytes()); + } + out.extend_from_slice(&stream); + Ok(out) +} + +/// Reconstructs the segment's little-endian f16 bytes from a container. +/// The output is within the quantization error of the original segment +/// and has exactly the declared decoded length. +pub fn decode_f16_segment(payload: &[u8]) -> Result> { + let (header, histogram, stream) = parse_container(payload)?; + let count = header.count; + let freqs = reference::histogram_to_freqs(&histogram, count)?; + let table = SymbolTable::from_freqs(&freqs).expect("histogram_to_freqs yields a valid table"); + + let mut decoder = RansDecoder::new(stream) + .ok_or_else(|| anyhow!("cachegen stream shorter than its initial state"))?; + // The reservation uses the parse-validated count (checked product, + // under the format ceiling), so a hostile header cannot inflate it; + // the decoded byte size is computed with checked math as well. + let mut symbols = Vec::with_capacity(count); + for _ in 0..count { + let symbol = decoder + .get(&table) + .ok_or_else(|| anyhow!("cachegen stream exhausted before tile completed"))?; + symbols.push(symbol as u8); + } + reference::delta_decode(&mut symbols, header.dims)?; + let values = reference::dequantize(&header.calibration, &symbols); + let decoded_bytes = decoded_byte_len(count)?; + let mut out = Vec::with_capacity(decoded_bytes); + for value in values { + out.extend_from_slice(&f32_to_f16_bits(value).to_le_bytes()); + } + Ok(out) +} + +/// Checked byte length of a decoded tile: `values * 2` (little-endian +/// f16), computed without trusting the caller-provided count. +fn decoded_byte_len(count: usize) -> Result { + count + .checked_mul(2) + .ok_or_else(|| anyhow!("cachegen decoded size exceeds the address space")) +} + +/// Admission check for the format ceiling, shared by both codec +/// directions: a tile of `count` values is representable only at or +/// under [`MAX_DECODED_VALUES`]. `parse_container` and +/// `encode_f16_segment` both route through this so the bound is +/// enforced symmetrically by construction — and so the boundary is +/// unit-testable without allocating a real over-ceiling tile. +fn checked_tile_len(count: usize) -> Result<()> { + if count > MAX_DECODED_VALUES { + bail!( + "cachegen tile declares {count} values, above the format ceiling of {MAX_DECODED_VALUES}" + ); + } + Ok(()) +} + +/// Value count a container decodes to (`rows * dims`), for capability +/// negotiation against a segment's declared decoded length before any +/// decode work happens. The count is the parse-validated, bounded +/// product — never an unchecked multiply of header fields. +pub fn decoded_value_count(payload: &[u8]) -> Result { + let (header, _, _) = parse_container(payload)?; + Ok(header.count) +} + +/// Calibration and tile shape from a container's header, without touching +/// the stream. Lets a caller digest the calibration into its lookup +/// identity before committing to a decode. +pub fn container_calibration(payload: &[u8]) -> Result<(Calibration, usize)> { + let (header, _, _) = parse_container(payload)?; + Ok((header.calibration, header.dims)) +} + +struct ContainerHeader { + dims: usize, + /// Parse-validated `rows * dims`: checked product, under + /// [`MAX_DECODED_VALUES`], and equal to the histogram total. + count: usize, + calibration: Calibration, +} + +fn parse_container(payload: &[u8]) -> Result<(ContainerHeader, Vec, &[u8])> { + if payload.len() < HEADER_LEN { + bail!("cachegen container shorter than its fixed header"); + } + if payload[0..4] != MAGIC { + bail!("not a cachegen container (bad magic)"); + } + if payload[FIXED_HEADER_LEN - 2..FIXED_HEADER_LEN] != [0u8; 2] { + bail!("cachegen container reserved bytes are not zero"); + } + let dims = u16::from_le_bytes([payload[4], payload[5]]) as usize; + let rows = u32::from_le_bytes([payload[6], payload[7], payload[8], payload[9]]) as usize; + let min_bits = u32::from_le_bytes(payload[10..14].try_into().expect("4 bytes")); + let scale_bits = u32::from_le_bytes(payload[14..18].try_into().expect("4 bytes")); + let stream_len = u32::from_le_bytes(payload[18..22].try_into().expect("4 bytes")) as usize; + if dims == 0 || rows == 0 { + bail!("cachegen container declares an empty tile"); + } + // The declared tile must be a real shape, not a hostile product: the + // checked count bounds the decoded-size math below, and the format + // ceiling bounds it further. Without these, a 92-byte payload could + // claim billions of decoded values and drive the allocation at decode + // time. + let count = rows + .checked_mul(dims) + .ok_or_else(|| anyhow!("cachegen tile shape overflows the address space"))?; + checked_tile_len(count)?; + // Calibration crosses the trust boundary as raw f32 bits; it must be + // a usable affine map before anything derives from it. NaN/inf poison + // every dequantized value, and a negative scale inverts the quantizer + // ring. scale == 0 stays legal: it is the flat-tile encoding. + let min = f32::from_bits(min_bits); + let scale = f32::from_bits(scale_bits); + if !min.is_finite() || !scale.is_finite() { + bail!("cachegen calibration carries a non-finite value"); + } + if scale < 0.0 { + bail!("cachegen calibration scale is negative"); + } + let histogram: Vec = payload[FIXED_HEADER_LEN..HEADER_LEN] + .as_chunks::<4>() + .0 + .iter() + .map(|bytes| u32::from_le_bytes(*bytes)) + .collect(); + // The histogram must account for exactly the declared tile. A total + // that disagrees with `rows * dims` is a corrupt (or hostile) header: + // accepting it would build a CDF for a tile that never existed and + // drive unbounded normalization repair at decode time. + let histogram_total: u64 = histogram.iter().map(|&count| u64::from(count)).sum(); + let expected_total = (rows as u64) * (dims as u64); + if histogram_total != expected_total { + bail!( + "cachegen histogram totals {histogram_total} but the tile declares {expected_total} symbols" + ); + } + if payload.len() != HEADER_LEN + stream_len { + bail!( + "cachegen container length {} disagrees with its declared stream length {stream_len}", + payload.len() + ); + } + Ok(( + ContainerHeader { + dims, + count, + calibration: Calibration { + min_bits, + scale_bits, + }, + }, + histogram, + &payload[HEADER_LEN..], + )) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Deterministic xorshift* so fixtures never depend on an RNG crate. + struct Xorshift(u64); + + impl Xorshift { + fn next_unit(&mut self) -> f32 { + self.0 ^= self.0 >> 12; + self.0 ^= self.0 << 25; + self.0 ^= self.0 >> 27; + (self.0.wrapping_mul(0x2545F4914F6CDD1D) >> 40) as f32 / 16_777_216.0 + } + } + + /// A smooth KV-like tile: slowly varying signal plus tiny jitter, the + /// shape the CacheGen paper's gains come from. + fn smooth_tile(rows: usize, dims: usize, seed: u64) -> Vec { + let mut rng = Xorshift(seed); + let mut out = Vec::with_capacity(rows * dims * 2); + for row in 0..rows { + for column in 0..dims { + let phase = (row * dims + column) as f32; + let value = (phase * 0.01).sin() * 0.4 + 0.5; + let jitter = (rng.next_unit() - 0.5) * 0.01; + out.extend_from_slice(&f32_to_f16_bits(value + jitter).to_le_bytes()); + } + } + out + } + + #[test] + fn smooth_segment_compresses_and_round_trips_within_quantization_error() { + let dims = 16; + let raw = smooth_tile(128, dims, 0xDECAFBAD); + let encoded = encode_f16_segment(&raw, dims).expect("encode"); + assert!( + encoded.len() < raw.len(), + "smooth KV must compress: {} vs {} bytes", + encoded.len(), + raw.len() + ); + let decoded = decode_f16_segment(&encoded).expect("decode"); + assert_eq!(decoded.len(), raw.len()); + // Per-value error is bounded by the quantization step plus the + // final f16 rounding; smooth data at 4 bits stays well inside one + // step. + for (original, restored) in raw + .as_chunks::<2>() + .0 + .iter() + .zip(decoded.as_chunks::<2>().0) + { + let original = f16_bits_to_f32(u16::from_le_bytes(*original)); + let restored = f16_bits_to_f32(u16::from_le_bytes(*restored)); + assert!( + (original - restored).abs() < 0.05, + "{original} rebuilt as {restored}" + ); + } + } + + #[test] + fn encoding_is_deterministic() { + let dims = 8; + let raw = smooth_tile(64, dims, 42); + let first = encode_f16_segment(&raw, dims).expect("encode"); + let second = encode_f16_segment(&raw, dims).expect("encode"); + assert_eq!(first, second); + } + + #[test] + fn calibration_digest_tracks_shape_and_values() { + let raw = smooth_tile(32, 8, 7); + let encoded = encode_f16_segment(&raw, 8).expect("encode"); + let (calibration, dims) = container_calibration(&encoded).expect("header"); + let digest = calibration_digest(&calibration, dims).expect("digest"); + assert_eq!(digest.len(), 64); + // Different shape: same values, different digest. + let wider = encode_f16_segment(&raw, 16).expect("encode"); + let (wider_calibration, wider_dims) = container_calibration(&wider).expect("header"); + assert_ne!( + digest, + calibration_digest(&wider_calibration, wider_dims).expect("digest") + ); + // Different values: different calibration, different digest. + let other = smooth_tile(32, 8, 8); + let (other_calibration, other_dims) = + container_calibration(&encode_f16_segment(&other, 8).expect("encode")).expect("header"); + assert_ne!( + digest, + calibration_digest(&other_calibration, other_dims).expect("digest") + ); + } + + #[test] + fn segment_identity_is_lossy_and_namespaced() { + let identity = segment_identity(1024, "digest".to_string()); + assert_eq!(identity.name, CACHEGEN_CODEC_NAME); + assert_eq!(identity.version, CACHEGEN_CODEC_VERSION); + assert_eq!(identity.class, CodecClass::Lossy); + assert_eq!(identity.decoded_len, 1024); + assert!(!identity.is_supported()); + assert!(identity.is_self_consistent(2048)); + assert_ne!( + identity.name, + crate::l3::CODEC_RAW, + "must stay out of the raw namespace" + ); + } + + /// scama re-review blocker 1, reproduced exactly: a NaN calibration + /// minimum, a negative scale, and a 92-byte container claiming + /// billions of decoded values must all be refused at the boundary — + /// `container_calibration` and `decoded_value_count` are capability + /// probes, so they must reject, not just `decode_f16_segment`. + #[test] + fn hostile_calibration_and_shapes_are_refused_at_the_boundary() { + let dims = 16usize; + let raw = smooth_tile(4, dims, 5); + let encoded = encode_f16_segment(&raw, dims).expect("encode"); + + let mutate = |position: usize, bytes: [u8; 4]| { + let mut payload = encoded.clone(); + payload[position..position + 4].copy_from_slice(&bytes); + payload + }; + // Header layout: min at 10, scale at 14. + let nan_min = mutate(10, f32::NAN.to_bits().to_le_bytes()); + let inf_min = mutate(10, f32::INFINITY.to_bits().to_le_bytes()); + let nan_scale = mutate(14, f32::NAN.to_bits().to_le_bytes()); + let negative_scale = mutate(14, (-1.0f32).to_bits().to_le_bytes()); + for (name, hostile) in [ + ("NaN min", &nan_min), + ("inf min", &inf_min), + ("NaN scale", &nan_scale), + ("negative scale", &negative_scale), + ] { + assert!( + container_calibration(hostile).is_err(), + "{name} must be refused by container_calibration" + ); + assert!( + decoded_value_count(hostile).is_err(), + "{name} must be refused by decoded_value_count" + ); + assert!( + decode_f16_segment(hostile).is_err(), + "{name} must be refused by decode_f16_segment" + ); + } + + // dims=16, rows=u32::MAX, every histogram entry u32::MAX, a + // 4-byte stream: the exact shape that used to decode as + // 68,719,476,720 values. + let mut bomb = vec![0u8; HEADER_LEN + 4]; + bomb[0..4].copy_from_slice(&MAGIC); + bomb[4..6].copy_from_slice(&(dims as u16).to_le_bytes()); + bomb[6..10].copy_from_slice(&u32::MAX.to_le_bytes()); + bomb[18..22].copy_from_slice(&4u32.to_le_bytes()); + // Every histogram entry at u32::MAX: little-endian u32::MAX is + // four 0xFF bytes, so a fill covers all TOKEN_COUNT entries. + bomb[FIXED_HEADER_LEN..HEADER_LEN].fill(0xFF); + assert!(decoded_value_count(&bomb).is_err()); + assert!(decode_f16_segment(&bomb).is_err()); + // And a merely huge-but-well-formed total still exceeds the + // format ceiling even when the histogram would have to lie to + // support it. + let mut huge = encoded.clone(); + huge[6..10].copy_from_slice(&(MAX_DECODED_VALUES as u32).to_le_bytes()); + assert!(decoded_value_count(&huge).is_err()); + } + + /// A minimal well-formed header for a `(rows x dims)` tile: correct + /// magic, zero reserved bytes, a histogram whose total supports the + /// declared shape, and a 4-byte stream. The stream is not decodable — + /// callers assert admission (`decoded_value_count`, + /// `container_calibration`) or parse-time rejection, never a full + /// decode. + fn crafted_header(dims: usize, rows: u32) -> Vec { + let mut payload = vec![0u8; FIXED_HEADER_LEN]; + payload[0..4].copy_from_slice(&MAGIC); + payload[4..6].copy_from_slice(&(dims as u16).to_le_bytes()); + payload[6..10].copy_from_slice(&rows.to_le_bytes()); + payload[18..22].copy_from_slice(&4u32.to_le_bytes()); + // Every entry equal to count / TOKEN_COUNT sums to exactly + // rows * dims. Written the way the encoder writes histograms. + let per_token = (u64::from(rows) * dims as u64 / TOKEN_COUNT as u64) as u32; + for _ in 0..TOKEN_COUNT { + payload.extend_from_slice(&per_token.to_le_bytes()); + } + payload.extend_from_slice(&[0u8; 4]); + payload + } + + /// scama's remaining blocker, bounded at both edges of the ceiling. + /// Exactly `MAX_DECODED_VALUES` (65536 x 256) must be admitted by the + /// boundary probes; exactly one value more (16777217 == 97 x 172961, + /// a real shape, not a truncation artifact) must be refused by every + /// probe before any decode-sized allocation happens. + #[test] + fn ceiling_boundary_is_inclusive_and_rejects_exactly_one_more() { + // The inclusive edge: largest tile the format admits. + let boundary = crafted_header(256, (MAX_DECODED_VALUES / 256) as u32); + assert_eq!( + decoded_value_count(&boundary).expect("boundary tile is admitted"), + MAX_DECODED_VALUES + ); + assert!(container_calibration(&boundary).is_ok()); + + // One value above: refused everywhere, naming the ceiling. + // 16_777_217 = 97 * 172_961, so rows x dims hits MAX + 1 exactly. + assert_eq!(97usize * 172_961usize, MAX_DECODED_VALUES + 1); + let over = crafted_header(97, 172_961); + assert_eq!( + 97usize.checked_mul(172_961).expect("shape"), + MAX_DECODED_VALUES + 1 + ); + for (name, refused) in [ + ("decoded_value_count", decoded_value_count(&over).is_err()), + ( + "container_calibration", + container_calibration(&over).is_err(), + ), + ("decode_f16_segment", decode_f16_segment(&over).is_err()), + ] { + assert!(refused, "{name} must refuse MAX + 1 values"); + } + let error = decoded_value_count(&over).expect_err("over-ceiling tile"); + assert!( + error.to_string().contains("format ceiling"), + "rejection should name the format ceiling: {error}" + ); + } + + /// The encoder enforces the same ceiling before any work — the + /// symmetric half of the blocker. The direct probe covers the exact + /// +1 edge without materializing the tile; the end-to-end call proves + /// the encoder actually routes through it (33.5 MB of f16 input is + /// refused before the f32 conversion, so the test stays cheap). + #[test] + fn encode_refuses_over_ceiling_tiles_symmetrically() { + assert!(checked_tile_len(MAX_DECODED_VALUES).is_ok()); + let error = + checked_tile_len(MAX_DECODED_VALUES + 1).expect_err("one value above the ceiling"); + assert!( + error.to_string().contains("format ceiling"), + "rejection should name the format ceiling: {error}" + ); + + let dims = 16usize; + let over_ceiling_input = vec![0u8; 2 * (MAX_DECODED_VALUES + 1)]; + let error = encode_f16_segment(&over_ceiling_input, dims) + .expect_err("over-ceiling input must not encode"); + assert!( + error.to_string().contains("format ceiling"), + "encode must refuse before any work: {error}" + ); + } + + /// The flat tile (scale == 0) is a legal encoding and must stay one + /// after calibration validation. + #[test] + fn flat_tile_calibration_remains_legal() { + let values = vec![0.5f32; 4 * 8]; + let mut raw = Vec::new(); + for value in &values { + raw.extend_from_slice(&f32_to_f16_bits(*value).to_le_bytes()); + } + let encoded = encode_f16_segment(&raw, 8).expect("encode flat tile"); + let (calibration, _) = container_calibration(&encoded).expect("flat calibration"); + assert_eq!(calibration.scale_bits, 0.0f32.to_bits()); + let decoded = decode_f16_segment(&encoded).expect("decode flat tile"); + assert_eq!(decoded.len(), raw.len()); + } + + #[test] + fn corrupt_containers_are_rejected_cleanly() { + let dims = 8; + let raw = smooth_tile(16, dims, 9); + let encoded = encode_f16_segment(&raw, dims).expect("encode"); + + let mut bad_magic = encoded.clone(); + bad_magic[0] = b'X'; + assert!(decode_f16_segment(&bad_magic).is_err()); + + let mut reserved = encoded.clone(); + reserved[FIXED_HEADER_LEN - 1] = 1; + assert!(decode_f16_segment(&reserved).is_err()); + + assert!(decode_f16_segment(&encoded[..encoded.len() - 1]).is_err()); + assert!(decode_f16_segment(&encoded[..HEADER_LEN - 1]).is_err()); + assert!(decode_f16_segment(&[]).is_err()); + + let mut length_lie = encoded.clone(); + length_lie[18] ^= 0xff; + assert!(decode_f16_segment(&length_lie).is_err()); + + // A corrupt histogram count must not silently decode: a zeroed + // symbol frequency can leave symbols undecodable, and the derived + // table must refuse to lie about totals. + let mut bad_histogram = encoded.clone(); + bad_histogram[FIXED_HEADER_LEN] ^= 0xff; + assert!(decode_f16_segment(&bad_histogram).is_err()); + + // A histogram that plausibly formats but does not sum to the + // declared tile is refused at parse, before any table is built: + // this is the bounded-normalization contract (scama blocker 2). + let mut total_lie = encoded.clone(); + let lie_position = FIXED_HEADER_LEN + 4 * 3; + let current = u32::from_le_bytes( + total_lie[lie_position..lie_position + 4] + .try_into() + .expect("4 bytes"), + ); + total_lie[lie_position..lie_position + 4] + .copy_from_slice(¤t.wrapping_add(1).to_le_bytes()); + let error = decode_f16_segment(&total_lie).expect_err("histogram/shape disagreement"); + assert!( + error.to_string().contains("histogram totals"), + "rejection should name the histogram total: {error}" + ); + } + + /// The exact shape scama's review measured and the spike reports: + /// 4096 rows x 128 dims, end to end through the real CGv1 container, + /// not a bypass. This is the regression for the u16 histogram cap. + #[test] + fn the_measured_4096x128_tile_round_trips_through_the_container() { + let dims = 128; + let rows = 4096; + let raw = smooth_tile(rows, dims, 0xC0FFEE); + let encoded = encode_f16_segment(&raw, dims).expect("encode a full-size tile"); + assert_eq!( + decoded_value_count(&encoded).expect("count"), + rows * dims, + "container must declare the full measured tile" + ); + let decoded = decode_f16_segment(&encoded).expect("decode a full-size tile"); + assert_eq!(decoded.len(), raw.len()); + let max_error: f32 = raw + .as_chunks::<2>() + .0 + .iter() + .zip(decoded.as_chunks::<2>().0) + .map(|(original, restored)| { + (f16_bits_to_f32(u16::from_le_bytes(*original)) + - f16_bits_to_f32(u16::from_le_bytes(*restored))) + .abs() + }) + .fold(0.0, f32::max); + assert!( + max_error < 0.05, + "full-size tile decoded outside quantization error: {max_error}" + ); + } + + #[test] + fn shape_mismatch_is_refused_before_encoding() { + let raw = smooth_tile(4, 8, 3); + assert!(encode_f16_segment(&raw, 7).is_err()); + assert!(encode_f16_segment(&raw, 0).is_err()); + assert!(encode_f16_segment(&raw[..raw.len() - 1], 8).is_err()); + } + + #[test] + fn decoded_value_count_reports_the_tile_size() { + let dims = 8; + let raw = smooth_tile(24, dims, 11); + let encoded = encode_f16_segment(&raw, dims).expect("encode"); + assert_eq!(decoded_value_count(&encoded).expect("count"), 24 * dims); + } + + #[test] + fn random_noise_does_not_expand_but_still_round_trips() { + // Incompressible input: the container must not blow up (bounded + // overhead) and must still decode to quantization-accurate data. + let dims = 8; + let mut rng = Xorshift(0x5EED_5EED); + let raw: Vec = (0..64 * dims) + .flat_map(|_| f32_to_f16_bits(rng.next_unit()).to_le_bytes()) + .collect(); + let encoded = encode_f16_segment(&raw, dims).expect("encode"); + assert!( + encoded.len() <= raw.len() + 256, + "noise must not expand the payload materially: {} vs {}", + encoded.len(), + raw.len() + ); + let decoded = decode_f16_segment(&encoded).expect("decode"); + assert_eq!(decoded.len(), raw.len()); + } +} + +#[cfg(test)] +mod store_contract_tests { + use super::*; + use crate::l3::{HandoffManifest, HandoffSegmentRef, HandoffSegmentStore}; + + /// The full #1652 contract, end to end: a real CacheGen-encoded segment + /// carries a lossy identity the store refuses at commit. No lossy entry + /// can enter the exact pipeline while raw is the only supported class. + #[test] + fn a_real_cachegen_segment_never_commits_to_the_store() { + let dims = 8; + let mut raw = Vec::new(); + for row in 0..16 { + for column in 0..dims { + let value = ((row * dims + column) as f32 * 0.05).sin() * 0.3 + 0.5; + raw.extend_from_slice(&f32_to_f16_bits(value).to_le_bytes()); + } + } + let encoded = encode_f16_segment(&raw, dims).expect("encode"); + let (calibration, container_dims) = container_calibration(&encoded).expect("header"); + let digest = calibration_digest(&calibration, container_dims).expect("digest"); + let identity = segment_identity(raw.len() as u64, digest); + assert!(!identity.is_supported()); + + let root = std::env::temp_dir() + .join("skippy-cachegen-store-tests") + .join(format!("reject-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&root); + let store = HandoffSegmentStore::open(&root, 0).expect("open store"); + let stored = store.put_segment(&encoded).expect("put segment"); + let mut manifest = HandoffManifest::new("blake3:cachegen".to_string(), "full-state".into()); + manifest.segments.push(HandoffSegmentRef { + index: 0, + offset: 0, + bytes: encoded.len() as u64, + digest: stored.digest.clone(), + codec_identity: Some(identity), + meta_json: None, + }); + manifest.total_bytes = encoded.len() as u64; + manifest.payload_digest = "blake3:cachegen-payload".to_string(); + let error = store + .commit(&manifest) + .expect_err("a lossy segment must not commit to the exact pipeline"); + assert!( + error.to_string().contains("cachegen"), + "commit error should name the codec: {error}" + ); + std::fs::remove_dir_all(&root).ok(); + } +} diff --git a/crates/skippy-cache/src/cachegen/fixtures/generate_lmcache_compat.py b/crates/skippy-cache/src/cachegen/fixtures/generate_lmcache_compat.py new file mode 100644 index 0000000000..78d87a4d1e --- /dev/null +++ b/crates/skippy-cache/src/cachegen/fixtures/generate_lmcache_compat.py @@ -0,0 +1,273 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Generate language-neutral compatibility fixtures for the Rust CacheGen port. + +This is a scalar transcription of LMCache revision +b5d109ea99a89b4d8a670ee4fc2e8cb76411ee5c. The relevant upstream files are: + +* lmcache/storage_backend/serde/cachegen_encoder.py +* lmcache/storage_backend/serde/cachegen_decoder.py +* csrc/cuda/ac_enc.cu +* csrc/cuda/ac_dec.cu +* csrc/cuda/cal_cdf.cu + +Run from this directory with Python 3.11 or newer. It needs no third-party +packages and deliberately shares no Rust implementation code. +""" + +from pathlib import Path +import math +import struct + +MAGIC = b"LCG1" +FIXTURE_MAGIC = b"LCFX" +MAX_BINS = 32 +CDF_LEN = MAX_BINS + 1 +CDF_TOTAL = 1 << 16 +QUARTER = 0x40000000 +HALF = 0x80000000 +THREE_QUARTERS = 0xC0000000 +MASK32 = 0xFFFFFFFF + + +def f32(value: float) -> float: + return struct.unpack(" bytes: + return struct.pack(" float: + return struct.unpack(" None: + self.register = 0 + self.bits = 0 + self.output = bytearray() + + def repeated(self, bit: int, count: int) -> None: + while count: + take = min(count, 32 - self.bits) + self.register = (self.register << take) & MASK32 + if bit: + self.register |= MASK32 if take == 32 else (1 << take) - 1 + self.bits += take + count -= take + if self.bits == 32: + self.output.extend(struct.pack(">I", self.register)) + self.register = 0 + self.bits = 0 + + def append(self, bit: int, pending: int) -> None: + self.repeated(bit, 1) + self.repeated(1 - bit, pending) + + def finish(self) -> bytes: + if self.bits: + self.register = (self.register << (32 - self.bits)) & MASK32 + self.output.extend(struct.pack(">I", self.register)[: (self.bits + 7) // 8]) + return bytes(self.output) + + +class BitReader: + def __init__(self, stream: bytes) -> None: + self.stream = stream + self.bit = 0 + + def read_bit(self) -> int: + byte = self.stream[self.bit // 8] if self.bit // 8 < len(self.stream) else 0 + value = (byte >> (7 - self.bit % 8)) & 1 + self.bit += 1 + return value + + def read_u32(self) -> int: + value = 0 + for _ in range(32): + value = ((value << 1) | self.read_bit()) & MASK32 + return value + + +def encode_channel(symbols, rows, channels, channel, cdf) -> bytes: + low = 0 + high = MASK32 + pending = 0 + writer = BitWriter() + for row in range(rows): + symbol = symbols[row * channels + channel] + span = high - low + 1 + c_low = cdf[symbol] + c_high = CDF_TOTAL if symbol == MAX_BINS - 1 else cdf[symbol + 1] + high = ((low - 1) + ((span * c_high) >> 16)) & MASK32 + low = (low + ((span * c_low) >> 16)) & MASK32 + while True: + if high < HALF: + writer.append(0, pending) + pending = 0 + elif low >= HALF: + writer.append(1, pending) + pending = 0 + elif low >= QUARTER and high < THREE_QUARTERS: + pending += 1 + low = (low << 1) & 0x7FFFFFFF + high = ((high << 1) | 0x80000001) & MASK32 + continue + else: + break + low = (low << 1) & MASK32 + high = ((high << 1) | 1) & MASK32 + pending += 1 + writer.append(1 if low >= QUARTER else 0, pending) + return writer.finish() + + +def decode_channel(stream, rows, channels, channel, cdf, output) -> None: + reader = BitReader(stream) + low = 0 + high = MASK32 + value = reader.read_u32() + for row in range(rows): + span = high - low + 1 + count = (((value - low + 1) * CDF_TOTAL) - 1) // span + left = 0 + right = MAX_BINS + while left + 1 < right: + middle = (left + right) // 2 + if cdf[middle] < count: + left = middle + elif cdf[middle] > count: + right = middle + else: + left = middle + break + symbol = left + output[row * channels + channel] = symbol + if row + 1 == rows: + break + c_low = cdf[symbol] + c_high = CDF_TOTAL if symbol == MAX_BINS - 1 else cdf[symbol + 1] + high = ((low - 1) + ((span * c_high) >> 16)) & MASK32 + low = (low + ((span * c_low) >> 16)) & MASK32 + while True: + if low >= HALF or high < HALF: + low = (low << 1) & MASK32 + high = ((high << 1) | 1) & MASK32 + value = ((value << 1) | reader.read_bit()) & MASK32 + elif low >= QUARTER and high < THREE_QUARTERS: + low = ((low << 1) & 0x7FFFFFFF) & MASK32 + high = ((high << 1) | 0x80000001) & MASK32 + value = (value - QUARTER) & MASK32 + value = ((value << 1) | reader.read_bit()) & MASK32 + else: + break + + +def quantize(values, rows, channels, bins): + center = f32(bins // 2 - 1) + symbols = [] + maxes = [] + for row in range(rows): + values_row = values[row * channels : (row + 1) * channels] + maximum = max(abs(value) for value in values_row) + maxes.append(maximum) + if maximum == 0.0: + symbols.extend([int(center)] * channels) + continue + factor = f32(center / maximum) + for value in values_row: + scaled = f32(value * factor) + shifted = f32(scaled + center) + symbols.append(max(0, min(int(center * 2), round(shifted)))) + return symbols, maxes + + +def calculate_cdf(symbols, rows, channels): + cdfs = [] + for channel in range(channels): + histogram = [0] * CDF_LEN + for row in range(rows): + histogram[symbols[row * channels + channel] + 1] += 1 + running = 0 + for index in range(1, CDF_LEN): + count = histogram[index] + histogram[index] += running + running += count + cdf = [((0xFFFF - MAX_BINS) * count // running) + index for index, count in enumerate(histogram)] + assert cdf[0] == 0 and cdf[-1] == 0xFFFF + assert all(left < right for left, right in zip(cdf, cdf[1:])) + cdfs.append(cdf) + return cdfs + + +def encode(raw: bytes, rows: int, channels: int, bins: int) -> bytes: + values = [f32(f16_value(raw[index : index + 2])) for index in range(0, len(raw), 2)] + symbols, maxes = quantize(values, rows, channels, bins) + cdfs = calculate_cdf(symbols, rows, channels) + streams = [encode_channel(symbols, rows, channels, channel, cdfs[channel]) for channel in range(channels)] + payload = bytearray(MAGIC) + payload.extend(bytes((bins, 0))) + payload.extend(struct.pack(" bytes: + assert payload[:4] == MAGIC and payload[5] == 0 + bins = payload[4] + rows, channels, stream_len = struct.unpack(" bytes: + raw = bytearray() + for row in range(rows): + for channel in range(channels): + phase = f32(f32(row * channels + channel) * f32(0.03125)) + value = f32(f32(math.sin(phase)) * f32(0.75)) + value = f32(value + f32(f32(channel - 3) * f32(0.015625))) + raw.extend(f16_bytes(value)) + return bytes(raw) + + +def write_fixture(path: Path, bins: int) -> None: + rows = 17 + channels = 8 + raw = input_bytes(rows, channels) + encoded = encode(raw, rows, channels, bins) + decoded = decode(encoded) + fixture = FIXTURE_MAGIC + struct.pack("4E&T46$zE5nVS0a+)S9^hk^V!91a}*91i^K*}26r zVF29U1c2N@H)vl@x{|t>_Ea~or>S$gF8zB$bVg@}nSM*3VZ!KJbRr`%ImTgQI;@6g zjnl?Vco{x{^N>MR=ppr!~Y2;PEOd!4JOXal?Z}R>1yi??GNkU z;NVMz=9bVBXl}hfIQYYSt!JYNE>|cCkzmVTB0vNPCs>z!wo((YE?-EvXwBLNRjq?r zH34^IU*4OSGxrLGhE{)hM~gK-d9 zgK;y#G!R{7nhmQoX*?+>)y`y3qI17f=~TvgW4#~$(Mxpg_xJmoh^DKI5=DcVkk;}C DRJ~v3 literal 0 HcmV?d00001 diff --git a/crates/skippy-cache/src/cachegen/fixtures/lmcache_b5d109e_bins32.bin b/crates/skippy-cache/src/cachegen/fixtures/lmcache_b5d109e_bins32.bin new file mode 100644 index 0000000000000000000000000000000000000000..bc3ac5a2d57513645a275a0e8b213366b2bc0aa4 GIT binary patch literal 1241 zcmb7@TTB#J7{?DvaD`ay0}n>JfhD?JTwHdTne!bFdjSFQfOe^f4w}?hk*SS%cfG`t zYO<-}rJxuuX}6JmSR>7B%P>*cQ$<}0V z@}Hi$x?A__tMz96px&mRHO?EoM!zv&+&1nS!^VjDm-&bJ#C&MpH-9s4nl0_agKKSKeOogYB4+5DULWbfg|;yKZ9i`r^{#s zwP&ylYgv=BMX?Rt@!_G#8D#n{LnW=E%c*Tx$0lZob$+O!ny#SMEw(7l%3k$oWYUo5 zk?4oaps;2|RraYJ(TP30kHLD2iCZYFk*#O@)QI|1+}?vo-#x>`Yfd+OsccpDtCqom zX7T@P7AK99ibAO~;{(c!{!m(O*MV(ML+NvKe!eq`LWd9CJl*A}(wnYs$lV@Z;K-@G z*m!mDIcXhV(s$u{LOZ9qZ!PXCy5+sjgS*^2R{X%H3%w<|;usI^m68T=RE$baxmR}c z;Fin($&I)Xn>Z6Pd2qkRf8sEN;a84}xQGYNaMFoIgwp{XkU}h4-tZ&XIEc+CuI6!_ zq?X5aDlL*tie!@z62c-*m#>S*)(Y;if^YMqc^PiU!gNvjIr5Kj2_sXQ*IcVEzH*9FSmOV*oKe0C9r_0}CUFLmMRw0B|rH AApigX literal 0 HcmV?d00001 diff --git a/crates/skippy-cache/src/cachegen/lmcache.rs b/crates/skippy-cache/src/cachegen/lmcache.rs new file mode 100644 index 0000000000..1ac9458f88 --- /dev/null +++ b/crates/skippy-cache/src/cachegen/lmcache.rs @@ -0,0 +1,819 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Derived from LMCache CacheGen at revision +// b5d109ea99a89b4d8a670ee4fc2e8cb76411ee5c: +// - lmcache/storage_backend/serde/cachegen_encoder.py +// - lmcache/storage_backend/serde/cachegen_decoder.py +// - csrc/cuda/ac_enc.cu +// - csrc/cuda/ac_dec.cu +// - csrc/cuda/cal_cdf.cu +// +// LMCache and Mesh-LLM are both licensed under Apache-2.0. This module keeps +// the reference arithmetic and tensor transforms explicit so compatibility +// can be checked without requiring Python, PyTorch, or a CUDA device at run +// time. + +//! A portable Rust implementation of LMCache's CacheGen tensor codec. +//! +//! LMCache serializes PyTorch objects with pickle. Skippy uses a bounded, +//! language-neutral envelope, while preserving the codec inputs and outputs: +//! per-token maximum magnitude, model/layer-selected 16- or 32-bin symmetric +//! quantization, a 33-entry per-channel CDF, and one 32-bit arithmetic-coded +//! stream per channel. The scalar oracle retains that byte-compatible LCG1 +//! representation. Device archives use LCG2, which bit-packs the same symbols +//! in token-major order for direct parallel restore. Each segment is one K or +//! V layer with at most 256 token-major rows, matching LMCache's CUDA kernel +//! limit. + +use anyhow::{Result, anyhow, bail}; +use skippy_protocol::binary::{f16_bits_to_f32, f32_to_f16_bits}; + +use crate::l3::{CodecClass, SegmentCodecIdentity}; + +pub const CACHEGEN_CODEC_NAME: &str = "lmcache-cachegen"; +pub const CACHEGEN_CODEC_VERSION: u32 = 2; +pub const MAX_TOKENS_PER_CHUNK: usize = 256; + +const MAGIC_V1: [u8; 4] = *b"LCG1"; +const MAGIC_V2: [u8; 4] = *b"LCG2"; +const HEADER_LEN: usize = 16; +const MAX_BINS: usize = 32; +const CDF_LEN: usize = MAX_BINS + 1; +const CDF_TOTAL: u32 = 1 << 16; +const QUARTER: u32 = 0x4000_0000; +const HALF: u32 = 0x8000_0000; +const THREE_QUARTERS: u32 = 0xc000_0000; +const MAX_DECODED_VALUES: usize = 1 << 24; + +#[derive(Debug)] +struct Parsed<'a> { + bins: u8, + channels: usize, + rows: usize, + bits_per_symbol: u8, + maxes: Vec, + cdfs: Vec<[u16; CDF_LEN]>, + lengths: Vec, + streams: &'a [u8], +} + +/// Geometry and bounded byte ranges validated from a portable CacheGen +/// segment. Device backends use this before accepting work from an archive. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SegmentGeometry { + pub bins: u8, + pub rows: usize, + pub channels: usize, + pub metadata_bytes: usize, + pub stream_bytes: usize, +} + +/// LMCache's generic model-family bin schedule. +/// +/// Models with fewer than ten layers use 32 bins everywhere. Larger models +/// use 32 bins for K layers 0..10 and V layers 0..2, then 16 bins. +pub fn bins_for_layer(layer_index: usize, layer_count: usize, is_key: bool) -> u8 { + if layer_count < 10 || layer_index < if is_key { 10 } else { 2 } { + 32 + } else { + 16 + } +} + +pub fn segment_identity(decoded_len: u64, codec_digest: String) -> SegmentCodecIdentity { + SegmentCodecIdentity { + name: CACHEGEN_CODEC_NAME.to_string(), + version: CACHEGEN_CODEC_VERSION, + class: CodecClass::Lossy, + decoded_len, + calibration_digest: Some(codec_digest), + } +} + +/// Encodes one token-major F16 K or V layer with LMCache CacheGen semantics. +pub fn encode_f16_segment(raw: &[u8], channels: usize, bins: u8) -> Result> { + validate_input(raw, channels, bins)?; + let rows = raw.len() / 2 / channels; + if rows > MAX_TOKENS_PER_CHUNK { + bail!("LMCache CacheGen chunks are limited to {MAX_TOKENS_PER_CHUNK} tokens, got {rows}"); + } + + let values: Vec = raw + .as_chunks::<2>() + .0 + .iter() + .map(|bytes| f16_bits_to_f32(u16::from_le_bytes(*bytes))) + .collect(); + if values.iter().any(|value| !value.is_finite()) { + bail!("LMCache CacheGen input contains a non-finite F16 value"); + } + + let (symbols, maxes) = quantize(&values, rows, channels, bins); + let cdfs = calculate_cdfs(&symbols, rows, channels)?; + let mut lengths = Vec::with_capacity(channels); + let mut streams = Vec::new(); + for (channel, cdf) in cdfs.iter().enumerate() { + let stream = arithmetic_encode_channel(&symbols, rows, channels, channel, cdf); + let length = u16::try_from(stream.len()) + .map_err(|_| anyhow!("LMCache channel stream exceeds u16 envelope field"))?; + lengths.push(length); + streams.extend_from_slice(&stream); + } + + let channels_u32 = u32::try_from(channels).map_err(|_| anyhow!("channel count exceeds u32"))?; + let rows_u16 = u16::try_from(rows).expect("row limit fits u16"); + let stream_len = u32::try_from(streams.len()) + .map_err(|_| anyhow!("LMCache segment stream exceeds u32 envelope field"))?; + let metadata_len = maxes + .len() + .checked_mul(4) + .and_then(|value| value.checked_add(channels.checked_mul(CDF_LEN * 2)?)) + .and_then(|value| value.checked_add(channels.checked_mul(2)?)) + .ok_or_else(|| anyhow!("LMCache segment metadata length overflow"))?; + let mut out = Vec::with_capacity(HEADER_LEN + metadata_len + streams.len()); + out.extend_from_slice(&MAGIC_V1); + out.push(bins); + out.push(0); + out.extend_from_slice(&rows_u16.to_le_bytes()); + out.extend_from_slice(&channels_u32.to_le_bytes()); + out.extend_from_slice(&stream_len.to_le_bytes()); + for value in maxes { + out.extend_from_slice(&value.to_bits().to_le_bytes()); + } + for cdf in &cdfs { + for value in cdf { + out.extend_from_slice(&value.to_le_bytes()); + } + } + for length in lengths { + out.extend_from_slice(&length.to_le_bytes()); + } + out.extend_from_slice(&streams); + Ok(out) +} + +/// Encodes the LMCache-quantized symbols in a token-major packed layout for +/// parallel device restore. +pub fn encode_f16_segment_packed(raw: &[u8], channels: usize, bins: u8) -> Result> { + validate_input(raw, channels, bins)?; + let rows = raw.len() / 2 / channels; + if rows > MAX_TOKENS_PER_CHUNK { + bail!("LMCache CacheGen chunks are limited to {MAX_TOKENS_PER_CHUNK} tokens, got {rows}"); + } + let values: Vec = raw + .as_chunks::<2>() + .0 + .iter() + .map(|bytes| f16_bits_to_f32(u16::from_le_bytes(*bytes))) + .collect(); + if values.iter().any(|value| !value.is_finite()) { + bail!("LMCache CacheGen input contains a non-finite F16 value"); + } + let (symbols, maxes) = quantize(&values, rows, channels, bins); + let bits_per_symbol = if bins == 16 { 4 } else { 5 }; + let packed = pack_symbols(&symbols, bits_per_symbol)?; + let channels_u32 = u32::try_from(channels).map_err(|_| anyhow!("channel count exceeds u32"))?; + let rows_u16 = u16::try_from(rows).expect("row limit fits u16"); + let packed_len = u32::try_from(packed.len()) + .map_err(|_| anyhow!("packed CacheGen segment exceeds u32 envelope field"))?; + let mut out = Vec::with_capacity(HEADER_LEN + maxes.len() * 4 + packed.len()); + out.extend_from_slice(&MAGIC_V2); + out.push(bins); + out.push(bits_per_symbol); + out.extend_from_slice(&rows_u16.to_le_bytes()); + out.extend_from_slice(&channels_u32.to_le_bytes()); + out.extend_from_slice(&packed_len.to_le_bytes()); + for value in maxes { + out.extend_from_slice(&value.to_bits().to_le_bytes()); + } + out.extend_from_slice(&packed); + Ok(out) +} + +/// Decodes a portable segment produced by either encoder. +pub fn decode_f16_segment(payload: &[u8]) -> Result> { + let parsed = parse(payload)?; + let symbols = if parsed.bits_per_symbol == 0 { + let mut symbols = vec![0u8; parsed.rows * parsed.channels]; + let mut cursor = 0usize; + for channel in 0..parsed.channels { + let end = cursor + .checked_add(parsed.lengths[channel]) + .ok_or_else(|| anyhow!("LMCache stream range overflow"))?; + let stream = parsed + .streams + .get(cursor..end) + .ok_or_else(|| anyhow!("LMCache channel stream exceeds payload"))?; + arithmetic_decode_channel( + stream, + parsed.rows, + parsed.channels, + channel, + &parsed.cdfs[channel], + &mut symbols, + )?; + cursor = end; + } + if cursor != parsed.streams.len() { + bail!("LMCache segment contains trailing stream bytes"); + } + symbols + } else { + unpack_symbols( + parsed.streams, + parsed.rows * parsed.channels, + parsed.bits_per_symbol, + )? + }; + + let center = f32::from(parsed.bins / 2 - 1); + let max_symbol = parsed.bins - 2; + let mut raw = Vec::with_capacity(symbols.len() * 2); + for (row, max) in parsed.maxes.iter().copied().enumerate() { + for channel in 0..parsed.channels { + let symbol = symbols[row * parsed.channels + channel]; + if symbol > max_symbol { + bail!( + "LMCache decoded symbol {symbol} exceeds the configured {}-bin range", + parsed.bins + ); + } + let centered = f32::from(symbol) - center; + let normalized = centered / center; + let value = normalized * max; + raw.extend_from_slice(&f32_to_f16_bits(value).to_le_bytes()); + } + } + Ok(raw) +} + +fn pack_symbols(symbols: &[u8], bits_per_symbol: u8) -> Result> { + let bit_len = symbols + .len() + .checked_mul(usize::from(bits_per_symbol)) + .ok_or_else(|| anyhow!("packed CacheGen bit length overflow"))?; + let mut packed = vec![0u8; bit_len.div_ceil(8)]; + let mask = (1u16 << bits_per_symbol) - 1; + for (index, &symbol) in symbols.iter().enumerate() { + if u16::from(symbol) > mask { + bail!("CacheGen symbol does not fit the packed width"); + } + let bit = index * usize::from(bits_per_symbol); + let byte = bit / 8; + let shift = bit % 8; + let value = u16::from(symbol) << shift; + packed[byte] |= value as u8; + if shift + usize::from(bits_per_symbol) > 8 { + packed[byte + 1] |= (value >> 8) as u8; + } + } + Ok(packed) +} + +fn unpack_symbols(packed: &[u8], count: usize, bits_per_symbol: u8) -> Result> { + let expected_bytes = count + .checked_mul(usize::from(bits_per_symbol)) + .ok_or_else(|| anyhow!("packed CacheGen bit length overflow"))? + .div_ceil(8); + if packed.len() != expected_bytes { + bail!("packed CacheGen byte length is inconsistent"); + } + let mask = (1u16 << bits_per_symbol) - 1; + let mut symbols = Vec::with_capacity(count); + for index in 0..count { + let bit = index * usize::from(bits_per_symbol); + let byte = bit / 8; + let shift = bit % 8; + let word = u16::from(packed[byte]) | u16::from(*packed.get(byte + 1).unwrap_or(&0)) << 8; + symbols.push(((word >> shift) & mask) as u8); + } + Ok(symbols) +} + +/// Validates a portable segment without allocating its decoded F16 output. +pub fn validate_f16_segment(payload: &[u8]) -> Result { + let parsed = parse(payload)?; + Ok(SegmentGeometry { + bins: parsed.bins, + rows: parsed.rows, + channels: parsed.channels, + metadata_bytes: payload.len() - parsed.streams.len(), + stream_bytes: parsed.streams.len(), + }) +} + +fn validate_input(raw: &[u8], channels: usize, bins: u8) -> Result<()> { + if !matches!(bins, 16 | 32) { + bail!("LMCache CacheGen supports 16 or 32 bins, got {bins}"); + } + if channels == 0 || raw.is_empty() || !raw.len().is_multiple_of(2) { + bail!("LMCache CacheGen requires a non-empty whole-F16 segment"); + } + let values = raw.len() / 2; + if values > MAX_DECODED_VALUES { + bail!("LMCache CacheGen segment exceeds the decoded-value ceiling"); + } + if !values.is_multiple_of(channels) { + bail!("LMCache CacheGen segment is not token-major rows of {channels} values"); + } + Ok(()) +} + +fn quantize(values: &[f32], rows: usize, channels: usize, bins: u8) -> (Vec, Vec) { + let center = f32::from(bins / 2 - 1); + let mut symbols = Vec::with_capacity(values.len()); + let mut maxes = Vec::with_capacity(rows); + for row in values.chunks_exact(channels) { + let max = row.iter().copied().map(f32::abs).fold(0.0f32, f32::max); + maxes.push(max); + if max == 0.0 { + symbols.extend(std::iter::repeat_n(center as u8, channels)); + continue; + } + let factor = center / max; + symbols.extend(row.iter().map(|value| { + let scaled = *value * factor; + let shifted = scaled + center; + shifted.round_ties_even().clamp(0.0, center * 2.0) as u8 + })); + } + (symbols, maxes) +} + +fn calculate_cdfs(symbols: &[u8], rows: usize, channels: usize) -> Result> { + let mut out = Vec::with_capacity(channels); + for channel in 0..channels { + let mut histogram = [0u32; CDF_LEN]; + for row in 0..rows { + let symbol = usize::from(symbols[row * channels + channel]); + let bucket = histogram + .get_mut(symbol + 1) + .ok_or_else(|| anyhow!("LMCache quantized symbol exceeds the CDF alphabet"))?; + *bucket += 1; + } + let mut running = 0u32; + for bucket in histogram.iter_mut().skip(1) { + let count = *bucket; + *bucket += running; + running += count; + } + if running != rows as u32 { + bail!("LMCache CDF histogram does not cover every token"); + } + let mut cdf = [0u16; CDF_LEN]; + let normalization = u32::from(u16::MAX) - MAX_BINS as u32; + for (index, count) in histogram.into_iter().enumerate() { + let normalized = normalization * count / running + index as u32; + cdf[index] = normalized as u16; + } + validate_cdf(&cdf)?; + out.push(cdf); + } + Ok(out) +} + +fn validate_cdf(cdf: &[u16; CDF_LEN]) -> Result<()> { + if cdf[0] != 0 || cdf[CDF_LEN - 1] != u16::MAX { + bail!("LMCache CDF endpoints are invalid"); + } + if cdf.windows(2).any(|pair| pair[0] >= pair[1]) { + bail!("LMCache CDF is not strictly increasing"); + } + Ok(()) +} + +fn arithmetic_encode_channel( + symbols: &[u8], + rows: usize, + channels: usize, + channel: usize, + cdf: &[u16; CDF_LEN], +) -> Vec { + let mut low = 0u32; + let mut high = u32::MAX; + let mut pending = 0u64; + let mut writer = BitWriter::default(); + for row in 0..rows { + let symbol = usize::from(symbols[row * channels + channel]); + let span = u64::from(high) - u64::from(low) + 1; + let c_low = u64::from(cdf[symbol]); + let c_high = if symbol == MAX_BINS - 1 { + u64::from(CDF_TOTAL) + } else { + u64::from(cdf[symbol + 1]) + }; + high = low + .wrapping_sub(1) + .wrapping_add(((span * c_high) >> 16) as u32); + low = low.wrapping_add(((span * c_low) >> 16) as u32); + loop { + if high < HALF { + writer.append_with_pending(0, &mut pending); + } else if low >= HALF { + writer.append_with_pending(1, &mut pending); + } else if low >= QUARTER && high < THREE_QUARTERS { + pending += 1; + low = (low << 1) & 0x7fff_ffff; + high = (high << 1) | 0x8000_0001; + continue; + } else { + break; + } + low <<= 1; + high = (high << 1) | 1; + } + } + pending += 1; + writer.append_with_pending(u32::from(low >= QUARTER), &mut pending); + writer.finish() +} + +fn arithmetic_decode_channel( + stream: &[u8], + rows: usize, + channels: usize, + channel: usize, + cdf: &[u16; CDF_LEN], + output: &mut [u8], +) -> Result<()> { + let mut reader = BitReader::new(stream); + let mut low = 0u32; + let mut high = u32::MAX; + let mut value = reader.read_u32(); + for row in 0..rows { + let span = u64::from(high) - u64::from(low) + 1; + let count = + (((u64::from(value) - u64::from(low) + 1) * u64::from(CDF_TOTAL) - 1) / span) as u16; + let symbol = cdf + .partition_point(|boundary| *boundary <= count) + .saturating_sub(1); + if symbol >= MAX_BINS { + bail!("LMCache arithmetic stream decoded an out-of-range symbol"); + } + output[row * channels + channel] = symbol as u8; + if row + 1 == rows { + break; + } + let c_low = u64::from(cdf[symbol]); + let c_high = if symbol == MAX_BINS - 1 { + u64::from(CDF_TOTAL) + } else { + u64::from(cdf[symbol + 1]) + }; + high = low + .wrapping_sub(1) + .wrapping_add(((span * c_high) >> 16) as u32); + low = low.wrapping_add(((span * c_low) >> 16) as u32); + loop { + if low >= HALF || high < HALF { + low <<= 1; + high = (high << 1) | 1; + value = (value << 1) | u32::from(reader.read_bit()); + } else if low >= QUARTER && high < THREE_QUARTERS { + low = (low << 1) & 0x7fff_ffff; + high = (high << 1) | 0x8000_0001; + value = value.wrapping_sub(QUARTER); + value = (value << 1) | u32::from(reader.read_bit()); + } else { + break; + } + } + } + Ok(()) +} + +#[derive(Default)] +struct BitWriter { + register: u32, + bits: u32, + bytes: Vec, +} + +impl BitWriter { + fn append_with_pending(&mut self, bit: u32, pending: &mut u64) { + self.add_repeated(bit, 1); + self.add_repeated(1 - bit, *pending); + *pending = 0; + } + + fn add_repeated(&mut self, bit: u32, mut count: u64) { + while count > 0 { + let take = count.min(u64::from(32 - self.bits)) as u32; + self.register <<= take; + if bit == 1 { + self.register |= if take == 32 { + u32::MAX + } else { + (1u32 << take) - 1 + }; + } + self.bits += take; + count -= u64::from(take); + if self.bits == 32 { + self.bytes.extend_from_slice(&self.register.to_be_bytes()); + self.register = 0; + self.bits = 0; + } + } + } + + fn finish(mut self) -> Vec { + if self.bits > 0 { + self.register <<= 32 - self.bits; + let bytes = self.bits.div_ceil(8) as usize; + self.bytes + .extend_from_slice(&self.register.to_be_bytes()[..bytes]); + } + self.bytes + } +} + +struct BitReader<'a> { + stream: &'a [u8], + bit: usize, +} + +impl<'a> BitReader<'a> { + fn new(stream: &'a [u8]) -> Self { + Self { stream, bit: 0 } + } + + fn read_bit(&mut self) -> u8 { + let byte = self.stream.get(self.bit / 8).copied().unwrap_or(0); + let value = (byte >> (7 - self.bit % 8)) & 1; + self.bit += 1; + value + } + + fn read_u32(&mut self) -> u32 { + let mut value = 0u32; + for _ in 0..32 { + value = (value << 1) | u32::from(self.read_bit()); + } + value + } +} + +fn parse(payload: &[u8]) -> Result> { + if payload.len() < HEADER_LEN || (payload[..4] != MAGIC_V1 && payload[..4] != MAGIC_V2) { + bail!("not an LMCache CacheGen portable segment"); + } + let bins = payload[4]; + let packed = payload[..4] == MAGIC_V2; + let bits_per_symbol = if !packed { + if payload[5] != 0 { + bail!("invalid LMCache CacheGen segment header"); + } + 0 + } else { + payload[5] + }; + let expected_bits = if bins == 16 { 4 } else { 5 }; + if !matches!(bins, 16 | 32) || (packed && bits_per_symbol != expected_bits) { + bail!("invalid LMCache CacheGen segment header"); + } + let rows = usize::from(u16::from_le_bytes([payload[6], payload[7]])); + let channels = u32::from_le_bytes(payload[8..12].try_into().expect("four bytes")) as usize; + let stream_len = u32::from_le_bytes(payload[12..16].try_into().expect("four bytes")) as usize; + if rows == 0 || rows > MAX_TOKENS_PER_CHUNK || channels == 0 { + bail!("invalid LMCache CacheGen segment geometry"); + } + let values = rows + .checked_mul(channels) + .ok_or_else(|| anyhow!("LMCache CacheGen segment shape overflow"))?; + if values > MAX_DECODED_VALUES { + bail!("LMCache CacheGen segment exceeds the decoded-value ceiling"); + } + let max_bytes = rows + .checked_mul(4) + .ok_or_else(|| anyhow!("max metadata overflow"))?; + let cdf_bytes = if bits_per_symbol == 0 { + channels + .checked_mul(CDF_LEN * 2) + .ok_or_else(|| anyhow!("CDF metadata overflow"))? + } else { + 0 + }; + let length_bytes = if bits_per_symbol == 0 { + channels + .checked_mul(2) + .ok_or_else(|| anyhow!("length metadata overflow"))? + } else { + 0 + }; + let metadata_end = HEADER_LEN + .checked_add(max_bytes) + .and_then(|value| value.checked_add(cdf_bytes)) + .and_then(|value| value.checked_add(length_bytes)) + .ok_or_else(|| anyhow!("LMCache CacheGen metadata range overflow"))?; + if payload.len() + != metadata_end + .checked_add(stream_len) + .ok_or_else(|| anyhow!("payload overflow"))? + { + bail!("LMCache CacheGen payload length disagrees with its header"); + } + let mut cursor = HEADER_LEN; + let mut maxes = Vec::with_capacity(rows); + for _ in 0..rows { + let bits = u32::from_le_bytes(payload[cursor..cursor + 4].try_into().expect("four bytes")); + let value = f32::from_bits(bits); + if !value.is_finite() || value < 0.0 { + bail!("LMCache CacheGen segment contains an invalid row maximum"); + } + maxes.push(value); + cursor += 4; + } + if bits_per_symbol != 0 { + let expected_stream_len = values + .checked_mul(usize::from(bits_per_symbol)) + .ok_or_else(|| anyhow!("packed CacheGen bit length overflow"))? + .div_ceil(8); + if stream_len != expected_stream_len { + bail!("packed CacheGen byte length is inconsistent"); + } + return Ok(Parsed { + bins, + channels, + rows, + bits_per_symbol, + maxes, + cdfs: Vec::new(), + lengths: Vec::new(), + streams: &payload[cursor..], + }); + } + let mut cdfs = Vec::with_capacity(channels); + for _ in 0..channels { + let mut cdf = [0u16; CDF_LEN]; + for value in &mut cdf { + *value = u16::from_le_bytes(payload[cursor..cursor + 2].try_into().expect("two bytes")); + cursor += 2; + } + validate_cdf(&cdf)?; + cdfs.push(cdf); + } + let mut lengths = Vec::with_capacity(channels); + let mut total = 0usize; + for _ in 0..channels { + let length = usize::from(u16::from_le_bytes( + payload[cursor..cursor + 2].try_into().expect("two bytes"), + )); + if length == 0 || length > MAX_TOKENS_PER_CHUNK { + bail!("LMCache CacheGen channel stream length is invalid"); + } + total = total + .checked_add(length) + .ok_or_else(|| anyhow!("stream length overflow"))?; + lengths.push(length); + cursor += 2; + } + if total != stream_len { + bail!("LMCache CacheGen channel lengths do not sum to the stream size"); + } + Ok(Parsed { + bins, + channels, + rows, + bits_per_symbol, + maxes, + cdfs, + lengths, + streams: &payload[cursor..], + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + const FIXTURE_HEADER_LEN: usize = 16; + + fn input(rows: usize, channels: usize) -> Vec { + let mut raw = Vec::with_capacity(rows * channels * 2); + for row in 0..rows { + for channel in 0..channels { + let phase = (row * channels + channel) as f32 * 0.03125; + let value = phase.sin() * 0.75 + (channel as f32 - 3.0) * 0.015625; + raw.extend_from_slice(&f32_to_f16_bits(value).to_le_bytes()); + } + } + raw + } + + #[test] + fn generic_bin_schedule_matches_lmcache() { + assert_eq!(bins_for_layer(0, 8, true), 32); + assert_eq!(bins_for_layer(7, 8, false), 32); + assert_eq!(bins_for_layer(9, 28, true), 32); + assert_eq!(bins_for_layer(10, 28, true), 16); + assert_eq!(bins_for_layer(1, 28, false), 32); + assert_eq!(bins_for_layer(2, 28, false), 16); + } + + #[test] + fn segment_round_trips_through_lmcache_quantization() { + let raw = input(64, 8); + for bins in [16, 32] { + let encoded = encode_f16_segment(&raw, 8, bins).expect("encode"); + let decoded = decode_f16_segment(&encoded).expect("decode"); + assert_eq!(decoded.len(), raw.len()); + for (original, restored) in raw + .as_chunks::<2>() + .0 + .iter() + .zip(decoded.as_chunks::<2>().0) + { + let original = f16_bits_to_f32(u16::from_le_bytes(*original)); + let restored = f16_bits_to_f32(u16::from_le_bytes(*restored)); + assert!( + (original - restored).abs() < 0.12, + "{original} rebuilt as {restored}" + ); + } + } + } + + #[test] + fn packed_segment_preserves_lmcache_quantized_values() { + let raw = input(64, 8); + for bins in [16, 32] { + let oracle = encode_f16_segment(&raw, 8, bins).expect("oracle encode"); + let packed = encode_f16_segment_packed(&raw, 8, bins).expect("packed encode"); + assert_eq!(&packed[..4], b"LCG2"); + assert_eq!(packed[5], if bins == 16 { 4 } else { 5 }); + assert_eq!( + decode_f16_segment(&packed).expect("packed decode"), + decode_f16_segment(&oracle).expect("oracle decode") + ); + } + } + + #[test] + fn flat_rows_restore_exactly() { + let mut raw = Vec::new(); + for value in [0.0f32, 0.5, -0.75] { + for _ in 0..8 { + raw.extend_from_slice(&f32_to_f16_bits(value).to_le_bytes()); + } + } + let encoded = encode_f16_segment(&raw, 8, 32).expect("encode"); + assert_eq!(decode_f16_segment(&encoded).expect("decode"), raw); + } + + #[test] + fn malformed_envelopes_are_refused() { + let raw = input(8, 4); + let encoded = encode_f16_segment(&raw, 4, 32).expect("encode"); + assert!(decode_f16_segment(&encoded[..encoded.len() - 1]).is_err()); + let mut bad_cdf = encoded.clone(); + let cdf_start = HEADER_LEN + 8 * 4; + bad_cdf[cdf_start + 2..cdf_start + 4].copy_from_slice(&0u16.to_le_bytes()); + assert!(decode_f16_segment(&bad_cdf).is_err()); + + let packed = encode_f16_segment_packed(&raw, 4, 32).expect("packed encode"); + assert!(decode_f16_segment(&packed[..packed.len() - 1]).is_err()); + let mut bad_width = packed; + bad_width[5] = 4; + assert!(decode_f16_segment(&bad_width).is_err()); + } + + fn assert_lmcache_fixture(fixture: &[u8], bins: u8) { + assert_eq!(&fixture[..4], b"LCFX"); + let raw_len = + u32::from_le_bytes(fixture[4..8].try_into().expect("fixture raw length")) as usize; + let encoded_len = + u32::from_le_bytes(fixture[8..12].try_into().expect("fixture encoded length")) as usize; + let decoded_len = + u32::from_le_bytes(fixture[12..16].try_into().expect("fixture decoded length")) + as usize; + assert_eq!( + fixture.len(), + FIXTURE_HEADER_LEN + raw_len + encoded_len + decoded_len + ); + let raw = &fixture[FIXTURE_HEADER_LEN..FIXTURE_HEADER_LEN + raw_len]; + let encoded_start = FIXTURE_HEADER_LEN + raw_len; + let expected = &fixture[encoded_start..encoded_start + encoded_len]; + let expected_decoded = &fixture[encoded_start + encoded_len..]; + let encoded = encode_f16_segment(raw, 8, bins).expect("encode fixture"); + assert_eq!( + encoded, expected, + "Rust encoder diverged from the pinned LMCache scalar reference" + ); + assert_eq!( + decode_f16_segment(expected).expect("decode fixture"), + expected_decoded, + "Rust decoder diverged from the pinned LMCache scalar reference" + ); + } + + #[test] + fn matches_lmcache_bins16_fixture() { + assert_lmcache_fixture(include_bytes!("fixtures/lmcache_b5d109e_bins16.bin"), 16); + } + + #[test] + fn matches_lmcache_bins32_fixture() { + assert_lmcache_fixture(include_bytes!("fixtures/lmcache_b5d109e_bins32.bin"), 32); + } +} diff --git a/crates/skippy-cache/src/cachegen/mod.rs b/crates/skippy-cache/src/cachegen/mod.rs new file mode 100644 index 0000000000..3814335074 --- /dev/null +++ b/crates/skippy-cache/src/cachegen/mod.rs @@ -0,0 +1,12 @@ +//! CacheGen experiments and the pinned LMCache-compatible CPU reference +//! used by the acceptance gate in #1652. +//! +//! [`lmcache`] owns the active reference algorithm. [`container`], [`reference`], +//! and [`rans`] retain the earlier simplified prototype and its historical +//! fixtures; production and acceptance-gate code must not select that path. + +pub mod archive; +pub mod container; +pub mod lmcache; +pub mod rans; +pub mod reference; diff --git a/crates/skippy-cache/src/cachegen/rans.rs b/crates/skippy-cache/src/cachegen/rans.rs new file mode 100644 index 0000000000..9948d4db50 --- /dev/null +++ b/crates/skippy-cache/src/cachegen/rans.rs @@ -0,0 +1,311 @@ +//! Byte-aligned rANS entropy coder, the static-CDF variant CacheGen uses. +//! +//! Ported from Fabian Giesen's public-domain `ryg_rans` (`rans_byte.h`), +//! which is also the coder LMCache ships. This is the deterministic CPU +//! reference: the GPU slices in the later CubeCL work must produce +//! byte-identical streams. The straightforward divide/mod operations are +//! used on purpose — this crate is the correctness oracle, not the +//! performance path. +//! +//! Encoding processes symbols in reverse order and emits bytes backwards; +//! [`RansEncoder::finish`] returns the stream the decoder consumes +//! forwards. + +/// Lower bound of the normalization interval (`RANS_BYTE_L`). +pub const RANS_L: u32 = 1 << 23; +/// CDF precision. Frequencies sum to `1 << SCALE_BITS` (CacheGen uses 12). +pub const SCALE_BITS: u32 = 12; +/// Total frequency, `2^SCALE_BITS`. +pub const SCALE: u32 = 1 << SCALE_BITS; + +/// Forward symbol table: cumulative start and frequency per token. +#[derive(Debug, Clone)] +pub struct SymbolTable { + /// Cumulative frequency at each symbol's range start. + pub starts: Vec, + /// Frequency of each symbol. `starts[0] == 0` and + /// `starts.last() + freq.last() == SCALE`. + pub freqs: Vec, +} + +impl SymbolTable { + /// Builds the table from per-symbol frequencies. An all-zero histogram + /// is refused rather than silently producing an undecodable table. + pub fn from_freqs(freqs: &[u32]) -> Option { + if freqs.is_empty() || freqs.iter().any(|&freq| freq > SCALE) { + return None; + } + let total: u64 = freqs.iter().map(|&freq| u64::from(freq)).sum(); + if total != u64::from(SCALE) { + return None; + } + let mut starts = Vec::with_capacity(freqs.len()); + let mut running = 0u32; + for &freq in freqs { + starts.push(running); + running = running.checked_add(freq)?; + } + Some(Self { + starts, + freqs: freqs.to_vec(), + }) + } + + /// The symbol whose range contains `value` (`RansDecGet` output). + pub fn symbol_for(&self, value: u32) -> usize { + let position = self + .starts + .partition_point(|&start| start <= value) + .saturating_sub(1); + position.min(self.freqs.len() - 1) + } +} + +/// Static-frequency byte-rANS encoder. +#[derive(Debug)] +pub struct RansEncoder { + state: u32, + /// Emitted bytes in time order; reversed into the final stream. + emitted: Vec, +} + +impl RansEncoder { + pub fn new() -> Self { + Self { + state: RANS_L, + emitted: Vec::new(), + } + } + + /// Encodes one symbol. Symbols must be pushed in reverse stream order + /// (last symbol first). + pub fn put(&mut self, table: &SymbolTable, symbol: usize) { + let start = table.starts[symbol]; + let freq = table.freqs[symbol]; + let x_max = ((RANS_L >> SCALE_BITS) << 8) * freq; + while self.state >= x_max { + self.emitted.push((self.state & 0xff) as u8); + self.state >>= 8; + } + self.state = ((self.state / freq) << SCALE_BITS) + (self.state % freq) + start; + } + + /// Flushes the state and returns the wire stream. + pub fn finish(mut self) -> Vec { + let mut stream = self.state.to_le_bytes().to_vec(); + stream.extend(self.emitted.drain(..).rev()); + stream + } +} + +impl Default for RansEncoder { + fn default() -> Self { + Self::new() + } +} + +/// Static-frequency byte-rANS decoder over a complete stream. +pub struct RansDecoder<'a> { + state: u32, + bytes: &'a [u8], + cursor: usize, +} + +impl<'a> RansDecoder<'a> { + /// Starts decoding; consumes the initial 4-byte little-endian state. + pub fn new(bytes: &'a [u8]) -> Option { + if bytes.len() < 4 { + return None; + } + let state = u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]); + Some(Self { + state, + bytes, + cursor: 4, + }) + } + + /// Decodes the next symbol (stream order). + pub fn get(&mut self, table: &SymbolTable) -> Option { + let value = self.state & (SCALE - 1); + let symbol = table.symbol_for(value); + let start = table.starts[symbol]; + let freq = table.freqs[symbol]; + let mut x = freq * (self.state >> SCALE_BITS) + value - start; + while x < RANS_L { + let byte = *self.bytes.get(self.cursor)?; + self.cursor += 1; + x = (x << 8) | u32::from(byte); + } + self.state = x; + Some(symbol) + } +} + +#[cfg(test)] +mod rans_golden_fixture_tests { + use super::*; + + /// Path of the committed golden fixture, relative to the crate root. + pub(crate) const GOLDEN_FIXTURE_PATH: &str = "src/cachegen/fixtures/ryg_rans_golden.bin"; + + /// Generator for the committed golden fixture. Run explicitly: + /// + /// ```text + /// cargo test -p skippy-cache --lib -- --ignored ryg_rans_generate + /// ``` + /// + /// The generator is deterministic: rerunning it reproduces the + /// committed bytes exactly, so provenance is checkable. The committed + /// file is the frozen artifact the normal suite verifies against — + /// corruption or hand-editing breaks + /// [`ryg_rans_golden_fixture_decodes`], because the decoder's output + /// stops matching the declared symbol pattern. + #[test] + #[ignore = "generator for the committed golden fixture; run explicitly"] + fn ryg_rans_generate_golden_fixture() { + // Exercise a skewed distribution with long runs in both symbols — + // the shape where normalization and multi-byte emit paths both + // fire. Same pattern family as the round-trip test above. + let table = SymbolTable::from_freqs(&[SCALE / 2, SCALE - SCALE / 2]).expect("table"); + let symbols: Vec = (0..1000) + .map(|index| if (index / 7) % 5 == 0 { 1 } else { 0 }) + .collect(); + let mut encoder = RansEncoder::new(); + for symbol in symbols.iter().rev() { + encoder.put(&table, *symbol); + } + let stream = encoder.finish(); + + let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(GOLDEN_FIXTURE_PATH); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent).expect("create fixtures dir"); + } + let length = u32::try_from(stream.len()).expect("fixture length fits u32"); + let mut meta = Vec::with_capacity(8 + stream.len()); + meta.extend_from_slice(&1_000u32.to_le_bytes()); + meta.extend_from_slice(&length.to_le_bytes()); + meta.extend_from_slice(&stream); + std::fs::write(path, &meta).expect("write fixture meta+stream"); + } + + /// The independent-correctness gate (scama blocker 4): the committed + /// 129-byte stream was produced by canonical upstream `ryg_rans` + /// (`rans_byte.h`, upstream commit `c9d162d996fd600315af9ae8eb89d832 + /// 576cb32d`, independently reproduced in C by scama; stream-only + /// SHA-256 `7887612d9c251e19cfa54558718d3388dd7a5b5dfd7f8d0f419681ea + /// ddc62a32`), and the decoder must reproduce the declared symbol + /// sequence from it. + #[test] + fn ryg_rans_golden_fixture_decodes() { + let raw = std::fs::read( + std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(GOLDEN_FIXTURE_PATH), + ) + .expect("golden fixture must be committed; run the ryg_rans_generate test once"); + assert!(raw.len() >= 8, "fixture carries its meta header"); + let count = u32::from_le_bytes(raw[0..4].try_into().expect("4 bytes")) as usize; + let stream_len = u32::from_le_bytes(raw[4..8].try_into().expect("4 bytes")) as usize; + assert_eq!( + raw.len(), + 8 + stream_len, + "fixture length must match its declared stream" + ); + let table = SymbolTable::from_freqs(&[SCALE / 2, SCALE - SCALE / 2]).expect("table"); + let mut decoder = RansDecoder::new(&raw[8..]).expect("decoder"); + let expected: Vec = (0..count) + .map(|index| if (index / 7) % 5 == 0 { 1 } else { 0 }) + .collect(); + let mut decoded = Vec::with_capacity(count); + for _ in 0..count { + decoded.push(decoder.get(&table).expect("symbol within fixture")); + } + assert_eq!(decoded, expected, "decoder diverged from the golden stream"); + } + + /// The encoder half of the golden gate (scama re-review blocker 2): + /// the frozen 129-byte fixture came from canonical upstream + /// `ryg_rans` (`c9d162d9`, independently reproduced in C by scama; + /// stream-only SHA-256 `7887612d…ddc62a32`), so this test pins the + /// Rust encoder's byte output against those bytes. Unlike the + /// `--ignored` generator — which runs the same encoder under test + /// and cannot detect drift — this is a normal test and fails the + /// moment `RansEncoder` stops producing the canonical stream. + #[test] + fn ryg_rans_golden_fixture_pins_the_encoder() { + let raw = std::fs::read( + std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(GOLDEN_FIXTURE_PATH), + ) + .expect("golden fixture must be committed; run the ryg_rans_generate test once"); + let count = u32::from_le_bytes(raw[0..4].try_into().expect("4 bytes")) as usize; + let stream_len = u32::from_le_bytes(raw[4..8].try_into().expect("4 bytes")) as usize; + let frozen_stream = &raw[8..8 + stream_len]; + + // The pinned inputs: same symbols, same CDF as the fixture's + // provenance (50/50 split CDF, run-heavy pattern). + let table = SymbolTable::from_freqs(&[SCALE / 2, SCALE - SCALE / 2]).expect("table"); + let symbols: Vec = (0..count) + .map(|index| if (index / 7) % 5 == 0 { 1 } else { 0 }) + .collect(); + let mut encoder = RansEncoder::new(); + for symbol in symbols.iter().rev() { + encoder.put(&table, *symbol); + } + assert_eq!( + encoder.finish(), + frozen_stream, + "encoder drift: Rust stream no longer matches canonical ryg_rans bytes" + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn two_symbol_table() -> SymbolTable { + SymbolTable::from_freqs(&[SCALE / 2, SCALE - SCALE / 2]).expect("table") + } + + #[test] + fn round_trips_a_skewed_two_symbol_stream() { + let table = two_symbol_table(); + // Skewed pattern with long runs in both directions. + let symbols: Vec = (0..1000) + .map(|index| if (index / 7) % 5 == 0 { 1 } else { 0 }) + .collect(); + let mut encoder = RansEncoder::new(); + for symbol in symbols.iter().rev() { + encoder.put(&table, *symbol); + } + let stream = encoder.finish(); + + let mut decoder = RansDecoder::new(&stream).expect("decoder"); + let mut decoded = Vec::with_capacity(symbols.len()); + for _ in 0..symbols.len() { + decoded.push(decoder.get(&table).expect("symbol")); + } + assert_eq!(decoded, symbols); + } + + #[test] + fn rejects_frequencies_that_do_not_sum_to_scale() { + assert!(SymbolTable::from_freqs(&[1, 2, 3]).is_none()); + assert!(SymbolTable::from_freqs(&[]).is_none()); + assert!(SymbolTable::from_freqs(&[SCALE + 1]).is_none()); + assert!(SymbolTable::from_freqs(&[SCALE]).is_some()); + } + + #[test] + fn single_symbol_stream_still_round_trips() { + let table = SymbolTable::from_freqs(&[SCALE]).expect("table"); + let mut encoder = RansEncoder::new(); + for _ in 0..64 { + encoder.put(&table, 0); + } + let stream = encoder.finish(); + let mut decoder = RansDecoder::new(&stream).expect("decoder"); + for _ in 0..64 { + assert_eq!(decoder.get(&table).expect("symbol"), 0); + } + } +} diff --git a/crates/skippy-cache/src/cachegen/reference.rs b/crates/skippy-cache/src/cachegen/reference.rs new file mode 100644 index 0000000000..ec2c0dce4c --- /dev/null +++ b/crates/skippy-cache/src/cachegen/reference.rs @@ -0,0 +1,305 @@ +//! Deterministic quantization and token-axis delta reference for CacheGen. +//! +//! This is the CPU oracle the later GPU work must reproduce bit-for-bit. +//! Every operation is an IEEE-754 single-precision arithmetic op or an +//! integer op, so results are identical on any conformant platform and any +//! backend that performs the same operations in the same order — that +//! property is what makes this module a golden reference rather than an +//! implementation detail. +//! +//! Simplification versus the CacheGen paper, deliberate and documented: +//! calibration is per-segment min/max affine at 4 bits. The paper calibrates +//! once per model over a sample of requests and applies K/M-mixed +//! quantization across tensor dimensions; both are follow-up experiments +//! behind the same container, and neither changes the wire contract. + +use anyhow::{Result, bail}; + +/// Bits per quantized symbol (CacheGen's default 4-bit KV quantization). +pub const QUANT_BITS: u32 = 4; +/// Quantized symbol alphabet size. +pub const TOKEN_COUNT: usize = 1 << QUANT_BITS; +/// Mask extracting a symbol from a byte; the delta ring is modulo this. +pub(crate) const ALPHABET_MASK: u8 = (TOKEN_COUNT - 1) as u8; + +/// Per-tile affine calibration, carried bit-exactly in the container. +/// +/// `min` and `scale` are the f32 bit patterns of the calibration values: +/// storing bits rather than decimal text is what makes the container +/// deterministic. `scale = (max - min) / (TOKEN_COUNT - 1)` over the tile's +/// f16-decoded values. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Calibration { + /// f32 bits of the calibration minimum. + pub min_bits: u32, + /// f32 bits of the quantization step. + pub scale_bits: u32, +} + +/// Derives the affine calibration for one tile of f16-decoded values. +/// Refuses empty or non-finite input: an NaN/inf element would poison the +/// calibration, and the crate's activation policy refuses non-finite +/// payloads rather than encoding them. +pub fn calibrate(values: &[f32]) -> Result { + let Some(&first) = values.first() else { + bail!("cannot calibrate an empty tile"); + }; + let mut min = first; + let mut max = first; + for &value in values { + if !value.is_finite() { + bail!("tile contains a non-finite value; refusing to calibrate"); + } + min = min.min(value); + max = max.max(value); + } + let scale = (max - min) / f32::from(TOKEN_COUNT as u16 - 1); + Ok(Calibration { + min_bits: min.to_bits(), + scale_bits: scale.to_bits(), + }) +} + +/// Quantizes one tile to symbols in `[0, TOKEN_COUNT)`. Midpoint rounding, +/// clamped at both ends. Must be preceded by [`calibrate`] over the same +/// values. +pub fn quantize(calibration: &Calibration, values: &[f32]) -> Result> { + let min = f32::from_bits(calibration.min_bits); + let scale = f32::from_bits(calibration.scale_bits); + let max_symbol = (TOKEN_COUNT - 1) as f32; + values + .iter() + .map(|&value| { + if !value.is_finite() { + bail!("tile contains a non-finite value; refusing to quantize"); + } + if scale == 0.0 { + // Flat tile: every value equals the calibration minimum. + return Ok(0u8); + } + let scaled = ((value - min) / scale).round(); + let clamped = scaled.clamp(0.0, max_symbol); + Ok(clamped as u8) + }) + .collect() +} + +/// Reconstructs f32 values from symbols. The inverse of quantization up to +/// the quantization error itself (at most half a step before the final +/// f16 round-trip). +pub fn dequantize(calibration: &Calibration, symbols: &[u8]) -> Vec { + let min = f32::from_bits(calibration.min_bits); + let scale = f32::from_bits(calibration.scale_bits); + symbols + .iter() + .map(|&symbol| f32::from(symbol) * scale + min) + .collect() +} + +/// Token-axis delta transform, CacheGen's cheap decorrelation pass. +/// +/// Symbols along the token axis (contiguous rows of `dims` values) become +/// the difference to the previous row's symbol, reduced modulo the 4-bit +/// alphabet: the mask keeps every symbol inside `[0, TOKEN_COUNT)` so the +/// entropy stage still sees a 16-symbol alphabet, and mod-`TOKEN_COUNT` +/// arithmetic is an exact ring, so the inverse is lossless. +/// +/// `symbols.len()` must be a non-zero multiple of `dims`. +pub fn delta_encode(symbols: &mut [u8], dims: usize) -> Result<()> { + validate_shape(symbols.len(), dims)?; + for row in (dims..symbols.len()).rev() { + symbols[row] = symbols[row].wrapping_sub(symbols[row - dims]) & ALPHABET_MASK; + } + Ok(()) +} + +/// Inverse of [`delta_encode`]. +pub fn delta_decode(symbols: &mut [u8], dims: usize) -> Result<()> { + validate_shape(symbols.len(), dims)?; + for row in dims..symbols.len() { + symbols[row] = symbols[row].wrapping_add(symbols[row - dims]) & ALPHABET_MASK; + } + Ok(()) +} + +/// Builds the static CDF from a symbol histogram: floor allocation with a +/// minimum frequency of one token per symbol, then deterministic +/// redistribution (subtract from the largest, add to the histogram-max with +/// lowest index) until the frequencies sum to the rANS scale. The +/// redistribution order is fixed, so the same histogram always yields the +/// same table on every backend. +/// +/// `total` must equal `sum(histogram)` — the caller's contract, enforced +/// here. With it, both normalization loops are provably bounded: the +/// decrement loop removes at most one unit per min-clamped symbol (at most +/// `TOKEN_COUNT` iterations) and the increment loop adds at most +/// `SCALE - 1` missing units. Without it, a forged histogram/total pair +/// could drive unbounded repair work. +pub fn histogram_to_freqs(histogram: &[u32], total: usize) -> Result> { + if histogram.len() != TOKEN_COUNT { + bail!("histogram must cover every 4-bit symbol"); + } + if total == 0 { + bail!("cannot build a symbol table for an empty tile"); + } + let histogram_total: u64 = histogram.iter().map(|&count| u64::from(count)).sum(); + if histogram_total != total as u64 { + bail!("histogram totals {histogram_total} but the caller declares {total} symbols"); + } + let mut freqs: Vec = histogram + .iter() + .map(|&count| (u64::from(count) * u64::from(super::rans::SCALE) / total as u64) as u32) + .collect(); + // Every symbol must be decodable, even unused ones: a corrupt or + // truncated stream must fail cleanly, not index an empty range. + for freq in &mut freqs { + *freq = (*freq).max(1); + } + let mut sum: u64 = freqs.iter().map(|&freq| u64::from(freq)).sum(); + while sum > u64::from(super::rans::SCALE) { + let largest = freqs + .iter() + .enumerate() + .max_by_key(|(index, freq)| (*freq, std::cmp::Reverse(*index))) + .expect("non-empty") + .0; + freqs[largest] -= 1; + sum -= 1; + } + let mut index = 0usize; + while sum < u64::from(super::rans::SCALE) { + // Unreachable with a validated histogram (every other symbol holds + // at least one unit, so no target can sit at the full scale), but + // the bound makes the loop's totality explicit rather than argued. + if index > u64::from(super::rans::SCALE) as usize { + bail!("histogram normalization failed to converge"); + } + let candidate = histogram + .iter() + .enumerate() + .max_by_key(|(position, count)| (*count, std::cmp::Reverse(*position))) + .expect("non-empty") + .0; + let target = if index == 0 { + candidate + } else { + (candidate + index) % TOKEN_COUNT + }; + if freqs[target] < super::rans::SCALE { + freqs[target] += 1; + sum += 1; + } + index = index.wrapping_add(1); + } + Ok(freqs) +} + +fn validate_shape(len: usize, dims: usize) -> Result<()> { + if dims == 0 || !len.is_multiple_of(dims) { + bail!("tile shape mismatch: {len} values are not rows of {dims}"); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn quantization_round_trips_within_half_a_step() { + let values: Vec = (0..64).map(|index| index as f32 / 7.0).collect(); + let calibration = calibrate(&values).expect("calibrate"); + let symbols = quantize(&calibration, &values).expect("quantize"); + let rebuilt = dequantize(&calibration, &symbols); + let scale = f32::from_bits(calibration.scale_bits); + for (original, restored) in values.iter().zip(&rebuilt) { + assert!( + (original - restored).abs() <= scale, + "{original} rebuilt as {restored} with step {scale}" + ); + } + } + + #[test] + fn flat_tile_quantizes_to_zero_and_dequantizes_exactly() { + let values = vec![0.25f32; 32]; + let calibration = calibrate(&values).expect("calibrate"); + assert_eq!(calibration.scale_bits, 0.0f32.to_bits()); + let symbols = quantize(&calibration, &values).expect("quantize"); + assert!(symbols.iter().all(|&symbol| symbol == 0)); + let rebuilt = dequantize(&calibration, &symbols); + assert!(rebuilt.iter().all(|&value| value == 0.25)); + } + + #[test] + fn delta_round_trips_across_the_wrap() { + let dims = 3; + let mut symbols = vec![15u8, 0, 7, 1, 14, 8, 2, 13, 9]; + let original = symbols.clone(); + delta_encode(&mut symbols, dims).expect("encode"); + delta_decode(&mut symbols, dims).expect("decode"); + assert_eq!(symbols, original); + } + + #[test] + fn delta_rejects_misshaped_tiles() { + let mut symbols = vec![0u8; 7]; + assert!(delta_encode(&mut symbols, 3).is_err()); + assert!(delta_decode(&mut symbols, 0).is_err()); + assert!(delta_encode(&mut symbols, 7).is_ok()); + } + + #[test] + fn non_finite_values_are_refused() { + let mut values = vec![1.0f32, 2.0]; + values.push(f32::NAN); + assert!(calibrate(&values).is_err()); + let calibration = calibrate(&[1.0, 2.0]).expect("calibrate"); + assert!(quantize(&calibration, &[f32::INFINITY]).is_err()); + } + + #[test] + fn freqs_sum_to_scale_and_are_deterministic() { + let mut histogram = vec![0u32; TOKEN_COUNT]; + histogram[0] = 700; + histogram[3] = 200; + histogram[9] = 100; + let total: usize = histogram.iter().sum::() as usize; + let first = histogram_to_freqs(&histogram, total).expect("freqs"); + let second = histogram_to_freqs(&histogram, total).expect("freqs"); + assert_eq!(first, second); + let sum: u32 = first.iter().sum(); + assert_eq!(sum, super::super::rans::SCALE); + assert!(first.iter().all(|&freq| freq >= 1)); + } + + #[test] + fn a_total_that_disagrees_with_the_histogram_is_refused() { + let mut histogram = vec![0u32; TOKEN_COUNT]; + histogram[0] = 700; + histogram[3] = 200; + histogram[9] = 100; + // The forged-total cases: wildly low, wildly high, and off-by-one. + // Any of them used to seed the CDF from a tile that never existed. + for lie in [1usize, usize::MAX, 1001] { + let error = histogram_to_freqs(&histogram, lie) + .expect_err("a disagreeing total must be refused"); + assert!( + error.to_string().contains("declares"), + "error should name the disagreement: {error}" + ); + } + } + + #[test] + fn unused_symbols_still_get_a_decodable_frequency() { + let histogram = { + let mut histogram = vec![0u32; TOKEN_COUNT]; + histogram[2] = 4096; + histogram + }; + let freqs = histogram_to_freqs(&histogram, 4096).expect("freqs"); + assert!(freqs.iter().all(|&freq| freq >= 1)); + assert_eq!(freqs.iter().sum::(), super::super::rans::SCALE); + } +} diff --git a/crates/skippy-cache/src/fsinfo.rs b/crates/skippy-cache/src/fsinfo.rs new file mode 100644 index 0000000000..fd25d8b07e --- /dev/null +++ b/crates/skippy-cache/src/fsinfo.rs @@ -0,0 +1,482 @@ +//! Filesystem facts the L3 disk tier has to know before it writes. +//! +//! The tier promises a minimum-free-space reserve, owner-only permissions, a +//! refusal to run on network filesystems, and a recency signal cheap enough to +//! update on every cache hit. None of that is in `std`, so the platform calls +//! live here and the store stays free of `unsafe`. + +use std::{ + fs, + path::{Component, Path, PathBuf}, + time::SystemTime, +}; + +use anyhow::{Context, Result, bail}; + +/// Bytes an unprivileged writer can still add to the filesystem holding +/// `path`. This is `f_bavail`, not `f_bfree`: the reserve check must not spend +/// blocks only root can allocate. +pub fn available_bytes(path: &Path) -> Result { + fs2::available_space(path) + .with_context(|| format!("failed to stat available space for {}", path.display())) +} + +/// Whether `path` sits on a filesystem the tier refuses to manage. Network +/// filesystems break the atomic-rename and locking assumptions the store is +/// built on, so §10.10 rejects them where they are reliably detectable. +pub fn is_network_filesystem(path: &Path) -> Result { + #[cfg(windows)] + { + return windows::is_network_filesystem(path); + } + #[cfg(not(windows))] + let name = filesystem_type_name(path)?; + #[cfg(not(windows))] + { + Ok(matches!( + name.as_str(), + "nfs" | "smbfs" | "afpfs" | "webdav" | "ftp" | "cifs" | "fuse" | "fuse.sshfs" + )) + } +} + +/// The filesystem type as the kernel names it, for status reporting. +pub fn filesystem_type_name(path: &Path) -> Result { + #[cfg(windows)] + { + return windows::filesystem_type_name(path); + } + #[cfg(target_os = "macos")] + { + let stat = statfs(path)?; + let raw = stat.f_fstypename; + let bytes: Vec = raw + .iter() + .take_while(|byte| **byte != 0) + .map(|byte| *byte as u8) + .collect(); + Ok(String::from_utf8_lossy(&bytes).into_owned()) + } + #[cfg(all(not(target_os = "macos"), not(windows)))] + { + // Linux reports a magic number rather than a name. Only the values the + // tier actually refuses are worth naming; anything else is local + // enough to manage. + let stat = statfs(path)?; + Ok(match stat.f_type { + 0x6969 => "nfs".to_string(), + 0xFF53_4D42 => "cifs".to_string(), + // FUSE_SUPER_MAGIC. statfs cannot name the subtype, so every FUSE + // mount reports as plain "fuse": sshfs and a local fuse filesystem + // are indistinguishable here. The tier refuses the whole class + // rather than guess, which is the conservative reading of §10.10. + 0x6573_5546 => "fuse".to_string(), + other => format!("0x{other:x}"), + }) + } +} + +/// Mark an entry as used now, so eviction can order by last use rather than +/// last write. One `utimensat` per cache hit is the bounded metadata update +/// §13.4 allows; it writes no payload bytes and allocates no blocks. +pub fn touch(path: &Path) -> Result<()> { + fs::OpenOptions::new() + .write(true) + .open(path) + .and_then(|file| file.set_times(fs::FileTimes::new().set_modified(SystemTime::now()))) + .with_context(|| format!("failed to touch {}", path.display())) +} + +/// Publish a fully written temporary file at `destination`, replacing an +/// existing entry when necessary. Windows `rename` does not replace an +/// existing file, so use the platform primitive with explicit replacement. +pub fn replace_file(temp: &Path, destination: &Path) -> Result<()> { + #[cfg(windows)] + { + return windows::replace_file(temp, destination); + } + #[cfg(not(windows))] + { + fs::rename(temp, destination) + .with_context(|| format!("failed to publish {}", destination.display())) + } +} + +/// Restrict a cache directory to its owner. The first release has no at-rest +/// encryption, so local account permissions are the only confidentiality the +/// tier offers and it must actually apply them. +pub fn restrict_to_owner(path: &Path, mode: u32) -> Result<()> { + #[cfg(windows)] + { + let _ = mode; + return windows::restrict_to_owner(path); + } + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + let permissions = fs::Permissions::from_mode(mode); + fs::set_permissions(path, permissions) + .with_context(|| format!("failed to restrict permissions on {}", path.display())) + } +} + +/// Refuse a path that reaches the store through a symlink. Following one would +/// let anything with write access to the parent redirect committed cache +/// bytes outside the managed root, past every budget and reserve check. +pub fn refuse_symlink(path: &Path) -> Result<()> { + match fs::symlink_metadata(path) { + Ok(metadata) if is_link_or_reparse_point(&metadata) => { + bail!( + "{} is a symlink; the cache refuses to traverse it", + path.display() + ) + } + Ok(_) => Ok(()), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(error) => Err(error).with_context(|| format!("failed to stat {}", path.display())), + } +} + +/// Create an absolute directory tree without traversing an untrusted link. +/// Existing ancestors are inspected before any missing component is created, +/// closing the gap where `create_dir_all` could follow a redirected parent. +pub fn create_dir_all_without_links(path: &Path) -> Result<()> { + let mut current = PathBuf::new(); + for component in path.components() { + if matches!(component, Component::Prefix(_)) { + current.push(component.as_os_str()); + continue; + } + current.push(component.as_os_str()); + match fs::symlink_metadata(¤t) { + Ok(metadata) if is_link_or_reparse_point(&metadata) => { + if current != path && is_trusted_platform_directory_link(¤t) { + continue; + } + bail!( + "{} is a symlink; the cache refuses to traverse it", + current.display() + ); + } + Ok(metadata) if !metadata.is_dir() => { + bail!("{} is not a directory", current.display()); + } + Ok(_) => {} + Err(error) if error.kind() == std::io::ErrorKind::NotFound => { + match fs::create_dir(¤t) { + Ok(()) => {} + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => { + let metadata = fs::symlink_metadata(¤t).with_context(|| { + format!("failed to verify cache directory {}", current.display()) + })?; + if is_link_or_reparse_point(&metadata) || !metadata.is_dir() { + bail!( + "{} appeared during creation but is not a safe directory", + current.display() + ); + } + } + Err(error) => { + return Err(error).with_context(|| { + format!("failed to create cache directory {}", current.display()) + }); + } + } + } + Err(error) => { + return Err(error).with_context(|| format!("failed to stat {}", current.display())); + } + } + } + refuse_symlink(path) +} + +fn is_link_or_reparse_point(metadata: &fs::Metadata) -> bool { + if metadata.file_type().is_symlink() { + return true; + } + #[cfg(windows)] + { + use std::os::windows::fs::MetadataExt; + const FILE_ATTRIBUTE_REPARSE_POINT: u32 = 0x400; + return metadata.file_attributes() & FILE_ATTRIBUTE_REPARSE_POINT != 0; + } + #[cfg(not(windows))] + false +} + +fn is_trusted_platform_directory_link(_path: &Path) -> bool { + #[cfg(target_os = "macos")] + { + matches!(_path.to_str(), Some("/var" | "/tmp" | "/etc")) + } + #[cfg(not(target_os = "macos"))] + false +} + +/// Refuse a symlink anywhere in the portion of `path` below `root`. +/// +/// The system prefix above the cache root is not ours to police: on macOS +/// `/var` is itself a symlink to `/private/var`, so refusing every symlinked +/// ancestor would reject the default temp and cache locations. What must hold +/// is that nothing the store creates under its own resolved root redirects +/// bytes outside it. +pub fn refuse_symlinked_descendant(root: &Path, path: &Path) -> Result<()> { + let Ok(relative) = path.strip_prefix(root) else { + bail!( + "{} is not inside the cache root {}", + path.display(), + root.display() + ); + }; + let mut walked = root.to_path_buf(); + for component in relative.components() { + walked.push(component); + refuse_symlink(&walked)?; + } + Ok(()) +} + +#[cfg(not(windows))] +fn c_path(path: &Path) -> Result { + use std::os::unix::ffi::OsStrExt; + std::ffi::CString::new(path.as_os_str().as_bytes()) + .with_context(|| format!("path {} contains an interior NUL", path.display())) +} + +#[cfg(not(windows))] +fn statfs(path: &Path) -> Result { + let c_path = c_path(path)?; + let mut stat = std::mem::MaybeUninit::::uninit(); + // SAFETY: as `statvfs` above. + let status = unsafe { libc::statfs(c_path.as_ptr(), stat.as_mut_ptr()) }; + if status != 0 { + return Err(std::io::Error::last_os_error()) + .with_context(|| format!("failed to stat filesystem for {}", path.display())); + } + // SAFETY: statfs returned 0, so `stat` is initialized. + Ok(unsafe { stat.assume_init() }) +} + +#[cfg(windows)] +mod windows { + use super::*; + use std::ffi::c_void; + use std::mem::{align_of, size_of}; + use std::os::windows::ffi::OsStrExt; + use std::ptr::{null, null_mut}; + use windows_sys::Win32::Foundation::CloseHandle; + use windows_sys::Win32::Security::Authorization::{SE_FILE_OBJECT, SetNamedSecurityInfoW}; + use windows_sys::Win32::Security::{ + ACCESS_ALLOWED_ACE, ACL, ACL_REVISION, AddAccessAllowedAceEx, CONTAINER_INHERIT_ACE, + DACL_SECURITY_INFORMATION, GetLengthSid, GetTokenInformation, InitializeAcl, + OBJECT_INHERIT_ACE, OWNER_SECURITY_INFORMATION, PROTECTED_DACL_SECURITY_INFORMATION, PSID, + TOKEN_QUERY, TOKEN_USER, TokenUser, + }; + use windows_sys::Win32::Storage::FileSystem::{ + FILE_ALL_ACCESS, GetDriveTypeW, GetVolumeInformationW, GetVolumePathNameW, + MOVEFILE_REPLACE_EXISTING, MOVEFILE_WRITE_THROUGH, MoveFileExW, + }; + use windows_sys::Win32::System::Threading::{GetCurrentProcess, OpenProcessToken}; + use windows_sys::Win32::System::WindowsProgramming::DRIVE_REMOTE; + + pub(super) fn is_network_filesystem(path: &Path) -> Result { + let volume = volume_root(path)?; + Ok(unsafe { GetDriveTypeW(volume.as_ptr()) } == DRIVE_REMOTE) + } + + pub(super) fn filesystem_type_name(path: &Path) -> Result { + let volume = volume_root(path)?; + let mut filesystem = [0_u16; 64]; + let ok = unsafe { + GetVolumeInformationW( + volume.as_ptr(), + null_mut(), + 0, + null_mut(), + null_mut(), + null_mut(), + filesystem.as_mut_ptr(), + filesystem.len() as u32, + ) + }; + if ok == 0 { + return Err(std::io::Error::last_os_error()) + .with_context(|| format!("failed to inspect filesystem for {}", path.display())); + } + let length = filesystem + .iter() + .position(|unit| *unit == 0) + .unwrap_or(filesystem.len()); + Ok(String::from_utf16_lossy(&filesystem[..length]).to_ascii_lowercase()) + } + + pub(super) fn replace_file(temp: &Path, destination: &Path) -> Result<()> { + let temp_wide = to_wide(temp); + let destination_wide = to_wide(destination); + let result = unsafe { + MoveFileExW( + temp_wide.as_ptr(), + destination_wide.as_ptr(), + MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH, + ) + }; + if result == 0 { + return Err(std::io::Error::last_os_error()) + .with_context(|| format!("failed to publish {}", destination.display())); + } + Ok(()) + } + + fn volume_root(path: &Path) -> Result> { + let input = to_wide(path); + let mut volume = vec![0_u16; 260]; + let ok = + unsafe { GetVolumePathNameW(input.as_ptr(), volume.as_mut_ptr(), volume.len() as u32) }; + if ok == 0 { + return Err(std::io::Error::last_os_error()) + .with_context(|| format!("failed to resolve volume for {}", path.display())); + } + Ok(volume) + } + + pub(super) fn restrict_to_owner(path: &Path) -> Result<()> { + with_current_user_sid(|sid| { + let acl_bytes = size_of::() + size_of::() - size_of::() + + unsafe { GetLengthSid(sid) as usize }; + let words = acl_bytes.div_ceil(size_of::()); + let mut acl_storage = vec![0_u64; words]; + let acl = acl_storage.as_mut_ptr().cast::(); + let metadata = fs::metadata(path)?; + let ace_flags = if metadata.is_dir() { + OBJECT_INHERIT_ACE | CONTAINER_INHERIT_ACE + } else { + 0 + }; + unsafe { + if InitializeAcl(acl, acl_bytes as u32, ACL_REVISION) == 0 + || AddAccessAllowedAceEx(acl, ACL_REVISION, ace_flags, FILE_ALL_ACCESS, sid) + == 0 + { + return Err(std::io::Error::last_os_error().into()); + } + } + let mut wide = to_wide(path); + let result = unsafe { + SetNamedSecurityInfoW( + wide.as_mut_ptr(), + SE_FILE_OBJECT, + OWNER_SECURITY_INFORMATION + | DACL_SECURITY_INFORMATION + | PROTECTED_DACL_SECURITY_INFORMATION, + sid, + null_mut(), + acl, + null(), + ) + }; + if result != 0 { + bail!("Windows ACL update failed with error {result}"); + } + Ok(()) + }) + .with_context(|| format!("failed to restrict permissions on {}", path.display())) + } + + fn with_current_user_sid(f: impl FnOnce(PSID) -> Result) -> Result { + let mut token = null_mut(); + if unsafe { OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &mut token) } == 0 { + return Err(std::io::Error::last_os_error().into()); + } + let _token = Handle(token); + let mut bytes = 0_u32; + unsafe { + let _ = GetTokenInformation(token, TokenUser, null_mut(), 0, &mut bytes); + } + if bytes == 0 { + return Err(std::io::Error::last_os_error().into()); + } + let words = (bytes as usize).div_ceil(align_of::()); + let mut buffer = vec![0_usize; words]; + if unsafe { + GetTokenInformation( + token, + TokenUser, + buffer.as_mut_ptr().cast::(), + bytes, + &mut bytes, + ) + } == 0 + { + return Err(std::io::Error::last_os_error().into()); + } + let token_user = unsafe { &*buffer.as_ptr().cast::() }; + f(token_user.User.Sid) + } + + fn to_wide(path: &Path) -> Vec { + path.as_os_str().encode_wide().chain(Some(0)).collect() + } + + struct Handle(windows_sys::Win32::Foundation::HANDLE); + + impl Drop for Handle { + fn drop(&mut self) { + unsafe { + CloseHandle(self.0); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn available_bytes_reports_a_plausible_figure() { + let available = available_bytes(Path::new(".")).expect("stat working directory"); + assert!(available > 0, "working directory reports no free space"); + } + + #[test] + fn touch_moves_modification_time_forward() { + let directory = + std::env::temp_dir().join(format!("skippy-fsinfo-touch-{}", std::process::id())); + fs::create_dir_all(&directory).expect("create temp dir"); + let path = directory.join("entry"); + fs::write(&path, b"entry").expect("write entry"); + let before = fs::metadata(&path) + .expect("stat") + .modified() + .expect("mtime"); + std::thread::sleep(std::time::Duration::from_millis(20)); + touch(&path).expect("touch entry"); + let after = fs::metadata(&path) + .expect("stat") + .modified() + .expect("mtime"); + assert!(after > before, "touch did not move the modification time"); + fs::remove_dir_all(&directory).ok(); + } + + #[cfg(unix)] + #[test] + fn symlinks_are_refused() { + let directory = + std::env::temp_dir().join(format!("skippy-fsinfo-symlink-{}", std::process::id())); + fs::create_dir_all(&directory).expect("create temp dir"); + let target = directory.join("target"); + let link = directory.join("link"); + fs::write(&target, b"target").expect("write target"); + let _ = fs::remove_file(&link); + std::os::unix::fs::symlink(&target, &link).expect("create symlink"); + assert!(refuse_symlink(&link).is_err(), "symlink was accepted"); + assert!(refuse_symlink(&target).is_ok(), "regular file was refused"); + assert!( + refuse_symlink(&directory.join("absent")).is_ok(), + "absent path was refused" + ); + fs::remove_dir_all(&directory).ok(); + } +} diff --git a/crates/skippy-cache/src/identity.rs b/crates/skippy-cache/src/identity.rs index 77e361e96a..ad2506e248 100644 --- a/crates/skippy-cache/src/identity.rs +++ b/crates/skippy-cache/src/identity.rs @@ -151,18 +151,7 @@ fn update_weight_identity(hasher: &mut blake3::Hasher, config: &StageConfig) { } } hasher.update(b"checkpoint-quantization:"); - let quantization = config - .checkpoint_quantization - .as_deref() - .unwrap_or("preserve") - .chars() - .filter(|character| character.is_ascii_alphanumeric()) - .flat_map(char::to_uppercase) - .collect::(); - hasher.update(match quantization.as_str() { - "DIRECT" | "NONE" | "PRESERVE" => b"PRESERVE" as &[u8], - _ => quantization.as_bytes(), - }); + hasher.update(normalized_checkpoint_quantization(config).as_bytes()); hasher.update(b"checkpoint-imatrix:"); match config.checkpoint_imatrix_sha256.as_deref() { Some(digest) => hasher.update(digest.as_bytes()), @@ -180,6 +169,21 @@ fn update_weight_identity(hasher: &mut blake3::Hasher, config: &StageConfig) { }); } +fn normalized_checkpoint_quantization(config: &StageConfig) -> String { + let quantization = config + .checkpoint_quantization + .as_deref() + .unwrap_or("preserve") + .chars() + .filter(|character| character.is_ascii_alphanumeric()) + .flat_map(char::to_uppercase) + .collect::(); + match quantization.as_str() { + "DIRECT" | "NONE" | "PRESERVE" => "PRESERVE".to_string(), + _ => quantization, + } +} + pub fn prefix_hash_with_namespace( config: &StageConfig, token_start: u64, @@ -218,11 +222,10 @@ fn prefix_namespace_hasher( // process and would prevent otherwise compatible cache identities from // agreeing across runs. // - // Reuse depends on the *shape* of the stage: the model, owned layers, and - // pipeline position. Those are hashed below with the runtime layout fields - // in `update_layout_identity`. - hasher.update(config.stage_id.as_bytes()); - hasher.update(&config.stage_index.to_le_bytes()); + // Reuse depends on the numerical shape of the stage, not its placement. + // Replicas may have different stage ids/indexes while owning the same layer + // range; including those labels prevents both disk reuse and later remote + // handoff between otherwise compatible runtimes. hasher.update(&config.layer_start.to_le_bytes()); hasher.update(&config.layer_end.to_le_bytes()); hasher.update(NATIVE_KV_RUNTIME_ABI_VERSION.as_bytes()); @@ -239,6 +242,172 @@ fn prefix_namespace_hasher( hasher } +/// Numerical identity for an exact-state payload produced by a stage. +/// +/// This is the production counterpart to [`ExactStateIdentityParams`]. It +/// deliberately reuses the same exhaustive weight/layout hashing as radix +/// pages while excluding run, topology, stage id/index, addresses, and other +/// placement labels. +pub fn exact_state_identity_for_stage(config: &StageConfig, payload_kind: &str) -> String { + let mut hasher = blake3::Hasher::new(); + hasher.update(b"exact-state-stage-identity-v1"); + update_weight_identity(&mut hasher, config); + hasher.update(&config.layer_start.to_le_bytes()); + hasher.update(&config.layer_end.to_le_bytes()); + hasher.update(NATIVE_KV_RUNTIME_ABI_VERSION.as_bytes()); + hasher.update(&NATIVE_KV_LAYER_CONTIGUOUS_LAYOUT.to_le_bytes()); + hasher.update(NATIVE_KV_DTYPE.as_bytes()); + update_layout_identity(&mut hasher, config); + update_platform_identity(&mut hasher); + hasher.update(&config.ctx_size.to_le_bytes()); + hasher.update(&config.lane_count.to_le_bytes()); + hasher.update(b"payload:"); + hasher.update(payload_kind.as_bytes()); + format!("blake3:{}", hasher.finalize().to_hex()) +} + +/// Stable identity for model-scoped lifecycle operations. +/// +/// Unlike an exact-state identity this deliberately excludes stage layout, +/// payload format, placement, and device details: every split stage for the +/// same numerical model must be selected by one model-scoped prune or clear. +/// A source-model digest is authoritative when present. Older/direct inputs +/// without one fall back to their materialized manifest/package identity. +pub fn numerical_model_identity_for_stage(config: &StageConfig) -> String { + let mut hasher = blake3::Hasher::new(); + hasher.update(b"numerical-model-identity-v1"); + match config.source_model_sha256.as_deref() { + Some(digest) => { + hasher.update(b"source:"); + hasher.update(digest.as_bytes()); + } + None => { + hasher.update(b"source:"); + hasher.update(b"manifest:"); + hasher.update( + config + .manifest_sha256 + .as_deref() + .unwrap_or("") + .as_bytes(), + ); + hasher.update(b"package:"); + hasher.update( + config + .package_ref + .as_deref() + .unwrap_or("") + .as_bytes(), + ); + // The fallback must not collapse unrelated legacy configurations + // that have no content digest at all. + hasher.update(b"model:"); + hasher.update(config.model_id.as_bytes()); + } + } + hasher.update(b"checkpoint-quantization:"); + hasher.update(normalized_checkpoint_quantization(config).as_bytes()); + hasher.update(b"checkpoint-imatrix:"); + hasher.update( + config + .checkpoint_imatrix_sha256 + .as_deref() + .unwrap_or("") + .as_bytes(), + ); + format!("blake3:{}", hasher.finalize().to_hex()) +} + +/// The *numerical* identity of an exact-state (full-state or KV+recurrent) +/// payload for prefill/decode handoff. +/// +/// This is the numerical half of the numerical-vs-placement identity split: +/// it covers every input that changes the bytes or the interpretation of an +/// exported state blob — weights, cache dtypes, flash-attention layout, the +/// GPU layer split, backend, platform, layer range, and the context shape +/// (`ctx_size`, `lane_count`, which decide the KV buffer geometry a +/// full-state blob is laid out against). +/// +/// It deliberately excludes placement: `stage_id`, `stage_index`, +/// `topology_id`, `run_id`, and bind addresses. A prefill replica and a +/// decode replica differ in exactly those fields, and state must flow +/// between them whenever the numerical identity matches. +pub struct ExactStateIdentityParams<'a> { + pub model_id: &'a str, + pub model_revision: Option<&'a str>, + pub model_file: Option<&'a str>, + /// Content digests of the served weights, when known. `model_id` is a + /// display name — two runs can present the same id while serving + /// different tensors (requantized artifact, republished package, + /// swapped GGUF), and state crossing that boundary is silent numerical + /// corruption. Absent digests are tagged distinctly so `None` cannot + /// alias a real value. + pub manifest_sha256: Option<&'a str>, + pub source_model_sha256: Option<&'a str>, + pub package_ref: Option<&'a str>, + /// Stable name of the tensor assembly path. Runtime slices, artifact + /// slices, and layer packages can produce different tensor layouts from + /// the same model coordinate and must not exchange resident state. + pub load_mode: &'a str, + pub cache_type_k: &'a str, + pub cache_type_v: &'a str, + pub flash_attn_type: FlashAttentionType, + pub n_gpu_layers: i32, + pub backend_device: Option<&'a str>, + pub layer_start: u32, + pub layer_end: u32, + pub ctx_size: u32, + pub lane_count: u32, + pub payload_kind: &'a str, +} + +pub fn exact_state_identity(params: &ExactStateIdentityParams<'_>) -> String { + let mut hasher = blake3::Hasher::new(); + hasher.update(b"exact-state-identity-v1"); + hasher.update(NATIVE_KV_RUNTIME_ABI_VERSION.as_bytes()); + hasher.update(b"model:"); + hasher.update(params.model_id.as_bytes()); + for (tag, value) in [ + (&b"revision:"[..], params.model_revision), + (&b"file:"[..], params.model_file), + (&b"manifest:"[..], params.manifest_sha256), + (&b"source:"[..], params.source_model_sha256), + (&b"package:"[..], params.package_ref), + (&b"device:"[..], params.backend_device), + ] { + hasher.update(tag); + match value { + Some(value) => { + hasher.update(b"="); + hasher.update(value.as_bytes()); + } + None => { + hasher.update(b""); + } + } + } + hasher.update(b"kv:"); + hasher.update(b"load:"); + hasher.update(params.load_mode.as_bytes()); + hasher.update(params.cache_type_k.as_bytes()); + hasher.update(b"/"); + hasher.update(params.cache_type_v.as_bytes()); + hasher.update(match params.flash_attn_type { + FlashAttentionType::Auto => b"fa:auto", + FlashAttentionType::Disabled => b"fa:offf", + FlashAttentionType::Enabled => b"fa:onnn", + }); + hasher.update(¶ms.n_gpu_layers.to_le_bytes()); + hasher.update(¶ms.layer_start.to_le_bytes()); + hasher.update(¶ms.layer_end.to_le_bytes()); + hasher.update(¶ms.ctx_size.to_le_bytes()); + hasher.update(¶ms.lane_count.to_le_bytes()); + hasher.update(b"kind:"); + hasher.update(params.payload_kind.as_bytes()); + update_platform_identity(&mut hasher); + format!("blake3:{}", hasher.finalize().to_hex()) +} + pub fn page_id( config: &StageConfig, token_start: u64, @@ -260,6 +429,157 @@ pub fn activation_page_id(page_id: &str, activation_width: i32) -> String { format!("act:{}:w{}", page_id, activation_width.max(0)) } +#[cfg(test)] +mod exact_state_identity_tests { + use super::*; + + fn params() -> ExactStateIdentityParams<'static> { + ExactStateIdentityParams { + model_id: "org/model:Q4_K_M", + model_revision: Some("abc123"), + model_file: Some("model.gguf"), + manifest_sha256: Some("m".repeat(64).leak()), + source_model_sha256: Some("s".repeat(64).leak()), + package_ref: None, + load_mode: "runtime-slice", + cache_type_k: "f16", + cache_type_v: "f16", + flash_attn_type: FlashAttentionType::Auto, + n_gpu_layers: 99, + backend_device: Some("Metal"), + layer_start: 0, + layer_end: 28, + ctx_size: 8192, + lane_count: 2, + payload_kind: "full-state", + } + } + + /// Placement is excluded *by construction*: the params carry no stage, + /// topology, or run fields, so two replicas that differ only in placement + /// produce the same identity. + #[test] + fn identical_numerics_produce_identical_identity() { + assert_eq!( + exact_state_identity(¶ms()), + exact_state_identity(¶ms()) + ); + } + + #[test] + fn cache_dtype_changes_identity() { + let saver = ExactStateIdentityParams { + cache_type_k: "q4_0", + cache_type_v: "q4_0", + ..params() + }; + assert_ne!( + exact_state_identity(¶ms()), + exact_state_identity(&saver) + ); + } + + #[test] + fn context_shape_changes_identity() { + let wider_ctx = ExactStateIdentityParams { + ctx_size: 16384, + ..params() + }; + let more_lanes = ExactStateIdentityParams { + lane_count: 4, + ..params() + }; + assert_ne!( + exact_state_identity(¶ms()), + exact_state_identity(&wider_ctx) + ); + assert_ne!( + exact_state_identity(¶ms()), + exact_state_identity(&more_lanes) + ); + } + + #[test] + fn backend_and_weights_change_identity() { + let cuda = ExactStateIdentityParams { + backend_device: Some("CUDA0"), + ..params() + }; + let other_revision = ExactStateIdentityParams { + model_revision: Some("def456"), + ..params() + }; + let absent_revision = ExactStateIdentityParams { + model_revision: None, + ..params() + }; + assert_ne!(exact_state_identity(¶ms()), exact_state_identity(&cuda)); + assert_ne!( + exact_state_identity(¶ms()), + exact_state_identity(&other_revision) + ); + assert_ne!( + exact_state_identity(¶ms()), + exact_state_identity(&absent_revision) + ); + } + + #[test] + fn load_mode_changes_identity() { + let artifact_slice = ExactStateIdentityParams { + load_mode: "artifact-slice", + ..params() + }; + assert_ne!( + exact_state_identity(¶ms()), + exact_state_identity(&artifact_slice) + ); + } + + /// Weight content digests must separate state even when the display + /// model id matches — the same argument `update_weight_identity` makes + /// for KV pages. + #[test] + fn weight_digests_change_identity() { + let requantized = ExactStateIdentityParams { + source_model_sha256: Some("t".repeat(64).leak()), + ..params() + }; + let repacked = ExactStateIdentityParams { + manifest_sha256: Some("n".repeat(64).leak()), + ..params() + }; + let absent = ExactStateIdentityParams { + manifest_sha256: None, + ..params() + }; + assert_ne!( + exact_state_identity(¶ms()), + exact_state_identity(&requantized) + ); + assert_ne!( + exact_state_identity(¶ms()), + exact_state_identity(&repacked) + ); + assert_ne!( + exact_state_identity(¶ms()), + exact_state_identity(&absent) + ); + } + + #[test] + fn payload_kind_changes_identity() { + let pages = ExactStateIdentityParams { + payload_kind: "kv-recurrent", + ..params() + }; + assert_ne!( + exact_state_identity(¶ms()), + exact_state_identity(&pages) + ); + } +} + #[cfg(test)] mod identity_completeness_tests { use skippy_protocol::{LoadMode, StageDevice}; @@ -781,6 +1101,71 @@ mod identity_stability_tests { ); } + #[test] + fn placement_labels_do_not_change_numerical_identity() { + let first = config_with_topology("topology-a"); + let replica = StageConfig { + run_id: "other-run".to_string(), + topology_id: "topology-b".to_string(), + stage_id: "prefill-replica-7".to_string(), + stage_index: 9, + bind_addr: "127.0.0.1:9999".to_string(), + ..first.clone() + }; + let tokens = (0..256).collect::>(); + + assert_eq!( + prefix_namespace_hash(&first, 0, None), + prefix_namespace_hash(&replica, 0, None) + ); + assert_eq!( + exact_state_identity_for_stage(&first, "full-state"), + exact_state_identity_for_stage(&replica, "full-state") + ); + assert_eq!( + numerical_model_identity_for_stage(&first), + numerical_model_identity_for_stage(&replica) + ); + assert_eq!( + prefix_identity(&first, 0, &tokens).prefix_hash, + prefix_identity(&replica, 0, &tokens).prefix_hash + ); + } + + #[test] + fn model_identity_groups_split_packages_from_one_source_model() { + let mut first = config_with_topology("topology-a"); + first.source_model_sha256 = Some("source-digest".to_string()); + first.manifest_sha256 = Some("stage-a-manifest".to_string()); + first.package_ref = Some("stage-a-package".to_string()); + let second = StageConfig { + stage_id: "stage-b".to_string(), + stage_index: 1, + layer_start: 24, + layer_end: 48, + manifest_sha256: Some("stage-b-manifest".to_string()), + package_ref: Some("stage-b-package".to_string()), + ..first.clone() + }; + + assert_eq!( + numerical_model_identity_for_stage(&first), + numerical_model_identity_for_stage(&second) + ); + assert_ne!( + exact_state_identity_for_stage(&first, "full-state"), + exact_state_identity_for_stage(&second, "full-state") + ); + + let mut equivalent_spelling = first.clone(); + equivalent_spelling.checkpoint_quantization = Some("none".to_string()); + assert_eq!( + numerical_model_identity_for_stage(&first), + numerical_model_identity_for_stage(&equivalent_spelling), + "equivalent preserve-mode spellings must share lifecycle identity" + ); + } + /// Stage shape still has to match: a different layer range is a different /// page and must never collide. #[test] diff --git a/crates/skippy-cache/src/l2/mod.rs b/crates/skippy-cache/src/l2/mod.rs new file mode 100644 index 0000000000..3591b589bf --- /dev/null +++ b/crates/skippy-cache/src/l2/mod.rs @@ -0,0 +1,2445 @@ +//! Host-RAM L2 tier over the packed L3 segment format (#1651). +//! +//! The radix cache (L1) holds resident payloads; the L3 tier holds the same +//! state durably on disk as content-addressed packed segments. This module +//! adds the missing middle tier: a bounded host-RAM cache of *immutable +//! packed segments*, keyed and identified exactly like the L3 entries they +//! mirror. +//! +//! Contract (mirrors `crate::tier::L3Tier`): +//! +//! - **Identity**: entries are stamped with the tier's model and exact-state +//! identities. The cache key includes the identities, so a re-identity is +//! a wholesale miss, never a silent hit. +//! - **Coordinates**: entries are keyed by the same +//! `(namespace, token path)` coordinates L3 uses — +//! [`crate::tier::l3_prefix_key`] / [`crate::tier::l3_namespace_key`] — so +//! an L2 hit is interchangeable with the L3 entry it cached. +//! - **Segment sharing**: an entry stores immutable `Arc>` segment +//! handles keyed by their content digests plus a layout that maps the +//! entry's L3 manifest segment list onto those handles. A longer prefix +//! that extends a shorter one references the same segment handles, so +//! turn growth shares prefix bytes instead of duplicating them — L2 RAM +//! tracks *distinct segment* bytes, not per-prefix assembled bytes. +//! - **Integrity**: the whole concatenated L3 wire digest (the manifest key) +//! is verified exactly once, at admission, against the payload being +//! admitted. After admission the segment bytes are immutable, so every +//! later read is a digest lookup plus handle assembly — no re-hash. An +//! admission-time mismatch refuses the insert; L2 never holds bytes it +//! did not verify. +//! - **Bounded**: the byte budget is charged with each entry's *distinct* +//! segment bytes (bytes not already held by an in-flight insert) and +//! enforced by evicting in deterministic LRU order. A payload whose +//! distinct bytes exceed the whole budget is refused. A segment shared +//! with an already-admitted entry is shared for accounting too: only the +//! first admission pays for it. +//! - **Zero-copy reads**: `get` clones handles, not bytes; the returned +//! `CacheBytes` is a block-backed view over the shared segment storages, +//! contiguous in the single-segment case. +//! +//! This first slice is a standalone store with no wiring into the request +//! path; the benchmark harness drives it directly. L2 promotion/demotion +//! policy and server integration land in a later slice. +use std::{ + collections::HashMap, + ops::Range, + sync::{ + Arc, Mutex, + atomic::{AtomicU64, Ordering}, + }, +}; + +use crate::payload::{CacheBytes, ExactStatePayloadKind}; +use crate::{HandoffManifest, segment_digest}; +#[cfg(test)] +use crate::{HandoffSegmentRef, MANIFEST_VERSION, PayloadCodec, SegmentCodecIdentity}; + +/// Where an entry came from, for telemetry and promotion policy later. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum L2Origin { + /// Admitted from an assembled L3 fill (verified wire). + FromL3, + /// Admitted from another verified wire source (tests, prefetch). + Direct, +} + +/// LRU eviction accounting for one removed entry. +/// +/// `freed_bytes` is what removal actually released: segments whose last +/// referencing entry left the tier. `retained_bytes` is shared-segment +/// bytes that stay because another entry still references them. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct L2Eviction { + pub cache_key: String, + pub freed_bytes: u64, + pub retained_bytes: u64, +} + +impl L2Eviction { + /// Bytes charged to the budget for this entry (what its removal freed). + pub fn payload_bytes(&self) -> u64 { + self.freed_bytes + } +} + +/// Read path counters. One snapshot per `stats()` call. +#[derive(Debug, Default, Clone, PartialEq, Eq)] +pub struct L2Stats { + pub entries: u64, + /// Sum of entries' distinct segment charges — the live budget usage. + pub bytes: u64, + /// Sum of entry payload lengths including cross-entry sharing; larger + /// than `bytes` exactly when entries share prefix segments. + pub logical_bytes: u64, + /// Distinct immutable segment handles currently held. + pub segments: u64, + /// Bytes held in the segment pool (== `bytes` when the pool is live). + pub segment_bytes: u64, + /// Distinct segment bytes a single admission did not have to copy + /// because an earlier admission already held them. + pub shared_bytes_admitted: u64, + pub budget_bytes: u64, + pub hits: u64, + pub misses: u64, + pub inserts: u64, + pub evictions: u64, + /// Admissions refused because the payload digest did not match the + /// bytes (hash mismatch or malformed digest string). + pub admission_rejects: u64, + pub refused_bytes: u64, +} + +/// One assembled L2 entry: which stored segment handles make up the wire, +/// in manifest order, plus the payload split so a fill can be rebuilt. +#[derive(Debug, Clone)] +pub struct L2Layout { + pub payload_kind: ExactStatePayloadKind, + pub total_bytes: u64, + pub kv_bytes: u64, + pub recurrent_bytes: u64, + /// `(segment digest, byte range within the assembled wire)` per + /// manifest segment, in manifest order. Ranges concatenate to + /// `0..total_bytes` exactly as the L3 manifest tiles them. + pub segments: Vec<(String, Range)>, +} + +/// The L2 mirror of an assembled L3 entry: verified segment handles plus +/// the layout needed to rebuild a serving payload without disk I/O. +#[derive(Debug, Clone)] +pub enum ExactStatePayloadMirror { + FullState { layout: L2Layout }, + RecurrentOnly { layout: L2Layout }, + KvRecurrent { layout: L2Layout }, +} + +impl ExactStatePayloadMirror { + pub fn kind(&self) -> ExactStatePayloadKind { + match self { + Self::FullState { .. } => ExactStatePayloadKind::FullState, + Self::RecurrentOnly { .. } => ExactStatePayloadKind::RecurrentOnly, + Self::KvRecurrent { .. } => ExactStatePayloadKind::KvRecurrent, + } + } + + /// Total wire length of the entry (the assembled payload length). + pub fn byte_len(&self) -> u64 { + match self { + Self::FullState { layout } + | Self::RecurrentOnly { layout } + | Self::KvRecurrent { layout } => layout.total_bytes, + } + } + + /// Build a mirror from a captured L3 manifest. Callers must verify the + /// payload wire against `manifest.payload_digest` — `admit` does this — + /// before the mirror is stored. + pub fn from_manifest(manifest: &HandoffManifest) -> Result { + let kind = match manifest.payload_kind.as_str() { + "full-state" => ExactStatePayloadKind::FullState, + "recurrent-only" => ExactStatePayloadKind::RecurrentOnly, + "kv-recurrent" => ExactStatePayloadKind::KvRecurrent, + other => { + return Err(L2InsertRefusal::UnknownPayloadKind(other.to_string())); + } + }; + let mut offset = 0u64; + let segments = + manifest + .segments + .iter() + .enumerate() + .map(|(index, segment)| { + if segment.index != index as u32 { + return Err(L2InsertRefusal::MalformedManifest(format!( + "segment {} records index {} but sits at position {index}", + segment.digest, segment.index + ))); + } + if segment.offset != offset { + return Err(L2InsertRefusal::MalformedManifest(format!( + "segment {} records offset {} but tiles at {offset}", + segment.digest, segment.offset + ))); + } + let start = offset; + offset = offset.checked_add(segment.bytes).ok_or( + L2InsertRefusal::MalformedManifest("segment tiling overflows".to_string()), + )?; + Ok((segment.digest.clone(), start..offset)) + }) + .collect::, _>>()?; + if offset != manifest.total_bytes { + return Err(L2InsertRefusal::MalformedManifest(format!( + "segments tile {offset} bytes but the manifest records {}", + manifest.total_bytes + ))); + } + let layout = L2Layout { + payload_kind: kind, + total_bytes: manifest.total_bytes, + kv_bytes: manifest.kv_bytes, + recurrent_bytes: manifest.recurrent_bytes, + segments, + }; + Ok(match kind { + ExactStatePayloadKind::FullState => Self::FullState { layout }, + ExactStatePayloadKind::RecurrentOnly => Self::RecurrentOnly { layout }, + ExactStatePayloadKind::KvRecurrent => Self::KvRecurrent { layout }, + }) + } + + fn layout(&self) -> &L2Layout { + match self { + Self::FullState { layout } + | Self::RecurrentOnly { layout } + | Self::KvRecurrent { layout } => layout, + } + } + + /// Segment digests in wire order, deduplicated. + fn segment_digests(&self) -> Vec<&str> { + let mut seen = Vec::new(); + for (digest, _) in &self.layout().segments { + if !seen.contains(&digest.as_str()) { + seen.push(digest.as_str()); + } + } + seen + } +} + +/// A hit handed to the caller: the entry's layout plus `Arc` clones of the +/// segment handles the layout references, keyed by digest. The tier stores +/// this directly on `L2Hit` so payload assembly needs no tier lock. +#[derive(Debug, Clone)] +pub struct L2Hit { + pub payload: ExactStatePayloadMirror, + pub token_count: u64, + pub payload_digest: String, + /// Distinct segment handles referenced by the layout, keyed by digest. + pub(crate) segments: HashMap, +} + +impl L2Hit { + /// Rebuild a serving payload. Cheap in the common cases: the returned + /// `CacheBytes` is a block-backed view sharing the stored segment + /// storages (`Arc` clones, not byte copies); a single whole-storage + /// segment borrows it contiguously. Only a multi-segment read of + /// distinct storages materializes bytes, and only into the caller's + /// `Cow` on `as_cow`. + pub fn to_payload(&self) -> crate::payload::ExactStatePayload { + let layout = self.payload.layout(); + let wire = self.wire_view(0..layout.total_bytes); + match self.payload.kind() { + crate::payload::ExactStatePayloadKind::FullState => { + crate::payload::ExactStatePayload::FullState { bytes: wire } + } + crate::payload::ExactStatePayloadKind::RecurrentOnly => { + crate::payload::ExactStatePayload::RecurrentOnly { recurrent: wire } + } + crate::payload::ExactStatePayloadKind::KvRecurrent => { + // Split the wire at kv_bytes exactly like L3 load does: kv + // is the leading block-backed view, recurrent the tail. Both + // share the same storages; no bytes are copied here. + let kv_len = layout.kv_bytes.min(layout.total_bytes); + let kv = self.wire_view(0..kv_len); + let recurrent = self.wire_view(kv_len..layout.total_bytes); + crate::payload::ExactStatePayload::KvRecurrent { kv, recurrent } + } + } + } + + /// Block-backed `CacheBytes` over `range` of the assembled wire, in + /// wire order. Blocks outside `range` are skipped; edge blocks are + /// narrowed to the overlap. Byte-identical views share the same + /// segment storages; nothing is copied. + pub(crate) fn wire_view(&self, range: Range) -> CacheBytes { + let layout = self.payload.layout(); + let start = range.start.min(layout.total_bytes); + let end = range.end.min(layout.total_bytes); + let blocks = layout + .segments + .iter() + .filter_map(|(digest, segment_range)| { + let block_start = segment_range.start.max(start); + let block_end = segment_range.end.min(end); + if block_start >= block_end { + return None; + } + let storage = self + .segments + .get(digest) + .map(|handle| Arc::clone(&handle.bytes)) + .unwrap_or_else(|| Arc::new(Vec::new())); + let len = storage.len() as u64; + let from = (block_start - segment_range.start).min(len); + let to = (block_end - segment_range.start).min(len); + Some((digest.clone(), storage, (from as usize)..(to as usize))) + }) + .collect::>(); + CacheBytes::from_shared_blocks(end.saturating_sub(start), blocks) + } +} + +/// Presence probe result. Probing never changes LRU recency. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct L2Peek { + pub token_count: u64, + pub payload_digest: String, + /// Total wire length including segments shared with other entries. + pub payload_bytes: u64, + /// Distinct segment bytes charged to the budget for this entry. + pub distinct_bytes: u64, + pub origin: L2Origin, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum L2InsertRefusal { + EmptyPayload, + OverBudget { + payload_bytes: u64, + }, + /// Admission would push the pool past the budget with the incoming + /// entry's handles pinned: even evicting every other entry could not + /// free the excess, so the transaction was rolled back untouched. + ProtectedOvercommit { + pool_bytes: u64, + budget_bytes: u64, + }, + MalformedDigest, + /// Admission hashing found the wire's BLAKE3 different from the digest + /// the payload claims (the L3 manifest key). + DigestMismatch { + expected: String, + actual: String, + }, + /// A layout segment's claimed digest does not match the BLAKE3 of its + /// exact wire range: the pool identity would not describe the bytes it + /// is supposed to serve. + SegmentDigestMismatch { + digest: String, + expected: String, + actual: String, + }, + /// A layout claims a segment digest for two different wire ranges, or + /// claims a digest the pool already holds with different content that + /// another live entry still references. Content-addressed identity must + /// stay unambiguous. + ConflictingSegment { + digest: String, + detail: String, + }, + UnknownPayloadKind(String), + MalformedManifest(String), +} + +impl L2InsertRefusal { + pub fn reason(&self) -> String { + match self { + Self::EmptyPayload => "refusing to cache an empty exact-state payload".to_string(), + Self::OverBudget { payload_bytes } => format!( + "payload of {payload_bytes} distinct bytes exceeds the entire L2 budget; \ + caching it would evict everything else" + ), + Self::ProtectedOvercommit { + pool_bytes, + budget_bytes, + } => format!( + "admission would leave the pool at {pool_bytes} bytes against a \ + {budget_bytes}-byte budget even after evicting every unpinned entry: \ + the admission's own (shared or pinned) segments are not evictable, \ + so it was rolled back" + ), + Self::MalformedDigest => { + "payload digest is not a 64-hex-character blake3 string".to_string() + } + Self::DigestMismatch { expected, actual } => format!( + "admission digest check failed: wire hashes to {actual} but the payload \ + claims {expected}" + ), + Self::SegmentDigestMismatch { + digest, + expected, + actual, + } => format!( + "segment {digest} does not describe its wire range: range hashes to {actual} \ + but the layout claims {expected}" + ), + Self::ConflictingSegment { digest, detail } => { + format!("conflicting claims for segment {digest}: {detail}") + } + Self::UnknownPayloadKind(kind) => { + format!("manifest holds unknown payload kind {kind}") + } + Self::MalformedManifest(detail) => { + format!("malformed L3 manifest: {detail}") + } + } + } +} + +/// An immutable segment: content-addressed bytes shared by `Arc`. +#[derive(Debug, Clone)] +pub(crate) struct SegmentHandle { + pub bytes: Arc>, +} + +#[derive(Debug)] +struct L2Entry { + payload: ExactStatePayloadMirror, + token_count: u64, + payload_digest: String, + origin: L2Origin, + /// LRU clock, bumped on successful hits only (probes are side-effect + /// free). + last_used: u64, + /// Distinct segment bytes charged against the budget. Shared segments + /// already held by other entries are not charged here. + charge_bytes: u64, + /// Total wire length including shared segments (telemetry). + payload_bytes: u64, +} + +#[derive(Default)] +struct L2Inner { + map: HashMap, + /// Content-addressed pool of immutable segments. + segments: HashMap, + /// Distinct segment bytes in the pool — the real RAM footprint. + bytes: u64, + clock: u64, +} + +/// Undo log for one admission transaction. Every pool/map mutation an +/// admission performs — reserving new handles, overwriting a digest with +/// new content, releasing orphaned handles, removing the replaced entry, +/// and the entries eviction removes — is recorded here so a refused +/// admission (protected overcommit) can restore the tier exactly. Charges +/// are not journaled: every exit path recomputes them from the live map. +#[derive(Default)] +struct AdmitJournal { + reserved: Vec, + /// Previous `Arc` handles overwritten by this admission's same-digest + /// new-content installs, restored on rollback. + overwrites: Vec<(String, SegmentHandle)>, + /// Pool bytes that left with the overwritten handles. + overwritten_bytes: u64, + /// Released-orphan handles: `(digest, handle)` pairs to reinstall on + /// rollback. + released_orphans: Vec<(String, SegmentHandle)>, + removed_entries: Vec<(String, L2Entry)>, + reserved_bytes: u64, + released_bytes: u64, +} + +impl AdmitJournal { + fn rollback(self, inner: &mut L2Inner) { + for (digest, handle) in self.overwrites { + inner.segments.insert(digest, handle); + } + for digest in &self.reserved { + inner.segments.remove(digest); + } + for (digest, handle) in self.released_orphans { + inner.segments.insert(digest, handle); + } + for (key, entry) in self.removed_entries { + inner.map.insert(key, entry); + } + inner.bytes = inner + .bytes + .saturating_add(self.released_bytes) + .saturating_add(self.overwritten_bytes) + .saturating_sub(self.reserved_bytes); + } +} + +/// Counters kept outside the map lock so `stats()` never blocks hits. +#[derive(Default)] +struct L2AtomicStats { + hits: AtomicU64, + misses: AtomicU64, + inserts: AtomicU64, + evictions: AtomicU64, + admission_rejects: AtomicU64, + refused_bytes: AtomicU64, + shared_bytes_admitted: AtomicU64, +} + +/// Bounded host-RAM L2 over immutable packed L3 segments. +pub struct L2Tier { + inner: Mutex, + budget_bytes: u64, + stats: L2AtomicStats, +} + +/// A digest string must be a BLAKE3 hex digest: `blake3:`-prefixed (as L3 +/// digests are) or bare 64 hex characters. +fn is_valid_digest(digest: &str) -> bool { + let hex = digest.strip_prefix("blake3:").unwrap_or(digest); + hex.len() == 64 && hex.bytes().all(|b| b.is_ascii_hexdigit()) +} + +/// One validated layout segment: its claimed digest and the exact verified +/// slice of the wire it describes. +type ValidatedSegment<'a> = (&'a str, &'a [u8]); + +/// Fully validate a mirror's layout against the verified wire *before* any +/// pool mutation: +/// +/// - the tiling is non-empty and contiguous over `0..total_bytes`; +/// - `kv_bytes + recurrent_bytes == total_bytes`; +/// - every segment digest is well-formed and hashes its exact wire range. +/// +/// A layout that fails any check is refused (`MalformedManifest` or +/// `SegmentDigestMismatch`): the pool is content-addressed, so a handle's +/// digest must describe the bytes it serves. +fn validate_layout<'a>( + mirror: &'a ExactStatePayloadMirror, + wire: &'a [u8], +) -> Result>, L2InsertRefusal> { + let layout = mirror.layout(); + if layout.segments.is_empty() { + return Err(L2InsertRefusal::MalformedManifest( + "layout holds no segments".to_string(), + )); + } + if layout.total_bytes != wire.len() as u64 { + return Err(L2InsertRefusal::MalformedManifest(format!( + "layout claims {total} total bytes but the wire holds {len}", + total = layout.total_bytes, + len = wire.len() + ))); + } + if layout.kv_bytes.saturating_add(layout.recurrent_bytes) != layout.total_bytes { + return Err(L2InsertRefusal::MalformedManifest(format!( + "kv ({kv}) + recurrent ({rec}) bytes do not tile the {total}-byte payload", + kv = layout.kv_bytes, + rec = layout.recurrent_bytes, + total = layout.total_bytes + ))); + } + let mut expected_start = 0u64; + let mut validated = Vec::with_capacity(layout.segments.len()); + for (digest, range) in &layout.segments { + if !is_valid_digest(digest) { + return Err(L2InsertRefusal::MalformedManifest(format!( + "segment digest {digest:?} is not a blake3 hex digest" + ))); + } + if range.start != expected_start || range.end < range.start || range.end > wire.len() as u64 + { + return Err(L2InsertRefusal::MalformedManifest(format!( + "segment range {range:?} does not contiguously tile the wire at offset \ + {expected_start}" + ))); + } + expected_start = range.end; + let slice = &wire[range.start as usize..range.end as usize]; + let actual = segment_digest(slice); + if actual != *digest { + return Err(L2InsertRefusal::SegmentDigestMismatch { + digest: digest.clone(), + expected: digest.clone(), + actual, + }); + } + validated.push((digest.as_str(), slice)); + } + if expected_start != layout.total_bytes { + return Err(L2InsertRefusal::MalformedManifest(format!( + "segments tile {expected_start} bytes but the layout claims {}", + layout.total_bytes + ))); + } + Ok(validated) +} +impl L2Tier { + pub fn new(budget_bytes: u64) -> Self { + Self { + inner: Mutex::new(L2Inner::default()), + budget_bytes, + stats: L2AtomicStats::default(), + } + } + + pub fn budget_bytes(&self) -> u64 { + self.budget_bytes + } + + /// Admit an assembled entry. + /// + /// `wire` is the payload's concatenated L3 wire — the exact bytes whose + /// BLAKE3 is the manifest key. Admission verifies + /// `segment_digest(&wire) == payload_digest` and refuses the insert on + /// mismatch: L2 never holds bytes it did not verify. Every layout + /// segment's digest is verified against its exact wire range before the + /// pool is touched, so a handle can never serve bytes its digest does + /// not describe. After admission the segment bytes are immutable, so + /// reads are a digest lookup plus handle assembly — no re-hash. + /// + /// The admission is atomic with respect to the segment pool: the + /// incoming layout's segments are installed (or reserved) *before* the + /// previous entry at the same key is released and before eviction runs, + /// with those handles pinned against removal, so a replacement or a + /// sharing admission can never release a handle it is about to + /// reference. Returns the evictions the admission caused. The budget is + /// charged with the entry's *distinct* segment bytes: segments already + /// held by another entry are shared, not duplicated, and only the first + /// admission pays for them. + pub fn admit( + &self, + cache_key: String, + token_count: u64, + payload_digest: String, + wire: &[u8], + mirror: ExactStatePayloadMirror, + origin: L2Origin, + ) -> Result, L2InsertRefusal> { + if !is_valid_digest(&payload_digest) { + self.stats.admission_rejects.fetch_add(1, Ordering::Relaxed); + return Err(L2InsertRefusal::MalformedDigest); + } + // The one integrity check on the whole wire: it must hash to the + // claimed manifest-key digest. + let actual = segment_digest(wire); + if actual != payload_digest { + self.stats.admission_rejects.fetch_add(1, Ordering::Relaxed); + return Err(L2InsertRefusal::DigestMismatch { + expected: payload_digest, + actual, + }); + } + let payload_bytes = mirror.byte_len(); + if payload_bytes == 0 { + // Mirrors L3: an empty payload cannot represent state. + return Err(L2InsertRefusal::EmptyPayload); + } + if payload_bytes != wire.len() as u64 { + return Err(L2InsertRefusal::MalformedManifest(format!( + "mirror claims {payload_bytes} payload bytes but the verified wire holds {}", + wire.len() + ))); + } + // Every segment slice is hashed against its claimed digest before + // any pool mutation: pool reuse trusts content, not digest text. + let segments = validate_layout(&mirror, wire)?; + + let mut inner = self.inner.lock().expect("L2 map lock poisoned"); + // Distinct-byte charge: segments the pool already holds with + // exactly the verified content are shared (anything else conflicts + // or is new). `shared` pins every handle this admission references + // — including segments still owned by the entry being replaced — + // so release and eviction below cannot drop them out from under + // the transaction. + let mut shared: Vec = Vec::new(); + let mut new_segments: Vec<(String, SegmentHandle)> = Vec::new(); + let mut new_bytes = 0u64; + let mut shared_bytes = 0u64; + for &(digest, slice) in segments.iter() { + if new_segments.iter().any(|(d, _)| d == digest) { + // Within-admission duplicate digest: the validation pass + // already proved both slices have identical content, so + // keep the first copy. + shared_bytes += slice.len() as u64; + continue; + } + match inner.segments.get(digest) { + // Pool already holds this exact content: share it, whoever + // currently owns it. + Some(handle) if handle.bytes.as_ref() == slice => { + shared_bytes += slice.len() as u64; + shared.push(digest.to_string()); + } + // Same digest text, different bytes in the pool. + Some(_) => { + // Replacing an entry at the same key legitimately + // re-uses a digest with new content: the old owner is + // about to be released. Anything else is a conflict — + // a stale handle another live entry still references + // must never be swapped underneath it. + let replacing_same_key = inner + .map + .get(&cache_key) + .is_some_and(|entry| entry.payload.segment_digests().contains(&digest)); + if !replacing_same_key { + return Err(L2InsertRefusal::ConflictingSegment { + digest: digest.to_string(), + detail: "the pool holds different bytes under this digest \ + for another live entry" + .to_string(), + }); + } + new_bytes = new_bytes.saturating_add(slice.len() as u64); + new_segments.push(( + digest.to_string(), + SegmentHandle { + bytes: Arc::new(slice.to_vec()), + }, + )); + } + None => { + new_bytes = new_bytes.saturating_add(slice.len() as u64); + new_segments.push(( + digest.to_string(), + SegmentHandle { + bytes: Arc::new(slice.to_vec()), + }, + )); + } + } + } + if new_bytes > self.budget_bytes { + self.stats + .refused_bytes + .fetch_add(new_bytes, Ordering::Relaxed); + return Err(L2InsertRefusal::OverBudget { + payload_bytes: new_bytes, + }); + } + // Admission is all-or-nothing. Every mutation from here is + // journaled (`AdmitJournal`, module-level); if eviction cannot + // bring the final pool footprint under budget (the incoming + // entry's own handles are pinned, so an admission sharing bytes + // with its victim can exceed what eviction frees), the journal is + // rolled back and the admission is refused without touching the + // tier. + let mut journal = AdmitJournal { + reserved: Vec::new(), + overwrites: Vec::new(), + overwritten_bytes: 0, + released_orphans: Vec::new(), + removed_entries: Vec::new(), + reserved_bytes: 0, + released_bytes: 0, + }; + let protected_set: Vec = new_segments + .iter() + .map(|(digest, _)| digest.clone()) + .chain(shared.iter().cloned()) + .collect(); + // Reserve the new handles in the pool before releasing anything, + // so a digest re-used with new content is unambiguous from here on + // and the incoming entry's bytes cannot be dropped mid-transaction. + for (digest, handle) in &new_segments { + let handle_len = handle.bytes.len() as u64; + inner.bytes = inner.bytes.saturating_add(handle_len); + journal.reserved_bytes = journal.reserved_bytes.saturating_add(handle_len); + if let Some(previous) = inner.segments.insert(digest.clone(), handle.clone()) { + // Same digest text re-used with new content: only possible + // when replacing the same key, which still holds the old + // handle. The previous bytes leave the pool now (net + // reserved delta is `new − old`); journal them so rollback + // restores the original count. + let previous_len = previous.bytes.len() as u64; + inner.bytes = inner.bytes.saturating_sub(previous_len); + journal.overwritten_bytes = journal.overwritten_bytes.saturating_add(previous_len); + journal.overwrites.push((digest.clone(), previous)); + } else { + journal.reserved.push(digest.clone()); + } + } + // One entry per cache key: a re-admit at the same coordinates is a + // replacement (fresher state for the same prefix), not a duplicate. + // The old entry's segments survive release where the incoming + // layout shares them (`protected_set`), so identical-wire re-admits + // never delete their own handles. Released orphan handles are + // journaled so rollback reinstates them, and segments that stay + // are transfer-charged to their surviving owners before the new + // entry lands. + if let Some(existing) = inner.map.remove(&cache_key) { + let digests = existing.payload.segment_digests(); + Self::recompute_all_charges(&mut inner); + for digest in digests { + if protected_set.iter().any(|p| p == digest) { + continue; + } + let still_referenced = inner + .map + .values() + .any(|other| other.payload.segment_digests().contains(&digest)); + if still_referenced { + continue; + } + if let Some(handle) = inner.segments.remove(digest) { + let released = handle.bytes.len() as u64; + inner.bytes = inner.bytes.saturating_sub(released); + journal.released_bytes = journal.released_bytes.saturating_add(released); + journal.released_orphans.push((digest.to_string(), handle)); + } + } + journal.removed_entries.push((cache_key.clone(), existing)); + } + // Evict to make room: the reservation already counts toward + // `inner.bytes`, so the pool (including this admission's distinct + // bytes) must fit the whole budget. Shared handles are pinned and + // can keep a victim from freeing — those retained bytes transfer + // to this entry's charge below. + let evictions = self.evict_to_limit( + &mut inner, + self.budget_bytes, + &cache_key, + &protected_set, + &mut journal, + ); + // The hard byte cap: if eviction could not free the excess even by + // evicting every non-pinned entry, the admission is refused and + // rolled back — the tier never sits over budget after `admit`. + if inner.bytes > self.budget_bytes { + let pool_bytes = inner.bytes; + let rolled_back_evictions = evictions.len() as u64; + journal.rollback(&mut inner); + // Charges were recomputed during eviction; restore them to + // match the rolled-back state. + Self::recompute_all_charges(&mut inner); + self.stats + .evictions + .fetch_sub(rolled_back_evictions, Ordering::Relaxed); + let over = pool_bytes.saturating_sub(self.budget_bytes); + self.stats.refused_bytes.fetch_add(over, Ordering::Relaxed); + return Err(L2InsertRefusal::ProtectedOvercommit { + pool_bytes, + budget_bytes: self.budget_bytes, + }); + } + // The entry lands with a zero charge; the deterministic recompute + // below assigns it exactly the pooled segments it owns (segments + // whose lowest-key live reference it is) and refreshes every other + // entry's charge, so the sum of charges always equals the pool. + inner.clock = inner.clock.wrapping_add(1); + let last_used = inner.clock; + self.stats + .shared_bytes_admitted + .fetch_add(shared_bytes, Ordering::Relaxed); + inner.map.insert( + cache_key.clone(), + L2Entry { + payload: mirror, + token_count, + payload_digest, + origin, + last_used, + charge_bytes: 0, + payload_bytes, + }, + ); + // Assign every pooled segment exactly once to its lowest-key live + // reference — the inserted entry included. + Self::recompute_all_charges(&mut inner); + self.stats.inserts.fetch_add(1, Ordering::Relaxed); + Ok(evictions) + } + + /// A verified hit records recency and returns the entry's layout with + /// `Arc` clones of its segment handles (no byte copies). Digests are + /// not re-hashed: admission verified the wire, and segments are + /// immutable afterward. + /// + /// If a segment handle the entry references is missing from the pool, + /// the entry is corrupt: the hit is downgraded to a miss, the entry and + /// its surviving segments are removed, LRU recency never moves, and the + /// miss counter is incremented. + pub fn get(&self, cache_key: &str) -> Option { + let mut inner = self.inner.lock().expect("L2 map lock poisoned"); + let digests: Vec = match inner.map.get(cache_key) { + Some(entry) => entry + .payload + .segment_digests() + .into_iter() + .map(str::to_string) + .collect(), + // Absent key: a miss, never an LRU touch. + None => { + self.stats.misses.fetch_add(1, Ordering::Relaxed); + return None; + } + }; + let missing = digests + .iter() + .any(|digest| !inner.segments.contains_key(digest)); + if missing { + // Corrupt entry: a segment handle vanished without an entry + // removal. Serve a miss, never partial bytes; drop the entry + // and its surviving handles; recency stays untouched. + let removed = inner.map.remove(cache_key); + if let Some(entry) = removed { + Self::recompute_all_charges(&mut inner); + self.release_entry_segments(&mut inner, &entry, &[]); + } + self.stats.misses.fetch_add(1, Ordering::Relaxed); + return None; + } + let now = { + inner.clock = inner.clock.wrapping_add(1); + inner.clock + }; + let entry = inner.map.get_mut(cache_key).expect("entry checked above"); + entry.last_used = now; + let payload = entry.payload.clone(); + let token_count = entry.token_count; + let payload_digest = entry.payload_digest.clone(); + self.stats.hits.fetch_add(1, Ordering::Relaxed); + let mut segments = HashMap::with_capacity(digests.len()); + for digest in &digests { + let handle = inner + .segments + .get(digest) + .expect("all digests checked above"); + segments.insert(digest.clone(), handle.clone()); + } + Some(L2Hit { + payload, + token_count, + payload_digest, + segments, + }) + } + + /// Presence probe: side-effect free. It does not touch LRU recency — + /// prefix probing must not make entries hot — and returns only + /// metadata. Recency is updated by `get` after a successful verified + /// hit. + pub fn peek(&self, cache_key: &str) -> Option { + let inner = self.inner.lock().expect("L2 map lock poisoned"); + let entry = inner.map.get(cache_key)?; + Some(L2Peek { + token_count: entry.token_count, + payload_digest: entry.payload_digest.clone(), + payload_bytes: entry.payload_bytes, + distinct_bytes: entry.charge_bytes, + origin: entry.origin, + }) + } + + pub fn remove(&self, cache_key: &str) -> Option { + let mut inner = self.inner.lock().expect("L2 map poisoned"); + let removed = inner.map.remove(cache_key)?; + // Retained bytes come from the pool's actual references: segments + // of this entry that other entries still reference after the + // removal stay in the pool. Charged bytes are never transferred + // between entries — a survivor's charge already excludes shared + // segments — so this cannot drift the survivors' accounting below + // the physical pool they own. + let mut retained = 0u64; + for digest in removed.payload.segment_digests() { + let referenced_elsewhere = inner + .map + .values() + .any(|other| other.payload.segment_digests().contains(&digest)); + if referenced_elsewhere { + retained = retained.saturating_add( + inner + .segments + .get(digest) + .map(|handle| handle.bytes.len() as u64) + .unwrap_or(0), + ); + } + } + // Segments that stay are now physically owned by the survivors: + // recompute charges from the live map (the removed entry is + // already out of it). + Self::recompute_all_charges(&mut inner); + let before = inner.bytes; + self.release_entry_segments(&mut inner, &removed, &[]); + let freed = before.saturating_sub(inner.bytes); + Some(L2Eviction { + cache_key: cache_key.to_string(), + freed_bytes: freed, + retained_bytes: retained, + }) + } + + /// Drop everything; returns the distinct bytes released. + pub fn clear(&self) -> u64 { + let mut inner = self.inner.lock().expect("L2 map lock poisoned"); + let bytes = inner.bytes; + inner.map.clear(); + inner.segments.clear(); + inner.bytes = 0; + bytes + } + + pub fn len(&self) -> usize { + self.inner.lock().expect("L2 map lock poisoned").map.len() + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Test-only: drop every pool handle without touching the entry map, + /// to force the missing-handle corruption path in `get`. + #[cfg(test)] + fn clear_pool_for_test(&self) { + let mut inner = self.inner.lock().expect("L2 map lock poisoned"); + inner.segments.clear(); + inner.bytes = 0; + } + + /// Point-in-time snapshot combining atomics with the locked totals. + pub fn stats(&self) -> L2Stats { + let inner = self.inner.lock().expect("L2 map lock poisoned"); + let logical_bytes = inner + .map + .values() + .map(|entry| entry.payload_bytes) + .sum::(); + L2Stats { + entries: inner.map.len() as u64, + bytes: inner.bytes, + logical_bytes, + segments: inner.segments.len() as u64, + segment_bytes: inner.bytes, + shared_bytes_admitted: self.stats.shared_bytes_admitted.load(Ordering::Relaxed), + budget_bytes: self.budget_bytes, + hits: self.stats.hits.load(Ordering::Relaxed), + misses: self.stats.misses.load(Ordering::Relaxed), + inserts: self.stats.inserts.load(Ordering::Relaxed), + evictions: self.stats.evictions.load(Ordering::Relaxed), + admission_rejects: self.stats.admission_rejects.load(Ordering::Relaxed), + refused_bytes: self.stats.refused_bytes.load(Ordering::Relaxed), + } + } + + /// Recompute every entry's charge from the live map: each pooled + /// segment is assigned exactly once, to its lowest-cache-key live + /// reference, and each entry's `charge_bytes` is the sum of the + /// segments it owns. The sum of all charges therefore always equals + /// the physical pool bytes — after admission, eviction, removal, + /// corruption cleanup, and rollback alike. Deterministic: identical + /// map states produce identical ownership. + fn recompute_all_charges(inner: &mut L2Inner) { + let mut ownership: HashMap = HashMap::new(); + for digest in inner.segments.keys() { + let Some(bytes) = inner.segments.get(digest).map(|h| h.bytes.len() as u64) else { + continue; + }; + let owner = inner + .map + .iter() + .filter(|(_, entry)| entry.payload.segment_digests().contains(&digest.as_str())) + .map(|(key, _)| key.clone()) + .min(); + if let Some(owner) = owner { + *ownership.entry(owner).or_insert(0) += bytes; + } + } + for (key, entry) in inner.map.iter_mut() { + entry.charge_bytes = ownership.get(key).copied().unwrap_or(0); + } + } + + /// Drop an entry's segments from the pool, decrementing the pool byte + /// total. Segments still referenced by another live entry, or pinned by + /// an in-flight admission (`protected`), stay. Zero-byte segments are + /// dropped without accounting (a pool without payload bytes must never + /// charge the budget). + fn release_entry_segments(&self, inner: &mut L2Inner, entry: &L2Entry, protected: &[String]) { + for digest in entry.payload.segment_digests() { + if protected.iter().any(|p| p == digest) { + continue; + } + let still_referenced = inner + .map + .values() + .any(|other| other.payload.segment_digests().contains(&digest)); + if still_referenced { + continue; + } + if let Some(handle) = inner.segments.remove(digest) { + let released = handle.bytes.len() as u64; + inner.bytes = inner.bytes.saturating_sub(released); + } + } + } + + /// Evict in deterministic LRU order until the pool fits `limit` bytes. + /// Shared segments are released only with their last referencing + /// entry; handles pinned by the in-flight admission (`protected`) are + /// never released; a victim that frees nothing is still counted as an + /// eviction. Every mutation is recorded in `journal` so a refused + /// admission can roll the evictions back exactly. + fn evict_to_limit( + &self, + inner: &mut L2Inner, + limit: u64, + protect_key: &str, + protected: &[String], + journal: &mut AdmitJournal, + ) -> Vec { + let mut evictions = Vec::new(); + while inner.bytes > limit { + // Deterministic LRU: lowest last_used wins; ties break on cache + // key so identical operation sequences produce identical + // evictions. + let victim = inner + .map + .iter() + .filter(|(key, _)| key.as_str() != protect_key) + .min_by(|a, b| a.1.last_used.cmp(&b.1.last_used).then_with(|| a.0.cmp(b.0))) + .map(|(key, _)| key.clone()); + let Some(victim) = victim else { break }; + let Some(removed) = inner.map.remove(&victim) else { + break; + }; + // Retained bytes from actual pool references, computed before + // release: segments of the victim that survivors still + // reference, or that the in-flight admission pins (its entry + // is not in the map yet, but it will own them). + let mut retained = 0u64; + for digest in removed.payload.segment_digests() { + let referenced_elsewhere = inner + .map + .values() + .any(|other| other.payload.segment_digests().contains(&digest)) + || protected.iter().any(|p| p == digest); + if referenced_elsewhere { + retained = retained.saturating_add( + inner + .segments + .get(digest) + .map(|handle| handle.bytes.len() as u64) + .unwrap_or(0), + ); + } + } + let before = inner.bytes; + // Charges are recomputed from the live map (journaled state + // is restored exactly; ownership follows the lowest key). + Self::recompute_all_charges(inner); + for digest in removed.payload.segment_digests() { + if protected.iter().any(|p| p == digest) { + continue; + } + let still_referenced = inner + .map + .values() + .any(|other| other.payload.segment_digests().contains(&digest)); + if still_referenced { + continue; + } + if let Some(handle) = inner.segments.remove(digest) { + let released = handle.bytes.len() as u64; + inner.bytes = inner.bytes.saturating_sub(released); + journal.released_bytes = journal.released_bytes.saturating_add(released); + journal.released_orphans.push((digest.to_string(), handle)); + } + } + let freed = before.saturating_sub(inner.bytes); + self.stats.evictions.fetch_add(1, Ordering::Relaxed); + journal.removed_entries.push((victim.clone(), removed)); + evictions.push(L2Eviction { + cache_key: victim, + freed_bytes: freed, + retained_bytes: retained, + }); + } + evictions + } +} + +/// Build the L2 cache key from the same coordinates L3 uses, plus the +/// identities the tier serves. Same coordinates under different identities +/// get different keys: an identity change cannot cross-contaminate. +pub fn l2_cache_key( + model_identity: &str, + state_identity: &str, + namespace: &str, + token_ids: &[i32], +) -> String { + let namespace_key = crate::tier::l3_namespace_key(namespace); + let prefix_key = crate::tier::l3_prefix_key(namespace, token_ids); + let mut hasher = blake3::Hasher::new(); + hasher.update(b"l2-cache-key-v1"); + hasher.update(model_identity.as_bytes()); + hasher.update(b"\0"); + hasher.update(state_identity.as_bytes()); + hasher.update(b"\0"); + hasher.update(namespace_key.as_bytes()); + hasher.update(prefix_key.as_bytes()); + format!("blake3:{}", hasher.finalize().to_hex()) +} + +#[cfg(test)] +mod tests { + use super::*; + + const DIGEST_B: &str = "b616719e0a0d39dc0fe85cd2d0a5e0e2f5e6e10b6b5a0a6f1a1c1d3e5f708a90"; + + thread_local! { + static FULL_WIRE: std::cell::RefCell> = const { std::cell::RefCell::new(Vec::new()) }; + } + + /// Build a synthetic wire of `len` bytes and its true BLAKE3 digest. + fn wire(len: usize, fill: u8) -> (Vec, String) { + // Position-dependent bytes so equal-length slices are never equal + // content: segment digests stay distinct across the wire. + let bytes: Vec = (0..len) + .map(|i| (fill as usize + i) % 251) + .map(|v| v as u8) + .collect(); + let digest = segment_digest(&bytes); + FULL_WIRE.with(|cell| *cell.borrow_mut() = bytes.clone()); + (bytes, digest) + } + + /// Single-segment full-state mirror over `len` wire bytes. The segment + /// key is the content digest of the whole wire, as L3 would produce. + fn single_segment_mirror(w: &[u8]) -> ExactStatePayloadMirror { + let len = w.len() as u64; + ExactStatePayloadMirror::FullState { + layout: L2Layout { + payload_kind: ExactStatePayloadKind::FullState, + total_bytes: len, + kv_bytes: len, + recurrent_bytes: 0, + segments: vec![(segment_digest(w), 0..len)], + }, + } + } + + /// Manifest-shaped mirror: digest-keyed segments cut at every + /// `segment_len` boundary, matching how `from_manifest` tiles. + fn manifest_shaped_mirror(w: &[u8], segment_len: u64) -> ExactStatePayloadMirror { + // `w` is a suffix of the test's full wire: segment digests are + // keyed by offset in that full wire so entries sharing a prefix + // also share segment identity. + let full = FULL_WIRE.with(|cell| cell.borrow().clone()); + let len = w.len() as u64; + let mut segments = Vec::new(); + let mut offset = 0u64; + while offset < len { + let end = (offset + segment_len).min(len); + let digest = segment_digest(&full[offset as usize..end as usize]); + segments.push((digest, offset..end)); + offset = end; + } + ExactStatePayloadMirror::FullState { + layout: L2Layout { + payload_kind: ExactStatePayloadKind::FullState, + total_bytes: len, + kv_bytes: len, + recurrent_bytes: 0, + segments, + }, + } + } + + fn key(namespace: &str, tokens: &[i32]) -> String { + l2_cache_key("model-a", "state-a", namespace, tokens) + } + + #[test] + fn admit_get_round_trip_serves_wire_bytes() { + let tier = L2Tier::new(1 << 20); + let k = key("ns", &[1, 2, 3]); + let (w, digest) = wire(64, 7); + tier.admit( + k.clone(), + 3, + digest.clone(), + &w, + single_segment_mirror(&w), + L2Origin::FromL3, + ) + .expect("admission must fit"); + let hit = tier.get(&k).expect("admitted key must hit"); + assert_eq!(hit.token_count, 3); + assert_eq!(hit.payload.byte_len(), 64); + assert_eq!(hit.payload_digest, digest); + // Round-trips into a serving payload with the right byte count and + // exactly the admitted wire bytes. + let payload = hit.to_payload(); + assert_eq!(payload.byte_len(), 64); + assert_eq!( + payload.kind(), + crate::payload::ExactStatePayloadKind::FullState + ); + let (bytes, _) = payload.full_state_bytes_timed().expect("full state"); + assert_eq!(bytes.as_ref(), &w[..], "served bytes must equal the wire"); + } + + #[test] + fn admission_digest_mismatch_refuses_and_stores_nothing() { + let tier = L2Tier::new(1 << 20); + let k = key("ns", &[1, 2, 3]); + let (w, _) = wire(64, 7); + let err = tier + .admit( + k.clone(), + 3, + DIGEST_B.to_string(), + &w, + single_segment_mirror(&w), + L2Origin::Direct, + ) + .expect_err("a wire that does not hash to the claimed digest must be refused"); + assert!(matches!(err, L2InsertRefusal::DigestMismatch { .. })); + assert!(tier.peek(&k).is_none(), "refused bytes must not be stored"); + assert!(tier.get(&k).is_none()); + let stats = tier.stats(); + assert_eq!(stats.admission_rejects, 1); + assert_eq!(stats.entries, 0); + assert_eq!(stats.bytes, 0); + } + + #[test] + fn corrupted_wire_never_reaches_the_tier() { + // L2Origin::Direct with arbitrary bytes under a valid-looking + // digest is exactly the hole this closes: the digest check runs on + // the actual bytes. + let tier = L2Tier::new(1 << 20); + let k = key("ns", &[9]); + let (mut w, digest) = wire(128, 3); + w[42] ^= 0xff; // one flipped bit + let err = tier + .admit( + k.clone(), + 1, + digest, + &w, + single_segment_mirror(&w), + L2Origin::Direct, + ) + .expect_err("corrupted wire must be refused at admission"); + assert!(matches!(err, L2InsertRefusal::DigestMismatch { .. })); + assert_eq!(tier.stats().admission_rejects, 1); + assert!(tier.is_empty()); + } + + #[test] + fn peek_is_side_effect_free_for_lru() { + let tier = L2Tier::new(160); + let k1 = key("ns", &[1]); + let k2 = key("ns", &[2]); + let (w1, d1) = wire(64, 1); + let (w2, d2) = wire(64, 2); + tier.admit( + k1.clone(), + 1, + d1, + &w1, + single_segment_mirror(&w1), + L2Origin::FromL3, + ) + .expect("k1 fits"); + tier.admit( + k2.clone(), + 1, + d2, + &w2, + single_segment_mirror(&w2), + L2Origin::FromL3, + ) + .expect("k2 fits"); + // Probe k1 many times: recency must not move. + for _ in 0..10 { + assert!(tier.peek(&k1).is_some()); + } + // Insert a third entry: k1 (never truly used) must still be the + // LRU victim, not k2. + let k3 = key("ns", &[3]); + let (w3, d3) = wire(64, 3); + let evictions = tier + .admit( + k3.clone(), + 1, + d3, + &w3, + single_segment_mirror(&w3), + L2Origin::FromL3, + ) + .expect("k3 fits after eviction"); + assert_eq!(evictions.len(), 1); + assert_eq!(evictions[0].cache_key, k1, "probed-but-unused k1 is LRU"); + assert!(tier.peek(&k2).is_some(), "untouched k2 survives"); + } + + #[test] + fn recency_moves_only_on_verified_hit() { + let tier = L2Tier::new(160); + let k1 = key("ns", &[1]); + let k2 = key("ns", &[2]); + let (w1, d1) = wire(64, 1); + let (w2, d2) = wire(64, 2); + tier.admit( + k1.clone(), + 1, + d1, + &w1, + single_segment_mirror(&w1), + L2Origin::FromL3, + ) + .expect("k1 fits"); + tier.admit( + k2.clone(), + 1, + d2, + &w2, + single_segment_mirror(&w2), + L2Origin::FromL3, + ) + .expect("k2 fits"); + // A real hit on k1 makes k2 the victim of the next insertion. + assert!(tier.get(&k1).is_some()); + let k3 = key("ns", &[3]); + let (w3, d3) = wire(64, 3); + let evictions = tier + .admit(k3, 1, d3, &w3, single_segment_mirror(&w3), L2Origin::FromL3) + .expect("k3 fits after eviction"); + assert_eq!(evictions.len(), 1); + assert_eq!(evictions[0].cache_key, k2, "k2 is now LRU"); + assert!(tier.peek(&k1).is_some(), "recently hit k1 survives"); + } + + #[test] + fn prefix_growth_shares_segment_bytes_instead_of_duplicating() { + // 16 KiB of four 4 KiB segments; the shorter prefix shares the + // first three segments with the longer one. + let segment_len = 4096u64; + let total = segment_len * 4; + let tier = L2Tier::new(total * 2); + let short = key("ns", &[1, 2, 3]); + let long = key("ns", &[1, 2, 3, 4, 5]); + let (w, digest) = wire(total as usize, 5); + let short_len = segment_len * 3; + tier.admit( + short.clone(), + 3, + segment_digest(&w[..short_len as usize]), + &w[..short_len as usize], + manifest_shaped_mirror(&w[..short_len as usize], segment_len), + L2Origin::FromL3, + ) + .expect("short prefix admitted"); + + tier.admit( + long.clone(), + 5, + digest, + &w, + manifest_shaped_mirror(&w, segment_len), + L2Origin::FromL3, + ) + .expect("long prefix admitted"); + + let stats = tier.stats(); + // The long entry pays only for its one new (4th) segment. + assert_eq!( + stats.bytes, total, + "pool must hold distinct segment bytes once: got {}", + stats.bytes + ); + assert_eq!( + stats.logical_bytes, + total + short_len, + "logical bytes count both entries' full wires" + ); + assert_eq!(stats.segments, 4, "four distinct segments, not seven"); + assert_eq!(stats.shared_bytes_admitted, short_len); + // Both entries serve their own wire slices. + let hit = tier.get(&long).expect("long hit"); + let payload = hit.to_payload(); + let (bytes, _) = payload.full_state_bytes_timed().expect("bytes"); + assert_eq!(bytes.as_ref(), &w[..]); + let hit_short = tier.get(&short).expect("short hit"); + let payload_short = hit_short.to_payload(); + let (bytes_short, _) = payload_short.full_state_bytes_timed().expect("bytes"); + assert_eq!(bytes_short.as_ref(), &w[..short_len as usize]); + } + + #[test] + fn evicting_one_entry_keeps_shared_prefix_segments() { + let segment_len = 4096u64; + let total = segment_len * 4; + let tier = L2Tier::new(total * 2); + let short = key("ns", &[1]); + let long = key("ns", &[2]); + let (w, digest) = wire(total as usize, 6); + let short_len = segment_len * 3; + tier.admit( + short, + 3, + segment_digest(&w[..short_len as usize]), + &w[..short_len as usize], + manifest_shaped_mirror(&w[..short_len as usize], segment_len), + L2Origin::FromL3, + ) + .expect("short admitted"); + tier.admit( + long.clone(), + 5, + digest, + &w, + manifest_shaped_mirror(&w, segment_len), + L2Origin::FromL3, + ) + .expect("long admitted"); + // Removing the long entry frees only its exclusive tail segment. + let removed = tier.remove(&long).expect("long entry present"); + assert_eq!(removed.freed_bytes, segment_len); + assert_eq!( + removed.retained_bytes, short_len, + "shared prefix bytes are retained by the shorter entry" + ); + let stats = tier.stats(); + assert_eq!(stats.bytes, short_len); + assert_eq!(stats.segments, 3); + assert!(tier.get(&key("ns", &[1])).is_some(), "short entry intact"); + } + + #[test] + fn budget_evicts_lru_first_and_never_the_protected_entry() { + let tier = L2Tier::new(256); + let k1 = key("ns", &[1]); + let k2 = key("ns", &[2]); + let k3 = key("ns", &[3]); + let (w1, d1) = wire(100, 1); + let (w2, d2) = wire(100, 2); + let (w3, d3) = wire(100, 3); + tier.admit( + k1.clone(), + 1, + d1, + &w1, + single_segment_mirror(&w1), + L2Origin::FromL3, + ) + .expect("k1 fits"); + tier.admit( + k2.clone(), + 1, + d2, + &w2, + single_segment_mirror(&w2), + L2Origin::FromL3, + ) + .expect("k2 fits"); + // Touch k1 so k2 becomes the LRU victim. + assert!(tier.get(&k1).is_some()); + let evictions = tier + .admit( + k3.clone(), + 1, + d3, + &w3, + single_segment_mirror(&w3), + L2Origin::FromL3, + ) + .expect("k3 fits after eviction"); + assert_eq!( + evictions.len(), + 1, + "one entry must be evicted: {evictions:?}" + ); + assert_eq!(evictions[0].cache_key, k2, "LRU victim is k2"); + assert_eq!(evictions[0].freed_bytes, 100); + assert_eq!(evictions[0].retained_bytes, 0); + assert!(tier.peek(&k1).is_some(), "recently used k1 survives"); + assert!(tier.peek(&k3).is_some(), "just-admitted k3 survives"); + assert!(tier.peek(&k2).is_none(), "k2 was evicted"); + let stats = tier.stats(); + assert_eq!(stats.evictions, 1); + assert_eq!(stats.bytes, 200, "pool bytes must track survivors exactly"); + } + + #[test] + fn oversized_distinct_bytes_are_refused_without_evicting() { + let tier = L2Tier::new(128); + let k1 = key("ns", &[1]); + let (w1, d1) = wire(64, 1); + tier.admit( + k1.clone(), + 1, + d1, + &w1, + single_segment_mirror(&w1), + L2Origin::FromL3, + ) + .expect("fits"); + let (w2, d2) = wire(129, 2); + let err = tier + .admit( + key("ns", &[2]), + 1, + d2, + &w2, + single_segment_mirror(&w2), + L2Origin::FromL3, + ) + .expect_err("over-budget payload must be refused"); + assert_eq!(err, L2InsertRefusal::OverBudget { payload_bytes: 129 }); + assert!(tier.peek(&k1).is_some(), "refusal must not evict anything"); + assert_eq!(tier.stats().refused_bytes, 129); + } + + #[test] + fn empty_payload_is_refused_like_l3() { + let tier = L2Tier::new(1 << 20); + let empty_digest = segment_digest(&[]); + let err = tier + .admit( + key("ns", &[1]), + 1, + empty_digest, + &[], + single_segment_mirror(&[]), + L2Origin::FromL3, + ) + .expect_err("empty payloads must be refused"); + assert_eq!(err, L2InsertRefusal::EmptyPayload); + assert!(tier.is_empty()); + } + + #[test] + fn malformed_digest_is_refused_before_any_hashing() { + let tier = L2Tier::new(1 << 20); + let (w, _) = wire(16, 4); + let err = tier + .admit( + key("ns", &[1]), + 1, + "not-a-digest".to_string(), + &w, + single_segment_mirror(&w), + L2Origin::FromL3, + ) + .expect_err("malformed digest must be refused"); + assert_eq!(err, L2InsertRefusal::MalformedDigest); + } + + #[test] + fn readmit_replaces_and_keeps_accounting_exact() { + let tier = L2Tier::new(1 << 20); + let k = key("ns", &[9]); + let (w1, d1) = wire(100, 1); + let (w2, d2) = wire(40, 2); + tier.admit( + k.clone(), + 3, + d1, + &w1, + single_segment_mirror(&w1), + L2Origin::FromL3, + ) + .expect("first admission"); + let evictions = tier + .admit( + k.clone(), + 3, + d2.clone(), + &w2, + single_segment_mirror(&w2), + L2Origin::FromL3, + ) + .expect("replacement"); + assert!(evictions.is_empty()); + assert_eq!(tier.len(), 1); + assert_eq!( + tier.stats().bytes, + 40, + "replacement must release the old bytes" + ); + // New digest is the one served now. + let hit = tier.get(&k).expect("replacement hit"); + assert_eq!(hit.payload_digest, d2); + let payload = hit.to_payload(); + let (bytes, _) = payload.full_state_bytes_timed().expect("bytes"); + assert_eq!(bytes.as_ref(), &w2[..]); + } + + #[test] + fn identical_coordinates_under_different_identities_get_different_keys() { + let a = l2_cache_key("model-a", "state-a", "ns", &[1, 2]); + let b = l2_cache_key("model-b", "state-a", "ns", &[1, 2]); + let c = l2_cache_key("model-a", "state-b", "ns", &[1, 2]); + assert_ne!(a, b); + assert_ne!(a, c); + // Same coordinates, same identities: stable key. + let a2 = l2_cache_key("model-a", "state-a", "ns", &[1, 2]); + assert_eq!(a, a2); + // Different token paths differ. + assert_ne!(a, l2_cache_key("model-a", "state-a", "ns", &[1, 3])); + } + + #[test] + fn mirror_round_trips_every_payload_kind_from_manifest() { + // kv-recurrent: kv 24 bytes then recurrent 8, cut into two segments. + let (kv_wire, _) = wire(24, 3); + let (rec_wire, _) = wire(8, 4); + let wire_bytes: Vec = [kv_wire, rec_wire].concat(); + let digest = segment_digest(&wire_bytes); + let manifest = HandoffManifest { + version: MANIFEST_VERSION, + codec: Some(PayloadCodec::raw()), + model_identity: "blake3:model".to_string(), + state_identity: "blake3:state".to_string(), + payload_kind: "kv-recurrent".to_string(), + total_bytes: wire_bytes.len() as u64, + payload_digest: digest.clone(), + segments: vec![ + HandoffSegmentRef { + index: 0, + offset: 0, + bytes: 16, + digest: segment_digest(&wire_bytes[..16]), + codec_identity: Some(SegmentCodecIdentity::raw(16)), + meta_json: None, + }, + HandoffSegmentRef { + index: 1, + offset: 16, + bytes: 16, + digest: segment_digest(&wire_bytes[16..]), + codec_identity: Some(SegmentCodecIdentity::raw(16)), + meta_json: None, + }, + ], + kv_bytes: 24, + recurrent_bytes: 8, + kv_desc_json: None, + token_count: 4, + continuation_token: 0, + expected_tokens: Vec::new(), + }; + let mirror = ExactStatePayloadMirror::from_manifest(&manifest).expect("manifest parses"); + assert_eq!(mirror.byte_len(), 32); + assert_eq!(mirror.kind(), ExactStatePayloadKind::KvRecurrent); + + let tier = L2Tier::new(1 << 20); + let k = key("ns", &[7]); + tier.admit(k.clone(), 4, digest, &wire_bytes, mirror, L2Origin::FromL3) + .expect("fits"); + let hit = tier.get(&k).expect("hit"); + let payload = hit.to_payload(); + assert_eq!(payload.byte_len(), 32); + let served_kv = payload + .kv_bytes() + .expect("kv bytes") + .expect("kv-recurrent has kv") + .into_owned(); + let served_rec = payload + .recurrent_state_bytes() + .expect("recurrent bytes") + .into_owned(); + assert_eq!(served_kv, wire_bytes[..24], "kv slice must match the wire"); + assert_eq!(served_rec, wire_bytes[24..], "recurrent slice matches"); + } + + #[test] + fn from_manifest_rejects_unknown_kind_and_bad_tiling() { + let mut manifest = HandoffManifest { + version: MANIFEST_VERSION, + codec: Some(PayloadCodec::raw()), + model_identity: "m".to_string(), + state_identity: "s".to_string(), + payload_kind: "blob".to_string(), + total_bytes: 10, + payload_digest: "blake3:aa".to_string(), + segments: Vec::new(), + kv_bytes: 10, + recurrent_bytes: 0, + kv_desc_json: None, + token_count: 1, + continuation_token: 0, + expected_tokens: Vec::new(), + }; + let err = ExactStatePayloadMirror::from_manifest(&manifest) + .expect_err("unknown kind must be refused"); + assert!(matches!(err, L2InsertRefusal::UnknownPayloadKind(_))); + + manifest.payload_kind = "full-state".to_string(); + manifest.segments = vec![HandoffSegmentRef { + index: 0, + offset: 0, + bytes: 7, + digest: "blake3:seg".to_string(), + codec_identity: Some(SegmentCodecIdentity::raw(7)), + meta_json: None, + }]; + let err = ExactStatePayloadMirror::from_manifest(&manifest) + .expect_err("tiling mismatch must be refused"); + assert!(matches!(err, L2InsertRefusal::MalformedManifest(_))); + } + + #[test] + fn multi_segment_reads_reassemble_exact_wire() { + // 6 KiB in 1 KiB segments: every read path crosses many blocks. + let segment_len = 1024u64; + let total = segment_len * 6; + let (w, digest) = wire(total as usize, 9); + let tier = L2Tier::new(total * 2); + let k = key("ns", &[1]); + tier.admit( + k.clone(), + 6, + digest, + &w, + manifest_shaped_mirror(&w, segment_len), + L2Origin::FromL3, + ) + .expect("fits"); + let hit = tier.get(&k).expect("hit"); + let payload = hit.to_payload(); + let (bytes, reconstruct) = payload.full_state_bytes_timed().expect("bytes"); + assert_eq!(bytes.as_ref(), &w[..]); + // Multiple distinct segment storages materialize on read; the + // reconstruction length must equal the payload either way. + assert_eq!( + reconstruct.reconstruct_bytes, total, + "multi-segment reads materialize the wire exactly once" + ); + } + + #[test] + fn clear_releases_everything_and_reports_bytes() { + let tier = L2Tier::new(1 << 20); + for i in 0..5i32 { + let (w, d) = wire(64, i as u8 + 10); + tier.admit( + key("ns", &[i]), + 1, + d, + &w, + single_segment_mirror(&w), + L2Origin::FromL3, + ) + .expect("fits"); + } + let released = tier.clear(); + assert_eq!(released, 320); + assert!(tier.is_empty()); + assert_eq!(tier.stats().bytes, 0); + assert_eq!(tier.stats().segments, 0); + } + + #[test] + fn remove_is_exact() { + let tier = L2Tier::new(1 << 20); + let k = key("ns", &[4]); + let (w, d) = wire(64, 8); + tier.admit( + k.clone(), + 1, + d, + &w, + single_segment_mirror(&w), + L2Origin::FromL3, + ) + .expect("fits"); + let removed = tier.remove(&k).expect("present entry removes"); + assert_eq!(removed.freed_bytes, 64); + assert_eq!(removed.retained_bytes, 0); + assert!(tier.remove(&k).is_none(), "second remove is None"); + assert_eq!(tier.stats().bytes, 0); + assert_eq!(tier.stats().segments, 0); + } + + #[test] + fn identical_wire_same_key_readmit_keeps_its_own_segments() { + // The original failure: re-admitting an identical wire at the same + // key classified the existing pool segments as shared, released the + // old entry's last references, inserted no replacement handles, and + // left an entry whose segments were absent + // (`declared=14 restored=0 pool=0`). + let tier = L2Tier::new(1 << 20); + let k = key("ns", &[11]); + let segment_len = 16u64; + let (w, digest) = wire(64, 21); + for round in 0..3 { + let mirror = manifest_shaped_mirror(&w, segment_len); + tier.admit(k.clone(), 4, digest.clone(), &w, mirror, L2Origin::FromL3) + .expect("identical re-admit must be accepted"); + let hit = tier + .get(&k) + .unwrap_or_else(|| panic!("round {round}: re-admitted key must hit")); + let payload = hit.to_payload(); + let (bytes, _) = payload.full_state_bytes_timed().expect("bytes"); + assert_eq!( + bytes.as_ref(), + &w[..], + "round {round}: re-admitted entry must serve its full wire" + ); + let stats = tier.stats(); + assert_eq!(stats.entries, 1); + assert_eq!( + stats.segments, 4, + "round {round}: pool must still hold every segment" + ); + assert_eq!(stats.bytes, 64, "round {round}: pool bytes exact"); + } + } + + #[test] + fn eviction_cannot_release_segments_the_incoming_entry_shares() { + // Pressure case: the incoming entry shares its would-be victim's + // prefix segments. The victim is not protected by the cache-key + // filter (different key) and is not yet replaced in the map, so + // eviction could drop the shared handles before the new entry + // lands. The 64-byte budget equals the old entry's footprint, so + // the 16-byte tail reservation forces eviction mid-admission: + // only the victim's *exclusive* tail X frees (16 bytes), the + // shared prefix is pinned by the admission and survives + // byte-exact, and the pool lands exactly at budget with the new + // entry's segments. Layout: old = [S0 S1 S2 X], grown = + // [S0 S1 S2 T] with T different content from X. + let segment_len = 16u64; + let total = segment_len * 4; + let tier = L2Tier::new(total); + let (w, _) = wire(total as usize, 30); + let short_len = segment_len * 3; + + let old = key("ns", &[1]); + tier.admit( + old.clone(), + 4, + segment_digest(&w), + &w, + manifest_shaped_mirror(&w, segment_len), + L2Origin::FromL3, + ) + .expect("old entry admitted"); + + // The grown entry swaps the old tail for a different one. + let tail: Vec = (0..segment_len as usize) + .map(|i| (200usize + i) % 251) + .map(|v| v as u8) + .collect(); + let mut grown_wire = w[..short_len as usize].to_vec(); + grown_wire.extend_from_slice(&tail); + let grown_digest = segment_digest(&grown_wire); + let grown_segments = vec![ + (segment_digest(&w[..segment_len as usize]), 0..segment_len), + ( + segment_digest(&w[segment_len as usize..segment_len as usize * 2]), + segment_len..segment_len * 2, + ), + ( + segment_digest(&w[segment_len as usize * 2..short_len as usize]), + segment_len * 2..short_len, + ), + (segment_digest(&tail), short_len..total), + ]; + let grown_mirror = ExactStatePayloadMirror::FullState { + layout: L2Layout { + payload_kind: ExactStatePayloadKind::FullState, + total_bytes: total, + kv_bytes: total, + recurrent_bytes: 0, + segments: grown_segments, + }, + }; + let grown = key("ns", &[2]); + let evictions = tier + .admit( + grown.clone(), + 4, + grown_digest.clone(), + &grown_wire, + grown_mirror, + L2Origin::FromL3, + ) + .expect("admission must succeed by evicting the old entry"); + assert_eq!(evictions.len(), 1, "old entry is the victim"); + assert_eq!(evictions[0].cache_key, old); + assert_eq!( + evictions[0].freed_bytes, segment_len, + "only the victim's exclusive tail frees; the pinned prefix stays" + ); + assert_eq!( + evictions[0].retained_bytes, short_len, + "the shared prefix is retained by the incoming entry" + ); + + // The shared prefix segments must have survived the eviction. + let hit = tier.get(&grown).expect("grown entry hits"); + let payload = hit.to_payload(); + let (bytes, _) = payload.full_state_bytes_timed().expect("bytes"); + assert_eq!( + bytes.as_ref(), + &grown_wire[..], + "shared segments must survive the admission that evicted their old owner" + ); + let stats = tier.stats(); + assert_eq!(stats.entries, 1); + assert_eq!(stats.segments, 4); + assert_eq!( + stats.bytes, total, + "the pool is exactly the admitted entry's distinct bytes" + ); + assert!( + stats.bytes <= budget_from(&tier), + "the hard byte cap holds after a sharing admission" + ); + assert!(tier.peek(&old).is_none(), "old entry was evicted"); + } + + /// The budget the tier was built with (test-only mirror of the + /// constructor argument). + fn budget_from(tier: &L2Tier) -> u64 { + tier.stats().budget_bytes + } + + #[test] + fn admission_reusing_protected_bytes_cannot_exceed_the_budget() { + // 100-byte budget: X (60 bytes) is live, the incoming X+Y (120 + // bytes) shares X's segment. `new_bytes` is only Y's 60, X is + // pinned (shared), so eviction cannot free the excess: the + // admission must be refused and rolled back, never inserted at + // 120 bytes over a 100-byte budget. + let tier = L2Tier::new(100); + let k1 = key("ns", &[1]); + let k2 = key("ns", &[2]); + let (x, _) = wire(60, 1); + tier.admit( + k1.clone(), + 1, + segment_digest(&x), + &x, + single_segment_mirror(&x), + L2Origin::FromL3, + ) + .expect("X admitted"); + let (y, _) = wire(60, 2); + let mut xy = x.clone(); + xy.extend_from_slice(&y); + let xy_digest = segment_digest(&xy); + let mirror = ExactStatePayloadMirror::FullState { + layout: L2Layout { + payload_kind: ExactStatePayloadKind::FullState, + total_bytes: 120, + kv_bytes: 120, + recurrent_bytes: 0, + segments: vec![(segment_digest(&x), 0..60), (segment_digest(&y), 60..120)], + }, + }; + let err = tier + .admit(k2.clone(), 2, xy_digest, &xy, mirror, L2Origin::Direct) + .expect_err("protected overcommit must be refused"); + assert_eq!( + err, + L2InsertRefusal::ProtectedOvercommit { + pool_bytes: 120, + budget_bytes: 100, + }, + "the pool could only reach the budget by evicting the pinned X" + ); + // The tier is exactly as it was before the refused admission. + let stats = tier.stats(); + assert_eq!(stats.entries, 1); + assert_eq!(stats.bytes, 60); + assert_eq!(stats.segments, 1); + assert_eq!(stats.inserts, 1, "the refusal must not count an insert"); + assert_eq!( + stats.evictions, 0, + "rolled-back evictions must not be counted" + ); + assert!(tier.peek(&k1).is_some(), "X survived the refusal"); + assert!(tier.peek(&k2).is_none(), "the incoming entry was refused"); + // X still serves its exact bytes. + let hit = tier.get(&k1).expect("X hit"); + let payload = hit.to_payload(); + let (bytes, _) = payload.full_state_bytes_timed().expect("bytes"); + assert_eq!(bytes.as_ref(), &x[..]); + } + + #[test] + fn eviction_moves_shared_bytes_into_the_survivors_charge() { + // The short entry owns A,B; the long entry shares A,B and owns C. + // The long entry is charged only C. When the short entry is + // evicted, the pooled bytes do not change — so the survivor's + // charge must grow to the full pool, never drift below it. + let segment_len = 16u64; + let total = segment_len * 3; + let tier = L2Tier::new(1 << 20); + let short = key("ns", &[1]); + let long = key("ns", &[2]); + let (w, digest) = wire(total as usize, 70); + let short_len = segment_len * 2; + tier.admit( + short.clone(), + 2, + segment_digest(&w[..short_len as usize]), + &w[..short_len as usize], + manifest_shaped_mirror(&w[..short_len as usize], segment_len), + L2Origin::FromL3, + ) + .expect("short admitted"); + tier.admit( + long.clone(), + 3, + digest, + &w, + manifest_shaped_mirror(&w, segment_len), + L2Origin::FromL3, + ) + .expect("long admitted"); + assert_eq!( + tier.peek(&long).expect("long peeked").distinct_bytes, + total, + "before the eviction the long entry owns the shared pool outright \ + (lowest-key owner), so its charge covers the full physical pool" + ); + let removed = tier.remove(&short).expect("short present"); + assert_eq!( + removed.freed_bytes, 0, + "every short-entry segment stays in the pool under the long entry" + ); + assert_eq!( + removed.retained_bytes, short_len, + "retained bytes come from actual pool references" + ); + let stats = tier.stats(); + assert_eq!(stats.bytes, total); + assert_eq!( + tier.peek(&long).expect("long peeked").distinct_bytes, + total, + "the survivor's charge must cover the physical pool it now owns" + ); + // Aggregate invariant: the sum of all charges equals the pool. + let sum_of_charges = tier.peek(&long).expect("long").distinct_bytes; + assert_eq!(sum_of_charges, stats.bytes); + } + + #[test] + fn repeated_same_digest_segments_report_retained_bytes_from_the_pool() { + // One 16-byte segment laid out twice: the pool holds 16 bytes, the + // logical wire is 32. Removing the only entry frees 16 and retains + // nothing — `logical − freed` would have overstated retention. + let tier = L2Tier::new(1 << 20); + let k = key("ns", &[3]); + let (seg, _) = wire(16, 80); + let wire_bytes: Vec = [seg.clone(), seg.clone()].concat(); + let digest = segment_digest(&wire_bytes); + let mirror = ExactStatePayloadMirror::FullState { + layout: L2Layout { + payload_kind: ExactStatePayloadKind::FullState, + total_bytes: 32, + kv_bytes: 32, + recurrent_bytes: 0, + segments: vec![ + (segment_digest(&seg), 0..16), + (segment_digest(&seg), 16..32), + ], + }, + }; + tier.admit(k.clone(), 2, digest, &wire_bytes, mirror, L2Origin::FromL3) + .expect("admitted"); + let stats = tier.stats(); + assert_eq!(stats.bytes, 16, "the pool holds the segment once"); + assert_eq!( + tier.peek(&k).expect("peeked").distinct_bytes, + 16, + "the charge is the distinct pool bytes, not the logical wire" + ); + let removed = tier.remove(&k).expect("present"); + assert_eq!(removed.freed_bytes, 16); + assert_eq!(removed.retained_bytes, 0, "an emptied pool retains nothing"); + assert_eq!(tier.stats().bytes, 0); + assert_eq!(tier.stats().segments, 0); + } + + #[test] + fn removing_a_zero_charge_sharer_never_inflates_survivor_charges() { + // A, B, C all share segment X; only the lowest-key entry is ever + // charged for X. Removing the other sharers — whatever their key + // order — must leave the survivor's charge at exactly X's size, + // never 2x or 3x the physical pool. + let tier = L2Tier::new(1 << 20); + let ka = key("ns", &[1]); + let kb = key("ns", &[2]); + let kc = key("ns", &[3]); + let (x, dx) = wire(64, 90); + for k in [&ka, &kb, &kc] { + tier.admit( + k.clone(), + 1, + dx.clone(), + &x, + single_segment_mirror(&x), + L2Origin::FromL3, + ) + .expect("admitted"); + } + // After each removal, no surviving charge may exceed the pool. + let assert_charges_bounded = |tier: &L2Tier| { + let stats = tier.stats(); + for k in [ka.clone(), kb.clone(), kc.clone()] { + if let Some(peeked) = tier.peek(&k) { + assert!( + peeked.distinct_bytes <= stats.bytes, + "no charge may exceed the physical pool" + ); + } + } + stats + }; + // Remove the two zero-charge sharers in both orders. + tier.remove(&kb).expect("B removed"); + let stats = assert_charges_bounded(&tier); + assert_eq!(stats.bytes, 64); + tier.remove(&kc).expect("C removed"); + let stats = assert_charges_bounded(&tier); + assert_eq!(stats.bytes, 64); + // The remaining A owns X exactly once. + assert_eq!( + tier.peek(&ka).expect("A peeked").distinct_bytes, + 64, + "A's charge is X once, never the double- or triple-counted sum" + ); + assert_eq!(tier.stats().bytes, 64); + } + + #[test] + fn two_sharers_charge_moves_deterministically_to_the_lowest_key() { + let tier = L2Tier::new(1 << 20); + let ka = key("ns", &[10]); + let kb = key("ns", &[11]); + let (x, dx) = wire(48, 91); + for k in [&ka, &kb] { + tier.admit( + k.clone(), + 1, + dx.clone(), + &x, + single_segment_mirror(&x), + L2Origin::FromL3, + ) + .expect("admitted"); + } + // Exactly one of the two is charged (the lowest key), the other + // carries zero. + let charged_a = tier.peek(&ka).expect("A").distinct_bytes; + let charged_b = tier.peek(&kb).expect("B").distinct_bytes; + assert_eq!( + charged_a + charged_b, + 48, + "the sum of charges equals the physical pool" + ); + assert!(charged_a == 48 || charged_b == 48, "one owns X"); + assert!(charged_a == 0 || charged_b == 0, "the other pays nothing"); + // Remove whichever one that is not the owner: the charge sum is + // unchanged. + let (owner, zero) = if charged_a == 48 { + (ka.clone(), kb.clone()) + } else { + (kb.clone(), ka.clone()) + }; + tier.remove(&zero).expect("zero-charge sharer removed"); + assert_eq!(tier.stats().bytes, 48); + assert_eq!( + tier.peek(&owner).expect("owner peeked").distinct_bytes, + 48, + "the owner's charge is unchanged by a zero-charge removal" + ); + } + + #[test] + fn repeated_digest_with_survivor_reports_retained_once() { + // The removed entry references X twice in its layout; a survivor + // references X once. The pool retains X exactly once, so + // `retained_bytes` must be size(X) — never the double count the + // raw layout iteration would produce. + let tier = L2Tier::new(1 << 20); + let k_removed = key("ns", &[20]); + let k_survivor = key("ns", &[21]); + let (x, _) = wire(32, 95); + let x_digest = segment_digest(&x); + + // Survivor: single-segment mirror over X. + tier.admit( + k_survivor.clone(), + 1, + x_digest.clone(), + &x, + single_segment_mirror(&x), + L2Origin::FromL3, + ) + .expect("survivor admitted"); + + // Removed entry: X laid out twice (X X). + let wire_bytes: Vec = [x.clone(), x.clone()].concat(); + let digest = segment_digest(&wire_bytes); + let mirror = ExactStatePayloadMirror::FullState { + layout: L2Layout { + payload_kind: ExactStatePayloadKind::FullState, + total_bytes: 64, + kv_bytes: 64, + recurrent_bytes: 0, + segments: vec![(x_digest.clone(), 0..32), (x_digest, 32..64)], + }, + }; + tier.admit( + k_removed.clone(), + 2, + digest, + &wire_bytes, + mirror, + L2Origin::FromL3, + ) + .expect("removed entry admitted"); + let stats = tier.stats(); + assert_eq!(stats.bytes, 32, "the pool holds X once"); + + let removed = tier.remove(&k_removed).expect("removed entry present"); + assert_eq!( + removed.retained_bytes, 32, + "retained is size(X) once, not the twice-referenced 64" + ); + assert_eq!(removed.freed_bytes, 0, "the survivor keeps X"); + assert_eq!(tier.stats().bytes, 32); + assert_eq!(tier.stats().segments, 1); + let hit = tier.get(&k_survivor).expect("survivor intact"); + let payload = hit.to_payload(); + let (bytes, _) = payload.full_state_bytes_timed().expect("bytes"); + assert_eq!(bytes.as_ref(), &x[..]); + } + + #[test] + fn same_digest_different_bytes_is_rejected_unless_replacing_same_key() { + // A second wire claiming an existing segment digest with + // different-length content must not steal or replace the live + // pool handle. + let tier = L2Tier::new(1 << 20); + let (w1, _) = wire(32, 41); + let k1 = key("ns", &[1]); + tier.admit( + k1.clone(), + 2, + segment_digest(&w1), + &w1, + single_segment_mirror(&w1), + L2Origin::FromL3, + ) + .expect("first entry admitted"); + + // Forge a fake digest; the wire integrity check would reject a + // mismatched whole-wire digest, so claim the real segment digest + // of another wire as the *layout segment* digest instead. Build a + // second wire whose layout claims k1's segment digest. + let (w2, d2) = wire(48, 42); + let stolen = segment_digest(&w1); + let mirror = ExactStatePayloadMirror::FullState { + layout: L2Layout { + payload_kind: ExactStatePayloadKind::FullState, + total_bytes: 48, + kv_bytes: 48, + recurrent_bytes: 0, + segments: vec![(stolen, 0..24), (segment_digest(&w2[24..]), 24..48)], + }, + }; + let k2 = key("ns", &[2]); + let err = tier + .admit(k2, 3, d2, &w2, mirror, L2Origin::Direct) + .expect_err( + "a layout that claims another entry's digest with different bytes \ + must be refused", + ); + assert!( + matches!(err, L2InsertRefusal::SegmentDigestMismatch { .. }), + "expected segment digest mismatch, got: {err:?}" + ); + // The first entry's handle is untouched and still serves its bytes. + let hit = tier.get(&k1).expect("first entry intact"); + let payload = hit.to_payload(); + let (bytes, _) = payload.full_state_bytes_timed().expect("bytes"); + assert_eq!(bytes.as_ref(), &w1[..]); + } + + #[test] + fn get_with_missing_segment_handle_is_a_cold_miss_without_recency() { + let tier = L2Tier::new(1 << 20); + let k = key("ns", &[5]); + let (w, d) = wire(64, 51); + tier.admit( + k.clone(), + 4, + d, + &w, + manifest_shaped_mirror(&w, 16), + L2Origin::FromL3, + ) + .expect("admitted"); + + // Simulate corruption: drop one handle directly out of the pool + // (test-only access through the public remove on a scratch entry + // would also free it, but here we remove the pool entry via the + // tier's own release path by admitting an exclusive same-digest + // layout is impossible — so exercise via a second tier sharing + // nothing is not needed; directly verify the downgrade path by + // clearing the pool). + tier.clear_pool_for_test(); + + let hit = tier.get(&k); + assert!(hit.is_none(), "missing handles must downgrade to a miss"); + let stats = tier.stats(); + assert_eq!(stats.misses, 1, "the miss counter must move"); + assert_eq!(stats.hits, 0); + // The corrupt entry is removed: the next get is also a miss, not a + // partial hit, and no panic occurs. + assert!(tier.get(&k).is_none()); + assert_eq!(tier.stats().misses, 2); + assert!( + tier.peek(&k).is_none(), + "corrupt entry must be dropped, not left peekable" + ); + } + + #[test] + fn from_manifest_rejects_wrong_segment_index_and_offset() { + let (w, digest) = wire(32, 61); + let base = |index: u32, offset: u64, digest: String| HandoffManifest { + version: MANIFEST_VERSION, + codec: Some(PayloadCodec::raw()), + model_identity: "m".to_string(), + state_identity: "s".to_string(), + payload_kind: "full-state".to_string(), + total_bytes: 32, + payload_digest: digest.clone(), + segments: vec![HandoffSegmentRef { + index, + offset, + bytes: 32, + digest: segment_digest(&w), + codec_identity: Some(SegmentCodecIdentity::raw(32)), + meta_json: None, + }], + kv_bytes: 32, + recurrent_bytes: 0, + kv_desc_json: None, + token_count: 1, + continuation_token: 0, + expected_tokens: Vec::new(), + }; + let manifest = base(1, 0, digest.clone()); + let err = ExactStatePayloadMirror::from_manifest(&manifest) + .expect_err("wrong segment index must be refused"); + assert!(matches!(err, L2InsertRefusal::MalformedManifest(_))); + + let manifest = base(0, 8, digest); + let err = ExactStatePayloadMirror::from_manifest(&manifest) + .expect_err("wrong segment offset must be refused"); + assert!(matches!(err, L2InsertRefusal::MalformedManifest(_))); + } +} diff --git a/crates/skippy-cache/src/l3.rs b/crates/skippy-cache/src/l3.rs new file mode 100644 index 0000000000..353a5fa35f --- /dev/null +++ b/crates/skippy-cache/src/l3.rs @@ -0,0 +1,2215 @@ +//! L3 exact-state segment store. +//! +//! The durable tier under the radix cache: exported continuation state is cut +//! into content-addressed segments and committed under a manifest that records +//! ordering and completeness. Disk (this module) and the network handoff +//! stream are backends of the same contract: +//! +//! - **Segment identity**: every segment is addressed by the BLAKE3 digest of +//! its bytes; reads verify the digest, so corruption is detected, never +//! silently imported. +//! - **Ordering**: the manifest lists segments with explicit index/offset; +//! assembly validates both. +//! - **Completeness**: a manifest only commits after every referenced segment +//! is present and the assembled payload digest matches. Partial state can +//! never be loaded — there is nothing to load until commit. +//! - **Idempotency**: putting a segment that already exists is a no-op; +//! concurrent writers of the same bytes converge on one object via +//! temp-file + atomic rename. +//! - **Capped budget**: `enforce_budget` evicts oldest manifests first (the +//! newest is never evicted) and garbage-collects unreferenced segments. + +use std::{ + fs, + io::Write, + path::{Path, PathBuf}, + sync::{ + Mutex, RwLock, + atomic::{AtomicBool, AtomicU64, Ordering}, + }, +}; + +use anyhow::{Context, Result, bail}; +use serde::{Deserialize, Serialize}; + +use crate::fsinfo; + +mod packed; +use packed::{PACK_DIR, PACK_INDEX_DIR, PackedReadRequest, PackedSegmentStore}; + +const SEGMENT_DIR: &str = "segments"; +const MANIFEST_DIR: &str = "manifests"; +const PREFIX_INDEX_DIR: &str = "prefixes"; +const QUARANTINE_DIR: &str = "quarantine"; +const ROOT_LOCK_FILE: &str = ".owner.lock"; + +/// Evict to this percentage of the budget rather than exactly to it. +/// +/// An eviction pass is O(manifests x segments): it parses every manifest to +/// learn which segments would become unreferenced. Measured at 812 ms for 20 +/// manifests of ~9.5k segments, which is what a 19K-token prefix costs at a +/// 64-row window. Evicting exactly to the cap means a full cache pays that on +/// every commit; leaving headroom amortises it over the writes that fill the +/// headroom back up. +const EVICTION_LOW_WATER_PERCENT: u64 = 85; +/// On-disk format version stamped into every manifest. A released change to +/// the layout bumps this and makes older entries misses, never migrations. +/// +/// Version 4 requires explicit per-segment codec identity; a manifest at this +/// version whose segment lacks it is rejected rather than reinterpreted. +pub const MANIFEST_VERSION: u32 = 4; + +/// The pre-codec-identity manifest format. These entries predate codec +/// identity, are always raw by construction, and are decoded as raw through +/// the explicit legacy path in [`decode_manifest`]. New manifests are never +/// written at this version. +pub const LEGACY_MANIFEST_VERSION: u32 = 2; + +/// The payload-level-codec manifest format (#1750): codec identity is carried +/// once on the payload and every segment is implicitly that codec. Read for +/// compatibility, never written. New manifests stamp identity per segment. +pub const LEGACY_PAYLOAD_CODEC_MANIFEST_VERSION: u32 = 3; + +/// Identity of the payload envelope: raw, uncompressed exact state. Its bytes +/// are the segments verbatim, including segments whose representation is the +/// runtime's native KV page layout. +pub const CODEC_RAW: &str = "raw"; +/// Version of the raw codec's on-disk representation. Bumped only if the raw +/// byte layout itself changes; an unknown version is rejected, never migrated. +pub const CODEC_RAW_VERSION: u32 = 1; + +/// Exact KV bytes exported by the active native runtime and restored directly +/// into that runtime without a storage transcode. +pub const CODEC_NATIVE_KV_PAGE: &str = "native-kv-page"; +/// Version of the native KV segment contract. The concrete runtime ABI, +/// tensor types, geometry, and platform are bound by the exact-state identity +/// and the serialized runtime page descriptor carried by the manifest. +pub const CODEC_NATIVE_KV_PAGE_VERSION: u32 = 1; + +/// The codec used to encode a payload's segment bytes, stamped into the +/// manifest so representations are explicit and negotiable. +/// +/// Only [`CODEC_RAW`] is implemented today. The identity is recorded so future +/// lossless or compressed codecs are namespaced rather than guessed, and so an +/// older build refuses a payload it cannot decode instead of returning corrupt +/// state. Backward compatibility is handled by the manifest version, not a +/// serde default: a current-version manifest must stamp its codec explicitly +/// (so the field cannot be stripped to force a raw reinterpretation), while a +/// [`LEGACY_MANIFEST_VERSION`] manifest predates codec identity and is +/// normalized to raw on decode. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct PayloadCodec { + /// Codec name, e.g. `raw`. + pub name: String, + /// On-disk representation version within `name`. + pub version: u32, +} + +impl PayloadCodec { + /// The raw, uncompressed codec every current payload uses. + pub fn raw() -> Self { + Self { + name: CODEC_RAW.to_string(), + version: CODEC_RAW_VERSION, + } + } + + /// Whether this build can assemble a payload encoded with this codec. + /// Only the exact raw name and version are supported; any other name or a + /// future raw version is unknown and must be refused before assembly. + pub fn is_supported(&self) -> bool { + self.name == CODEC_RAW && self.version == CODEC_RAW_VERSION + } +} + +impl Default for PayloadCodec { + fn default() -> Self { + Self::raw() + } +} + +/// What a codec's output promises at assembly time. The class is contract, +/// not documentation: an exact entry can satisfy an exact lookup; a lossy +/// entry can only satisfy a lookup that accepts its codec family and +/// calibration. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CodecClass { + /// Decodes to the exact segment bytes the writer held. Payload digest + /// verification applies at assembly. + Exact, + /// Decodes to a calibrated approximation of the segment bytes. Never + /// verified against a payload digest; must never satisfy an exact + /// lookup. + Lossy, +} + +impl CodecClass { + pub fn as_str(self) -> &'static str { + match self { + Self::Exact => "exact", + Self::Lossy => "lossy", + } + } +} + +/// Per-segment codec identity: which codec, which representation version, +/// what it promises, and (for lossy codecs) which calibration produced it. +/// Required on every segment of a current-version manifest. +/// +/// The identity is namespaced per segment rather than per payload so one +/// payload can mix raw and compressed segments (e.g. a grown prefix whose +/// leading windows are compressed and whose new tail is raw), and so a v4 +/// manifest with the identity stripped from a segment is detectable and +/// rejected instead of silently inheriting the payload codec or raw. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct SegmentCodecIdentity { + /// Segment codec name, e.g. `raw`. Distinct namespace from the payload + /// codec so entry formats version independently. + pub name: String, + /// On-disk representation version within `name`. + pub version: u32, + /// What decoding promises: [`CodecClass::Exact`] entries verify against + /// the payload digest; [`CodecClass::Lossy`] entries never do. + pub class: CodecClass, + /// Byte length the codec decodes to. Must equal the assembled payload + /// region this segment covers; a mismatch is corruption, not a hint. + pub decoded_len: u64, + /// Calibration namespace for lossy codecs (e.g. the CacheGen CDF + /// calibration digest). Must be empty for exact codecs; lossy lookups + /// only ever match a payload calibrated with the same digest. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub calibration_digest: Option, +} + +impl SegmentCodecIdentity { + /// Identity for a segment stored verbatim: the exact raw codec at the + /// store's raw version, decoding to `segment.bytes`. + pub fn raw(encoded_len: u64) -> Self { + Self { + name: CODEC_RAW.to_string(), + version: CODEC_RAW_VERSION, + class: CodecClass::Exact, + decoded_len: encoded_len, + calibration_digest: None, + } + } + + /// Identity for a runtime-native KV page segment stored verbatim. The + /// store performs no encode/decode step; restore imports these same bytes + /// through the runtime page descriptor associated with the manifest. + pub fn native_kv_page(encoded_len: u64) -> Self { + Self { + name: CODEC_NATIVE_KV_PAGE.to_string(), + version: CODEC_NATIVE_KV_PAGE_VERSION, + class: CodecClass::Exact, + decoded_len: encoded_len, + calibration_digest: None, + } + } + + /// Whether this identity names the supported native KV passthrough + /// representation. This is stricter than a name check: an identity with + /// the wrong version, class, or calibration is not native passthrough. + pub fn is_native_kv_page(&self) -> bool { + self.name == CODEC_NATIVE_KV_PAGE + && self.version == CODEC_NATIVE_KV_PAGE_VERSION + && self.class == CodecClass::Exact + && self.calibration_digest.is_none() + } + + /// Whether this build can assemble a segment encoded with this identity. + /// Exact raw and native KV passthrough are implemented. Both assemble + /// verbatim; anything else is unknown and must be refused before assembly. + pub fn is_supported(&self) -> bool { + let supported_representation = (self.name == CODEC_RAW + && self.version == CODEC_RAW_VERSION) + || self.is_native_kv_page(); + supported_representation + && self.class == CodecClass::Exact + && self.calibration_digest.is_none() + } + + /// Internal consistency the capability negotiation relies on: an exact + /// identity must be raw-shaped (decoded length equals the stored bytes, + /// no calibration), and every identity must agree that decoding yields + /// what the manifest says it yields. + pub fn is_self_consistent(&self, encoded_len: u64) -> bool { + match self.class { + CodecClass::Exact => { + self.decoded_len == encoded_len && self.calibration_digest.is_none() + } + // A lossy identity must declare a calibration to decode against; + // its decoded length is checked against the payload layout at + // negotiation, not against the stored bytes. + CodecClass::Lossy => self.calibration_digest.is_some(), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct HandoffSegmentRef { + pub index: u32, + pub offset: u64, + pub bytes: u64, + pub digest: String, + /// Per-segment codec identity: what the stored bytes are, what decoding + /// them promises, and (for lossy codecs) which calibration applies. + /// Required on a current-version manifest; the `Option` exists only so a + /// v4 manifest with the identity stripped is *detectable* and rejected — + /// it never falls back to the payload codec or to raw. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub codec_identity: Option, + /// Per-segment metadata for page-stream payloads (serialized + /// `RuntimeKvPageDesc` plus token range), opaque to the store. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub meta_json: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HandoffManifest { + pub version: u32, + /// Numerical model identity used by exact model-scoped lifecycle + /// operations. This intentionally spans every split stage of one model. + pub model_identity: String, + /// `exact_state_identity` of the producing runtime — the numerical + /// identity a loader must match before importing this state. + pub state_identity: String, + pub payload_kind: String, + /// Codec identity of the segment payload bytes. Always present on a decoded + /// manifest: a current-version manifest must carry it explicitly, and a + /// [`LEGACY_MANIFEST_VERSION`] manifest is normalized to raw on decode. The + /// field is optional only so a version-3 manifest with the codec stripped + /// is detectable and rejected rather than silently defaulted to raw. + #[serde(default)] + pub codec: Option, + pub total_bytes: u64, + /// BLAKE3 of the assembled payload; also the manifest's key. + pub payload_digest: String, + pub segments: Vec, + pub kv_bytes: u64, + pub recurrent_bytes: u64, + /// Serialized `RuntimeKvPageDesc` for kv-recurrent payloads; opaque to + /// this crate so the store does not depend on the runtime. + #[serde(skip_serializing_if = "Option::is_none")] + pub kv_desc_json: Option, + pub token_count: u64, + pub continuation_token: i32, + /// Greedy continuation produced by the exporting session, when known — + /// lets an offline restore self-verify determinism. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub expected_tokens: Vec, +} + +impl HandoffManifest { + pub fn new(state_identity: String, payload_kind: String) -> Self { + Self::new_for_model(state_identity.clone(), state_identity, payload_kind) + } + + pub fn new_for_model( + model_identity: String, + state_identity: String, + payload_kind: String, + ) -> Self { + Self { + version: MANIFEST_VERSION, + model_identity, + state_identity, + payload_kind, + codec: Some(PayloadCodec::raw()), + total_bytes: 0, + payload_digest: String::new(), + segments: Vec::new(), + kv_bytes: 0, + recurrent_bytes: 0, + kv_desc_json: None, + token_count: 0, + continuation_token: 0, + expected_tokens: Vec::new(), + } + } + + /// Whether this manifest contains runtime-native KV page segments. + pub fn uses_native_kv_passthrough(&self) -> bool { + self.segments.iter().any(|segment| { + segment + .codec_identity + .as_ref() + .is_some_and(SegmentCodecIdentity::is_native_kv_page) + }) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SegmentPut { + pub new: bool, + pub bytes: u64, +} + +/// How a payload is laid out, so segments can be cut where a growing prefix +/// keeps its bytes still. +/// +/// The runtime exports exact state as a sequence of runs, each holding one +/// token-row per token in token order: every layer's K, then every layer's V, +/// then the indexer rows. Adding tokens extends every run, so a cut at a fixed +/// byte offset lands in a different place each turn and nothing dedupes — the +/// measured cost was 8x the newly committed bytes at 8 MiB segments. +/// +/// Cutting each run into fixed windows of token-rows instead means turn N+1's +/// segments are byte-identical to turn N's up to the last partial window, and +/// only genuinely new state is written. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PayloadGeometry { + /// Runs in wire order. + pub blocks: Vec, + /// Token-rows in every run. + pub rows: u64, + /// Rows per segment. Must depend only on the model's own shape, never on + /// how many tokens this particular entry holds, or the boundaries move + /// between turns and the dedupe is lost. + pub window_rows: u64, + /// Trailing bytes with no row structure (a recurrent snapshot), cut at the + /// store's default segment size. + pub tail_bytes: u64, +} + +/// One run of token-rows. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct GeometryBlock { + /// Bytes per token-row. + pub stride: u64, + /// What this run holds, for the segment metadata: `k`, `v` or `kidx`. + pub kind: GeometryKind, + /// Which layer, relative to the exported range. + pub layer: u32, + /// Sub-run within the layer. Always 0 except for transposed V, where each + /// embedding column is its own token-contiguous run. + pub column: u32, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum GeometryKind { + Key, + Value, + KeyIndex, +} + +impl GeometryKind { + pub fn as_str(self) -> &'static str { + match self { + Self::Key => "k", + Self::Value => "v", + Self::KeyIndex => "kidx", + } + } +} + +impl PayloadGeometry { + /// Bytes the geometry claims, for checking it against the payload it is + /// meant to describe. + pub fn total_bytes(&self) -> u64 { + self.blocks + .iter() + .map(|block| block.stride.saturating_mul(self.rows)) + .fold(0u64, u64::saturating_add) + .saturating_add(self.tail_bytes) + } + + /// A geometry that does not describe this payload exactly is not usable: + /// cutting to it would produce segments that reassemble to different + /// bytes. Callers fall back to fixed-size cutting. + pub fn matches(&self, payload_bytes: u64) -> bool { + self.rows > 0 + && self.window_rows > 0 + && !self.blocks.is_empty() + && self.blocks.iter().all(|block| block.stride > 0) + && self.total_bytes() == payload_bytes + } + + /// The cuts, as `(offset, len, label)` in wire order. + pub fn plan(&self, tail_segment_bytes: u64) -> Vec<(u64, u64, String)> { + let mut cuts = Vec::new(); + let mut offset = 0u64; + for block in &self.blocks { + let mut row = 0u64; + while row < self.rows { + let rows = self.window_rows.min(self.rows - row); + let len = rows.saturating_mul(block.stride); + cuts.push(( + offset, + len, + format!( + "{}:{}:{}:{row}", + block.kind.as_str(), + block.layer, + block.column + ), + )); + offset = offset.saturating_add(len); + row += rows; + } + } + let tail_cut = tail_segment_bytes.max(1); + let mut remaining = self.tail_bytes; + while remaining > 0 { + let len = tail_cut.min(remaining); + cuts.push((offset, len, "tail".to_string())); + offset = offset.saturating_add(len); + remaining -= len; + } + cuts + } +} + +/// What the store is allowed to occupy. Both bounds are hard: the budget caps +/// what the cache manages, the reserve caps what it may take from everything +/// else on the filesystem. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct StoreLimits { + /// Whole-root cap over segments, manifests, indexes and in-flight bytes. + /// Zero disables the cap; the public configuration surface rejects zero + /// rather than treating it as unlimited, so only internal callers and + /// tests can reach that state. + pub budget_bytes: u64, + /// Free space the store preserves for everything else on the filesystem. + pub minimum_free_bytes: u64, +} + +impl StoreLimits { + pub fn new(budget_bytes: u64, minimum_free_bytes: u64) -> Self { + Self { + budget_bytes, + minimum_free_bytes, + } + } +} + +/// Why a write was not admitted. Every refusal is a miss with a stable reason, +/// never a partial or best-effort write. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WriteRefusal { + /// The entry alone is larger than the whole budget. Storing it could only + /// succeed by evicting everything else and still overflowing. + SkippedOversize, + /// The budget is full and eviction cannot free enough (everything else is + /// pinned or the entry needs more than the whole cap). + InsufficientSpace, + /// The filesystem is at the minimum-free reserve. Reads still serve; the + /// tier stops writing until space comes back. + ReadOnlyLowSpace, +} + +impl WriteRefusal { + /// Stable reason code for the status surface. + pub fn reason(self) -> &'static str { + match self { + Self::SkippedOversize => "skipped_oversize", + Self::InsufficientSpace => "insufficient_space", + Self::ReadOnlyLowSpace => "read_only_low_space", + } + } +} + +/// Capacity held for bytes that are being written but not yet committed. +/// Counted against the budget for as long as it lives, so two concurrent +/// writers cannot each pass an admission check and together overflow the cap. +#[derive(Debug)] +pub struct Reservation<'store> { + store: &'store HandoffSegmentStore, + bytes: u64, +} + +impl Drop for Reservation<'_> { + fn drop(&mut self) { + self.store + .reserved_inflight + .fetch_sub(self.bytes, Ordering::AcqRel); + } +} + +/// Holds one segment against collection while its manifest is being built. +/// +/// A writer puts every segment before committing the manifest that binds them, +/// so for that window the bytes are unreferenced and eviction would collect +/// them out from under the commit about to name them. The hold lasts exactly +/// as long as the writer keeps the guard, so an abandoned write releases on +/// drop rather than leaking until restart. +#[derive(Debug)] +pub struct SegmentHold<'store> { + store: &'store HandoffSegmentStore, + digest: String, +} + +impl Drop for SegmentHold<'_> { + fn drop(&mut self) { + let mut holds = self + .store + .inflight_segments + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(count) = holds.get_mut(&self.digest) { + *count = count.saturating_sub(1); + if *count == 0 { + holds.remove(&self.digest); + } + } + } +} + +/// A stored segment and the hold that keeps it collectable-proof until the +/// caller's manifest commits. +#[derive(Debug)] +pub struct StoredSegment<'store> { + pub digest: String, + pub put: SegmentPut, + _hold: SegmentHold<'store>, +} + +/// Holds one manifest against eviction while it is being read or written. +#[derive(Debug)] +pub struct ManifestPin<'store> { + store: &'store HandoffSegmentStore, + key: String, +} + +impl Drop for ManifestPin<'_> { + fn drop(&mut self) { + let mut pins = self + .store + .pins + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(count) = pins.get_mut(&self.key) { + *count = count.saturating_sub(1); + if *count == 0 { + pins.remove(&self.key); + } + } + } +} + +/// What the store currently holds, for the status contract. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +pub struct StoreUsage { + pub budget_bytes: u64, + pub used_bytes: u64, + pub reserved_inflight_bytes: u64, + pub filesystem_available_bytes: u64, + pub minimum_free_bytes: u64, + pub manifests: u64, + pub unique_segments: u64, + pub evicted_manifests: u64, + pub quarantined_objects: u64, +} + +/// Repairs performed before a node starts serving a cache root. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize)] +pub struct StoreReconciliation { + pub removed_temporary_files: u64, + pub quarantined_manifests: u64, + pub removed_prefix_links: u64, + pub removed_orphan_bytes: u64, +} + +#[derive(Debug)] +pub struct HandoffSegmentStore { + root: PathBuf, + limits: RwLock, + /// Exclusive process-level ownership of this cache root. The manager + /// shares one store between all stages in the node; a second process must + /// not build an independent reservation/pin universe over the same files. + _root_lock: fs::File, + /// Makes the admission check and reservation increment one transaction. + /// Without this, two stages can both observe the same free budget before + /// either publishes its reservation. + admission: Mutex<()>, + /// Bytes reserved by in-flight writes. Part of the managed total, so the + /// cap holds across concurrent writers rather than only at rest. + reserved_inflight: AtomicU64, + /// Manifests an active reader or writer is using. A pinned manifest is + /// never evicted, pruned or cleared out from under the operation. + pins: Mutex>, + /// Manifests removed by budget enforcement or prune since open. + evicted_manifests: AtomicU64, + /// Objects moved to quarantine after failing verification since open. + quarantined_objects: AtomicU64, + /// Last computed on-disk managed total, excluding in-flight reservations. + /// + /// `reserve` runs per segment put, and a full scan stats every file under + /// the root, so recomputing it each time makes one spill of a long prefix + /// O(segments^2) syscalls under the admission lock. The hot path keeps + /// this total incrementally; every bulk mutation invalidates it and the + /// next read pays for one authoritative scan. + usage_bytes: AtomicU64, + /// Whether `usage_bytes` can be trusted without rescanning. + usage_valid: AtomicBool, + /// Segments written but not yet referenced by a committed manifest. + /// + /// A writer puts every segment before committing the manifest that binds + /// them, so for that window the bytes are unreferenced — and eviction + /// triggered by another writer (or by this one needing room) would + /// collect them out from under the commit that is about to reference + /// them. Left unprotected this fails as "manifest references missing + /// segment" under exactly the pressure the cache is for. + inflight_segments: Mutex>, + packed: PackedSegmentStore, +} + +pub fn segment_digest(bytes: &[u8]) -> String { + blake3::hash(bytes).to_hex().to_string() +} + +fn acquire_root_lock(root: &Path) -> Result { + let path = root.join(ROOT_LOCK_FILE); + let file = fs::OpenOptions::new() + .create(true) + .truncate(false) + .read(true) + .write(true) + .open(&path) + .with_context(|| format!("failed to open cache root lock {}", path.display()))?; + fsinfo::restrict_to_owner(&path, 0o600)?; + + if let Err(error) = fs2::FileExt::try_lock_exclusive(&file) { + if error.kind() == std::io::ErrorKind::WouldBlock { + bail!( + "cache root {} is already owned by another manager", + root.display() + ); + } + return Err(error).with_context(|| format!("failed to lock cache root {}", root.display())); + } + + Ok(file) +} + +impl HandoffSegmentStore { + /// Open (creating if needed) a store rooted at `root`, capped by + /// `budget_bytes` with no free-space reserve. Prefer + /// [`Self::open_with_limits`]; this exists for callers that predate the + /// reserve. + pub fn open(root: impl Into, budget_bytes: u64) -> Result { + Self::open_with_limits(root, StoreLimits::new(budget_bytes, 0)) + } + + /// Open (creating if needed) a store rooted at `root` under `limits`. + /// + /// Refuses a root reached through a symlink and a root on a network + /// filesystem: both break the atomic-rename and containment assumptions + /// every later guarantee rests on, and neither is worth a partial mode. + pub fn open_with_limits(root: impl Into, limits: StoreLimits) -> Result { + let store = Self::open_unreconciled_with_limits(root, limits)?; + store.reconcile_startup()?; + Ok(store) + } + + pub(crate) fn open_unreconciled_with_limits( + root: impl Into, + limits: StoreLimits, + ) -> Result { + let root = root.into(); + if !root.is_absolute() { + bail!("cache root must be absolute: {}", root.display()); + } + fsinfo::create_dir_all_without_links(&root)?; + // Anchor everything below to the resolved root, so a symlink crossed on + // the way in cannot make containment checks disagree with where bytes + // actually land. + let root = fs::canonicalize(&root) + .with_context(|| format!("failed to resolve cache root {}", root.display()))?; + if fsinfo::is_network_filesystem(&root).unwrap_or(false) { + let name = fsinfo::filesystem_type_name(&root).unwrap_or_default(); + bail!( + "{} is on an unsupported network filesystem ({name}); the disk cache requires local storage", + root.display() + ); + } + for directory in [ + SEGMENT_DIR, + MANIFEST_DIR, + PREFIX_INDEX_DIR, + PACK_DIR, + PACK_INDEX_DIR, + ] { + let path = root.join(directory); + fsinfo::refuse_symlinked_descendant(&root, &path)?; + fs::create_dir_all(&path).with_context(|| { + format!("failed to create {directory} dir under {}", root.display()) + })?; + fsinfo::restrict_to_owner(&path, 0o700)?; + } + fsinfo::restrict_to_owner(&root, 0o700)?; + let root_lock = acquire_root_lock(&root)?; + let packed = PackedSegmentStore::open(&root)?; + Ok(Self { + root, + limits: RwLock::new(limits), + _root_lock: root_lock, + admission: Mutex::new(()), + reserved_inflight: AtomicU64::new(0), + // Left invalid so the first read pays for one authoritative scan + // of whatever the previous process left behind. + usage_bytes: AtomicU64::new(0), + usage_valid: AtomicBool::new(false), + pins: Mutex::new(std::collections::BTreeMap::new()), + evicted_manifests: AtomicU64::new(0), + quarantined_objects: AtomicU64::new(0), + inflight_segments: Mutex::new(std::collections::BTreeMap::new()), + packed, + }) + } + + pub fn root(&self) -> &Path { + &self.root + } + + pub fn limits(&self) -> StoreLimits { + *self + .limits + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + /// Replace the live hard cap and filesystem reserve. Admission is paused + /// while the pair changes, so a writer can never observe half of the new + /// policy. Shrinking evicts inactive entries immediately; pinned entries + /// remain valid and subsequent writes stay refused until usage fits. + pub fn update_limits(&self, limits: StoreLimits) -> Result { + let _admission = self + .admission + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut current = self + .limits + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let previous = *current; + *current = limits; + drop(current); + if limits.budget_bytes > 0 { + self.enforce_budget_to(limits.budget_bytes)?; + } + Ok(previous) + } + + /// Repair incomplete or invalid state left by an interrupted writer before + /// the manager exposes this root to any stage. + pub fn reconcile_startup(&self) -> Result { + let mut report = StoreReconciliation::default(); + for directory in [ + SEGMENT_DIR, + MANIFEST_DIR, + PREFIX_INDEX_DIR, + PACK_DIR, + PACK_INDEX_DIR, + ] { + report.removed_temporary_files = report + .removed_temporary_files + .saturating_add(remove_temporary_files(&self.root.join(directory))?); + } + + for key in self.list_manifests()? { + if self.validate_committed_manifest(&key).is_err() { + self.quarantine(&self.manifest_path(&key))?; + report.quarantined_manifests += 1; + } + } + self.rebuild_packed_index()?; + report.removed_prefix_links = self.remove_dangling_prefix_links()?; + report.removed_orphan_bytes = self.collect_unreferenced_segments()?; + Ok(report) + } + + fn validate_committed_manifest(&self, key: &str) -> Result<()> { + let manifest = self.load_manifest(key)?; + if manifest.payload_digest != key { + bail!( + "manifest key {key} disagrees with payload digest {}", + manifest.payload_digest + ); + } + let mut expected_offset = 0u64; + let locations = self + .packed + .load_manifest_index(key, manifest.segments.len())?; + for (position, (segment, location)) in manifest.segments.iter().zip(locations).enumerate() { + if segment.index as usize != position || segment.offset != expected_offset { + bail!("manifest {key} has invalid segment ordering"); + } + self.validate_segment_ref(segment, location.as_ref()) + .with_context(|| format!("manifest {key} references a missing segment"))?; + expected_offset = expected_offset + .checked_add(segment.bytes) + .context("manifest segment offsets overflow")?; + } + if expected_offset != manifest.total_bytes { + bail!("manifest {key} does not tile its payload"); + } + Ok(()) + } + + fn remove_dangling_prefix_links(&self) -> Result { + // Files are about to be removed or rewritten in bulk. + self.invalidate_usage(); + let root = self.root.join(PREFIX_INDEX_DIR); + let mut removed = 0u64; + for path in files_recursive(&root)? { + let digest = match fs::read_to_string(&path) { + Ok(value) => value, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => continue, + Err(error) => return Err(error.into()), + }; + if digest.is_empty() || !self.manifest_path(&digest).is_file() { + match fs::remove_file(&path) { + Ok(()) => removed += 1, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => return Err(error.into()), + } + } + } + Ok(removed) + } + + fn segment_path(&self, digest: &str) -> PathBuf { + self.root.join(SEGMENT_DIR).join(format!("{digest}.seg")) + } + + fn validate_segment_ref( + &self, + segment: &HandoffSegmentRef, + location: Option<&packed::PackedSegmentLocation>, + ) -> Result<()> { + if let Some(location) = location { + return self.packed.validate(location, segment.bytes); + } + let metadata = fs::metadata(self.segment_path(&segment.digest))?; + if metadata.len() != segment.bytes { + bail!("segment {} is truncated", segment.digest); + } + Ok(()) + } + + fn rebuild_packed_index(&self) -> Result<()> { + let mut entries = Vec::new(); + for key in self.list_manifests()? { + let Ok(manifest) = self.load_manifest(&key) else { + continue; + }; + let locations = self + .packed + .load_manifest_index(&key, manifest.segments.len())?; + entries.extend(manifest.segments.into_iter().zip(locations).filter_map( + |(segment, location)| location.map(|location| (segment.digest, location)), + )); + } + self.packed.rebuild(entries); + Ok(()) + } + + fn manifest_path(&self, payload_digest: &str) -> PathBuf { + self.root + .join(MANIFEST_DIR) + .join(format!("{payload_digest}.json")) + } + + fn namespace_dir(&self, namespace_key: &str) -> PathBuf { + let key = namespace_key + .strip_prefix("blake3:") + .unwrap_or(namespace_key); + self.root.join(PREFIX_INDEX_DIR).join(key) + } + + fn prefix_entry_path(&self, namespace_key: &str, token_len: u64, prefix_key: &str) -> PathBuf { + // Zero-padded length keeps directory listings sorted and lets the + // lookup filter by length without parsing every name. + let key = prefix_key.strip_prefix("blake3:").unwrap_or(prefix_key); + self.namespace_dir(namespace_key) + .join(format!("{token_len:012}-{key}.key")) + } + + /// Bind a (namespace, token-length, prefix-hash) coordinate to a + /// committed manifest. Entries at many lengths coexist, which is what + /// makes longest-recorded-prefix lookup work: each spill is a complete + /// state for its own length, and later, longer prompts find the longest + /// spilled length that is a prefix of theirs. + pub fn link_prefix( + &self, + namespace_key: &str, + token_len: u64, + prefix_key: &str, + payload_digest: &str, + ) -> Result<()> { + let path = self.prefix_entry_path(namespace_key, token_len, prefix_key); + let bytes = payload_digest.as_bytes(); + let replaced_bytes = + fs::metadata(&path) + .map(|metadata| metadata.len()) + .or_else(|error| { + if error.kind() == std::io::ErrorKind::NotFound { + Ok(0) + } else { + Err(error) + } + })?; + let growth_bytes = (bytes.len() as u64).saturating_sub(replaced_bytes); + let reservation = match self.reserve_write(bytes.len() as u64, growth_bytes)? { + Ok(reservation) => reservation, + Err(refusal) => bail!("cannot store prefix link: {}", refusal.reason()), + }; + let parent = path.parent().context("prefix entry has no parent")?; + fsinfo::create_dir_all_without_links(parent)?; + fsinfo::restrict_to_owner(parent, 0o700)?; + // Replacement can shrink or grow the entry, so rescan rather than + // trying to update the cached total from a racy pre-write stat. + self.invalidate_usage(); + write_atomically(&path, bytes)?; + fsinfo::restrict_to_owner(&path, 0o600)?; + drop(reservation); + self.enforce_budget()?; + Ok(()) + } + + /// Recorded token lengths for a namespace, longest first, deduplicated. + pub fn recorded_prefix_lengths(&self, namespace_key: &str) -> Result> { + let dir = self.namespace_dir(namespace_key); + let entries = match fs::read_dir(&dir) { + Ok(entries) => entries, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(error) => return Err(error.into()), + }; + let mut lengths = Vec::new(); + for entry in entries { + let entry = entry?; + let name = entry.file_name(); + let Some(name) = name.to_str() else { continue }; + let Some((length, _)) = name.split_once('-') else { + continue; + }; + if let Ok(length) = length.parse::() { + lengths.push(length); + } + } + lengths.sort_unstable_by(|a, b| b.cmp(a)); + lengths.dedup(); + Ok(lengths) + } + + /// The manifest recorded at exactly (namespace, token-length, + /// prefix-hash), pruning links whose manifest was evicted. Transient + /// I/O errors are surfaced, not treated as absence, so a briefly + /// unreadable disk cannot delete healthy links. + pub fn manifest_for_prefix( + &self, + namespace_key: &str, + token_len: u64, + prefix_key: &str, + ) -> Result> { + let path = self.prefix_entry_path(namespace_key, token_len, prefix_key); + let bytes = match fs::read(&path) { + Ok(bytes) => bytes, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(error) => return Err(error.into()), + }; + let payload_digest = String::from_utf8(bytes).context("malformed prefix link")?; + let manifest_path = self.manifest_path(&payload_digest); + match fs::read(&manifest_path) { + Ok(bytes) => { + let manifest = match decode_manifest(&payload_digest, &bytes) { + Ok(manifest) => manifest, + Err(_) => { + self.quarantine(&manifest_path)?; + let _ = fs::remove_file(&path); + self.invalidate_usage(); + return Ok(None); + } + }; + // A hit is a use. Recording it here is what makes eviction + // least-recently-*used* rather than least-recently-written. + self.touch_manifest(&payload_digest); + Ok(Some(manifest)) + } + Err(error) if error.kind() == std::io::ErrorKind::NotFound => { + // The manifest was evicted after the link was written; only + // this definite absence prunes the link. + let _ = fs::remove_file(&path); + self.invalidate_usage(); + Ok(None) + } + Err(error) => Err(error.into()), + } + } + + /// Content-addressed, idempotent put. Concurrent writers of the same + /// bytes race benignly: both write temp files, both rename onto the same + /// final path. + /// + /// Stores one segment, or reports why it was refused. + /// + /// A segment already present is a no-op that costs nothing and is always + /// admitted: content addressing means the bytes on disk are the bytes + /// being written, so re-putting cannot grow the store. + pub fn put_segment(&self, bytes: &[u8]) -> Result> { + match self.try_put_segment(bytes)? { + Ok(stored) => Ok(stored), + Err(refusal) => bail!("cannot store segment: {}", refusal.reason()), + } + } + + /// As [`Self::put_segment`], returning the refusal rather than an error so + /// callers can record a miss with its reason and carry on. + pub fn try_put_segment(&self, bytes: &[u8]) -> Result, WriteRefusal>> { + let digest = segment_digest(bytes); + // Establish the GC hold before observing or publishing the file. This + // closes the rename-to-hold window where a concurrent clear could see + // an unreferenced segment and remove it before its manifest commits. + let hold = self.hold_segment(&digest); + let path = self.segment_path(&digest); + if path.exists() { + return Ok(Ok(StoredSegment { + digest, + put: SegmentPut { + new: false, + bytes: bytes.len() as u64, + }, + _hold: hold, + })); + } + // Reserve before any temporary bytes exist on disk, so two writers + // cannot both pass the check and together exceed the cap. + let reservation = match self.reserve(bytes.len() as u64)? { + Ok(reservation) => reservation, + Err(refusal) => return Ok(Err(refusal)), + }; + if let Err(error) = + write_atomically(&path, bytes).and_then(|()| fsinfo::restrict_to_owner(&path, 0o600)) + { + // A failed cleanup or a permissions failure may leave physical + // bytes behind. Force the next admission to reconcile with disk. + self.invalidate_usage(); + return Err(error); + } + self.add_usage_bytes(bytes.len() as u64); + drop(reservation); + Ok(Ok(StoredSegment { + digest, + put: SegmentPut { + new: true, + bytes: bytes.len() as u64, + }, + _hold: hold, + })) + } + + /// Store one spill's logical segments in one immutable physical pack. + pub fn try_put_segments<'store>( + &'store self, + segments: &[&[u8]], + ) -> Result>, WriteRefusal>> { + let digests = segments + .iter() + .map(|bytes| segment_digest(bytes)) + .collect::>(); + let holds = digests + .iter() + .map(|digest| self.hold_segment(digest)) + .collect::>(); + let inputs = digests + .iter() + .zip(segments) + .map(|(digest, bytes)| (digest.as_str(), *bytes)) + .collect::>(); + let estimated = self.packed.estimated_new_bytes(&inputs); + let reservation = match self.reserve(estimated)? { + Ok(reservation) => reservation, + Err(refusal) => return Ok(Err(refusal)), + }; + let (packed, new_bytes) = match self.packed.write_batch(&inputs) { + Ok(result) => result, + Err(error) => { + self.invalidate_usage(); + return Err(error); + } + }; + self.add_usage_bytes(new_bytes); + drop(reservation); + Ok(Ok(digests + .into_iter() + .zip(holds) + .zip(packed) + .zip(segments) + .map(|(((digest, hold), packed), bytes)| StoredSegment { + digest, + put: SegmentPut { + new: packed.new, + bytes: bytes.len() as u64, + }, + _hold: hold, + }) + .collect())) + } + + pub fn has_segment(&self, digest: &str) -> bool { + self.segment_path(digest).exists() || self.packed.location(digest).is_some() + } + + /// Read one segment, verifying its content digest. + pub fn read_segment(&self, digest: &str) -> Result> { + if let Some(location) = self.packed.location(digest) { + let requests = [PackedReadRequest { + digest, + bytes: location.bytes, + location: &location, + output_offset: 0, + }]; + let len = usize::try_from(location.bytes).context("segment exceeds usize")?; + let mut bytes = Vec::with_capacity(len); + return match self.packed.append_many(&requests, &mut bytes, None) { + Ok(()) => Ok(bytes), + Err(failure) => { + let path = self.packed.pack_path(&failure.pack_digest); + let _ = self.quarantine(&path); + Err(failure.error) + } + }; + } + let path = self.segment_path(digest); + let bytes = fs::read(&path).with_context(|| format!("failed to read segment {digest}"))?; + if segment_digest(&bytes) != digest { + self.quarantine(&path)?; + bail!("segment {digest} failed digest verification on read and was quarantined"); + } + Ok(bytes) + } + + /// Move a corrupt or truncated object out of the managed tree. + /// + /// A segment that fails verification can never become valid: every + /// manifest referencing it is already a miss, and leaving it in place + /// means paying to read and reject it again on the next hit. Quarantine + /// keeps the bytes under `quarantine/` for inspection while taking them + /// out of every lookup path. + fn quarantine(&self, path: &Path) -> Result<()> { + // Files are about to be removed or rewritten in bulk. + self.invalidate_usage(); + let directory = self.root.join(QUARANTINE_DIR); + fs::create_dir_all(&directory).with_context(|| { + format!( + "failed to create quarantine dir under {}", + self.root.display() + ) + })?; + fsinfo::restrict_to_owner(&directory, 0o700)?; + let name = path + .file_name() + .context("quarantined path has no file name")?; + self.quarantined_objects.fetch_add(1, Ordering::Relaxed); + match fs::rename(path, directory.join(name)) { + Ok(()) => Ok(()), + // Losing the evidence beats serving corrupt state, so a failed + // move falls back to removal. + Err(_) => fs::remove_file(path) + .with_context(|| format!("failed to quarantine {}", path.display())), + } + } + + /// Commit a manifest. Fails unless every referenced segment is present + /// with the recorded size and offsets tile the payload exactly — the + /// completeness gate that makes partial state unloadable. + pub fn commit(&self, manifest: &HandoffManifest) -> Result<()> { + match self.try_commit(manifest)? { + Ok(()) => Ok(()), + Err(refusal) => bail!( + "cannot commit manifest {}: {}", + manifest.payload_digest, + refusal.reason() + ), + } + } + + /// As [`Self::commit`], preserving an admission refusal as structured + /// state for the manager's effective-state contract. + pub fn try_commit(&self, manifest: &HandoffManifest) -> Result> { + // The manifest and its prefix index are new files on disk. + self.invalidate_usage(); + if manifest.payload_digest.is_empty() { + bail!("manifest has no payload digest"); + } + validate_manifest_compatibility(manifest)?; + let mut expected_offset = 0u64; + let locations = manifest + .segments + .iter() + .map(|segment| self.packed.location(&segment.digest)) + .collect::>(); + for (position, (segment, location)) in manifest.segments.iter().zip(&locations).enumerate() + { + if segment.index as usize != position { + bail!( + "manifest segment order broken: index {} at position {position}", + segment.index + ); + } + if segment.offset != expected_offset { + bail!( + "manifest segment {} offset {} does not tile payload (expected {expected_offset})", + segment.index, + segment.offset + ); + } + self.validate_segment_ref(segment, location.as_ref()) + .with_context(|| { + format!( + "manifest references missing segment {} ({})", + segment.index, segment.digest + ) + })?; + expected_offset = expected_offset + .checked_add(segment.bytes) + .context("manifest offsets overflow")?; + } + if expected_offset != manifest.total_bytes { + bail!( + "manifest segments cover {expected_offset} bytes but total_bytes is {}", + manifest.total_bytes + ); + } + let limits = self.limits(); + if limits.budget_bytes > 0 && manifest.total_bytes > limits.budget_bytes { + return Ok(Err(WriteRefusal::SkippedOversize)); + } + // Compact, not pretty: a 64-row window turns a long prefix into + // thousands of segment refs, and eviction parses every manifest to + // build its reference map. Indentation is pure cost on a file nobody + // reads by hand. + let serialized = serde_json::to_vec(manifest).context("failed to serialize manifest")?; + let segment_digests = manifest + .segments + .iter() + .map(|segment| segment.digest.clone()) + .collect::>(); + let packed_index = self + .packed + .encode_manifest_index(&manifest.payload_digest, &segment_digests)?; + // The manifest is what makes the segments loadable, so it is pinned + // while it lands: eviction triggered by its own admission check must + // not remove the entry being committed. + let _pin = self.pin(&manifest.payload_digest); + let manifest_path = self.manifest_path(&manifest.payload_digest); + let replaced_bytes = fs::metadata(&manifest_path) + .map(|metadata| metadata.len()) + .or_else(|error| { + if error.kind() == std::io::ErrorKind::NotFound { + Ok(0) + } else { + Err(error) + } + })?; + let growth_bytes = (serialized.len() as u64).saturating_sub(replaced_bytes); + let packed_index_path = self.packed.index_path(&manifest.payload_digest); + let replaced_index_bytes = fs::metadata(&packed_index_path) + .map(|metadata| metadata.len()) + .unwrap_or(0); + let index_bytes = packed_index.as_ref().map_or(0, |bytes| bytes.len() as u64); + let write_bytes = (serialized.len() as u64).saturating_add(index_bytes); + let total_growth = + growth_bytes.saturating_add(index_bytes.saturating_sub(replaced_index_bytes)); + match self.reserve_write(write_bytes, total_growth)? { + Ok(_reservation) => { + if let Some(index) = &packed_index { + self.packed + .publish_manifest_index(&manifest.payload_digest, index)?; + } else { + self.packed + .remove_manifest_index(&manifest.payload_digest)?; + } + write_atomically(&manifest_path, &serialized)?; + fsinfo::restrict_to_owner(&manifest_path, 0o600)?; + } + Err(refusal) => return Ok(Err(refusal)), + } + self.enforce_budget()?; + Ok(Ok(())) + } + + /// Record that an entry was used, so eviction can order by last use. + /// Best-effort: a store on a read-only path still serves reads, and + /// failing a restore because recency could not be recorded would trade a + /// working cache for a bookkeeping detail. + pub fn touch_manifest(&self, payload_digest: &str) { + let _ = fsinfo::touch(&self.manifest_path(payload_digest)); + } + + pub fn load_manifest(&self, payload_digest: &str) -> Result { + let bytes = fs::read(self.manifest_path(payload_digest)) + .with_context(|| format!("failed to read manifest {payload_digest}"))?; + decode_manifest(payload_digest, &bytes) + } + + /// Manifest keys, newest first by modification time. + pub fn list_manifests(&self) -> Result> { + let mut entries = Vec::new(); + for entry in fs::read_dir(self.root.join(MANIFEST_DIR))? { + let entry = entry?; + let path = entry.path(); + if path.extension().is_none_or(|extension| extension != "json") { + continue; + } + let Some(stem) = path + .file_stem() + .map(|stem| stem.to_string_lossy().into_owned()) + else { + continue; + }; + let modified = entry.metadata()?.modified()?; + entries.push((modified, stem)); + } + entries.sort_by_key(|entry| std::cmp::Reverse(entry.0)); + Ok(entries.into_iter().map(|(_, stem)| stem).collect()) + } + + /// Assemble the full payload for a manifest, verifying every segment + /// digest, the tiling, and the whole-payload digest. + pub fn assemble(&self, manifest: &HandoffManifest) -> Result> { + validate_manifest_compatibility(manifest)?; + let total = usize::try_from(manifest.total_bytes).context("payload exceeds usize")?; + let mut payload = Vec::with_capacity(total); + let mut payload_hasher = blake3::Hasher::new(); + let locations = self + .packed + .load_manifest_index(&manifest.payload_digest, manifest.segments.len())?; + let mut expected_offset = 0u64; + for segment in &manifest.segments { + if segment.offset != expected_offset { + bail!( + "segment {} offset {} does not match assembled length {expected_offset}", + segment.index, + segment.offset, + ); + } + expected_offset = expected_offset + .checked_add(segment.bytes) + .context("assembled payload size overflows")?; + } + if expected_offset != manifest.total_bytes { + bail!( + "assembled {expected_offset} bytes but manifest records {}", + manifest.total_bytes + ); + } + let mut packed_requests = Vec::new(); + for (segment, location) in manifest.segments.iter().zip(&locations) { + if let Some(location) = location.as_ref() { + packed_requests.push(PackedReadRequest { + digest: &segment.digest, + bytes: segment.bytes, + location, + output_offset: segment.offset, + }); + } else { + self.append_packed(&packed_requests, &mut payload, &mut payload_hasher)?; + packed_requests.clear(); + let bytes = self.read_segment(&segment.digest)?; + payload_hasher.update(&bytes); + payload.extend_from_slice(&bytes); + } + } + self.append_packed(&packed_requests, &mut payload, &mut payload_hasher)?; + if payload.len() != total { + bail!( + "assembled {} bytes but manifest records {total}", + payload.len() + ); + } + if payload_hasher.finalize().to_hex().as_str() != manifest.payload_digest { + bail!("assembled payload failed manifest digest verification"); + } + Ok(payload) + } + + fn append_packed( + &self, + requests: &[PackedReadRequest<'_>], + payload: &mut Vec, + payload_hasher: &mut blake3::Hasher, + ) -> Result<()> { + match self + .packed + .append_many(requests, payload, Some(payload_hasher)) + { + Ok(()) => Ok(()), + Err(failure) => { + let path = self.packed.pack_path(&failure.pack_digest); + let _ = self.quarantine(&path); + Err(failure.error) + } + } + } + + pub fn segment_footprint_bytes(&self) -> Result { + Ok(directory_bytes(&self.root.join(SEGMENT_DIR))? + .saturating_add(self.packed.footprint_bytes()?)) + } + + /// Every byte the store manages: committed segments, the manifests that + /// make them loadable, the prefix index that finds them, and capacity + /// reserved by writes in flight. + /// + /// The budget is a physical cap on the cache root, so it has to count the + /// bookkeeping too. Counting segments alone lets a store with many small + /// entries sit well over its stated cap in manifests and index files. + pub fn managed_usage_bytes(&self) -> Result { + let on_disk = if self.usage_valid.load(Ordering::Acquire) { + self.usage_bytes.load(Ordering::Acquire) + } else { + self.rescan_usage_bytes()? + }; + Ok(on_disk.saturating_add(self.reserved_inflight.load(Ordering::Acquire))) + } + + /// Stat every managed file and adopt the result as the running total. + fn rescan_usage_bytes(&self) -> Result { + let mut total = directory_bytes(&self.root.join(SEGMENT_DIR))?; + total = total.saturating_add(self.packed.footprint_bytes()?); + total = total.saturating_add(self.packed.index_footprint_bytes()?); + total = total.saturating_add(directory_bytes(&self.root.join(MANIFEST_DIR))?); + total = total.saturating_add(directory_bytes_recursive( + &self.root.join(PREFIX_INDEX_DIR), + )?); + total = total.saturating_add(directory_bytes_recursive(&self.root.join(QUARANTINE_DIR))?); + self.usage_bytes.store(total, Ordering::Release); + self.usage_valid.store(true, Ordering::Release); + Ok(total) + } + + /// Drop the cached total. Every mutation that is not a plain segment put + /// goes through here, so the next read rescans rather than trusting a + /// figure that a bulk removal or rewrite may have invalidated. + fn invalidate_usage(&self) { + self.usage_valid.store(false, Ordering::Release); + } + + /// Account for bytes a segment put just added, keeping the hot path free + /// of a rescan. A no-op while the total is already known to be stale. + fn add_usage_bytes(&self, bytes: u64) { + if self.usage_valid.load(Ordering::Acquire) { + self.usage_bytes.fetch_add(bytes, Ordering::AcqRel); + } + } + + /// Reserve capacity for `bytes` about to be written. + /// + /// Reservation happens before any temporary payload bytes exist on disk, + /// so the cap and the free-space reserve hold even while several writers + /// are mid-write. The returned guard releases the reservation on drop, + /// including on the error paths. + pub fn reserve(&self, bytes: u64) -> Result, WriteRefusal>> { + self.reserve_write(bytes, bytes) + } + + /// Reserve free space for the temporary write and budget capacity only + /// for its projected net growth. Replacements need room for a complete + /// temporary file, but charging their full size against the cache cap + /// would evict healthy entries even when the final file has the same size. + fn reserve_write( + &self, + write_bytes: u64, + growth_bytes: u64, + ) -> Result, WriteRefusal>> { + let _admission = self + .admission + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let limits = self.limits(); + if limits.budget_bytes > 0 && growth_bytes > limits.budget_bytes { + return Ok(Err(WriteRefusal::SkippedOversize)); + } + let available = fsinfo::available_bytes(&self.root)?; + if available.saturating_sub(write_bytes) < limits.minimum_free_bytes { + return Ok(Err(WriteRefusal::ReadOnlyLowSpace)); + } + if limits.budget_bytes > 0 { + let used = self.managed_usage_bytes()?; + if used.saturating_add(growth_bytes) > limits.budget_bytes { + // Make room from inactive entries before refusing: a full + // cache is the normal steady state, not an error. Clear to the + // low-water mark, or further when this write alone needs more. + let target = self + .low_water_bytes() + .min(limits.budget_bytes.saturating_sub(growth_bytes)); + self.enforce_budget_to(target)?; + // Never refuse on a cached figure: rescan so an incremental + // total that has drifted low cannot turn a writable cache into + // a reported InsufficientSpace. + let used = self + .rescan_usage_bytes()? + .saturating_add(self.reserved_inflight.load(Ordering::Acquire)); + if used.saturating_add(growth_bytes) > limits.budget_bytes { + return Ok(Err(WriteRefusal::InsufficientSpace)); + } + } + } + self.reserved_inflight + .fetch_add(growth_bytes, Ordering::AcqRel); + Ok(Ok(Reservation { + store: self, + bytes: growth_bytes, + })) + } + + /// Pin a manifest for the duration of a read or write. Eviction, prune + /// and clear all skip pinned entries, so an in-progress restore never has + /// its segments removed underneath it. + pub fn pin(&self, payload_digest: &str) -> ManifestPin<'_> { + let mut pins = self + .pins + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + *pins.entry(payload_digest.to_string()).or_insert(0) += 1; + ManifestPin { + store: self, + key: payload_digest.to_string(), + } + } + + fn is_pinned(&self, payload_digest: &str) -> bool { + self.pins + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get(payload_digest) + .is_some_and(|count| *count > 0) + } + + /// What the store holds right now, for the status contract. + pub fn usage(&self) -> Result { + let limits = self.limits(); + let manifests = self.list_manifests()?; + let mut segments = std::collections::HashSet::new(); + for key in &manifests { + if let Ok(manifest) = self.load_manifest(key) { + for segment in manifest.segments { + segments.insert(segment.digest); + } + } + } + Ok(StoreUsage { + budget_bytes: limits.budget_bytes, + used_bytes: self.managed_usage_bytes()?, + reserved_inflight_bytes: self.reserved_inflight.load(Ordering::Acquire), + filesystem_available_bytes: fsinfo::available_bytes(&self.root)?, + minimum_free_bytes: limits.minimum_free_bytes, + manifests: manifests.len() as u64, + unique_segments: segments.len() as u64, + evicted_manifests: self.evicted_manifests.load(Ordering::Relaxed), + quarantined_objects: self.quarantined_objects.load(Ordering::Relaxed), + }) + } + + /// Namespaces with at least one indexed prefix. Each namespace is one + /// numerical identity (model, layout, layer range, load mode), so this is + /// how many distinct configurations the root currently serves. + pub fn namespace_count(&self) -> Result { + let entries = match fs::read_dir(self.root.join(PREFIX_INDEX_DIR)) { + Ok(entries) => entries, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(0), + Err(error) => return Err(error.into()), + }; + let mut count = 0u64; + for entry in entries { + if entry?.file_type()?.is_dir() { + count += 1; + } + } + Ok(count) + } + + /// Evict least-recently-used manifests and collect unreferenced segments + /// until managed usage fits the budget. Returns bytes freed. + /// + /// Recency is last *use*, not last write: reads touch their manifest, so a + /// prefix that is hit constantly but never re-recorded is the last thing + /// evicted rather than the first. + pub fn enforce_budget(&self) -> Result { + let limits = self.limits(); + if limits.budget_bytes == 0 { + return Ok(0); + } + // Over the cap, evict below it: the pass is expensive enough that + // paying it once per headroom refill beats paying it per commit. + if self.managed_usage_bytes()? <= limits.budget_bytes { + return Ok(0); + } + self.enforce_budget_to(self.low_water_bytes()) + } + + /// The level eviction drops to once it runs. + fn low_water_bytes(&self) -> u64 { + self.limits() + .budget_bytes + .saturating_mul(EVICTION_LOW_WATER_PERCENT) + / 100 + } + + /// Evict until managed usage is at or below `target_bytes`. + /// + /// Single pass regardless of how many manifests evict: usage, the manifest + /// list and the reference map are each scanned once, eviction runs against + /// the in-memory model, and one final GC removes what became unreferenced. + /// Pinned manifests are never evicted, so an in-flight read or write keeps + /// its state loadable even under pressure. + pub fn enforce_budget_to(&self, target_bytes: u64) -> Result { + self.enforce_budget_to_model(target_bytes, None) + } + + fn enforce_budget_to_model( + &self, + target_bytes: u64, + model_identity: Option<&str>, + ) -> Result { + let usage_before = self.managed_usage_bytes()?; + if usage_before <= target_bytes { + return Ok(0); + } + // Least-recently-used last, so eviction pops from the back. + let keys = self.list_manifests()?; + let mut manifests = Vec::with_capacity(keys.len()); + let mut reference_counts: std::collections::HashMap = + std::collections::HashMap::new(); + for key in &keys { + let manifest = self.load_manifest(key).with_context(|| { + format!("failed to build eviction reference map from manifest {key}") + })?; + let locations = self + .packed + .load_manifest_index(key, manifest.segments.len())?; + for (segment, location) in manifest.segments.iter().zip(&locations) { + let (object, bytes) = location.as_ref().map_or_else( + || (segment.digest.clone(), segment.bytes), + |location| { + ( + format!("pack:{}", location.pack_digest), + fs::metadata(self.packed.pack_path(&location.pack_digest)) + .map(|metadata| metadata.len()) + .unwrap_or(0), + ) + }, + ); + let entry = reference_counts.entry(object).or_insert((0, bytes)); + entry.0 += 1; + } + manifests.push((manifest, locations)); + } + let mut freeable = usage_before; + let mut evicted_any = false; + while freeable > target_bytes { + let Some(position) = manifests.iter().rposition(|(manifest, _)| { + !self.is_pinned(&manifest.payload_digest) + && model_identity.is_none_or(|identity| manifest.model_identity == identity) + }) else { + // Everything left is in use. Refusing the write is correct; + // tearing state out from under a live operation is not. + break; + }; + let (evicted, locations) = manifests.remove(position); + let manifest_path = self.manifest_path(&evicted.payload_digest); + let manifest_bytes = fs::metadata(&manifest_path) + .map(|metadata| metadata.len()) + .unwrap_or(0); + self.invalidate_usage(); + fs::remove_file(&manifest_path) + .with_context(|| format!("failed to evict manifest {}", evicted.payload_digest))?; + let index_bytes = self.packed.remove_manifest_index(&evicted.payload_digest)?; + self.evicted_manifests.fetch_add(1, Ordering::Relaxed); + freeable = freeable + .saturating_sub(manifest_bytes) + .saturating_sub(index_bytes); + for (segment, location) in evicted.segments.iter().zip(locations) { + let object = location.map_or_else( + || segment.digest.clone(), + |location| format!("pack:{}", location.pack_digest), + ); + if let Some(entry) = reference_counts.get_mut(&object) { + entry.0 = entry.0.saturating_sub(1); + if entry.0 == 0 { + freeable = freeable.saturating_sub(entry.1); + } + } + } + evicted_any = true; + } + if !evicted_any { + return Ok(0); + } + self.remove_dangling_prefix_links()?; + self.collect_unreferenced_segments()?; + let usage_after = self.managed_usage_bytes()?; + Ok(usage_before.saturating_sub(usage_after)) + } + + /// Evict least-recently-used inactive manifests until managed usage fits + /// `target_bytes`. The user-facing prune: it never removes an active + /// entry, and it reports what it actually freed rather than what it + /// intended to. + pub fn prune_to(&self, target_bytes: u64) -> Result { + self.enforce_budget_to(target_bytes) + } + + pub(crate) fn prune_model_to(&self, model_identity: &str, target_bytes: u64) -> Result { + self.enforce_budget_to_model(target_bytes, Some(model_identity)) + } + + /// Remove every manifest that is not pinned, then collect the segments + /// that became unreferenced. Returns bytes freed. + /// + /// Requests in flight keep serving: clearing removes stored state, and a + /// miss falls back to cold prefill. + pub fn clear(&self) -> Result { + self.clear_model_inner(None) + } + + pub(crate) fn clear_model(&self, model_identity: &str) -> Result { + self.clear_model_inner(Some(model_identity)) + } + + fn clear_model_inner(&self, model_identity: Option<&str>) -> Result { + // Files are about to be removed or rewritten in bulk. + self.invalidate_usage(); + let before = self.managed_usage_bytes()?; + for key in self.list_manifests()? { + if self.is_pinned(&key) { + continue; + } + if let Some(model_identity) = model_identity { + let manifest = self.load_manifest(&key).with_context(|| { + format!("failed to inspect manifest {key} while clearing model state") + })?; + if manifest.model_identity != model_identity { + continue; + } + } + let path = self.manifest_path(&key); + match fs::remove_file(&path) { + Ok(()) => {} + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => { + return Err(error).with_context(|| format!("failed to clear manifest {key}")); + } + } + self.packed.remove_manifest_index(&key)?; + } + self.remove_dangling_prefix_links()?; + self.collect_unreferenced_segments()?; + let after = self.managed_usage_bytes()?; + Ok(before.saturating_sub(after)) + } + + fn hold_segment(&self, digest: &str) -> SegmentHold<'_> { + let mut holds = self + .inflight_segments + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + *holds.entry(digest.to_string()).or_insert(0) += 1; + SegmentHold { + store: self, + digest: digest.to_string(), + } + } + + /// Remove segments referenced by no manifest. Returns bytes freed. + pub fn collect_unreferenced_segments(&self) -> Result { + // Files are about to be removed or rewritten in bulk. + self.invalidate_usage(); + let mut referenced = std::collections::HashSet::new(); + let manifest_keys = self.list_manifests()?; + for key in &manifest_keys { + let manifest = self + .load_manifest(key) + .with_context(|| format!("failed to build GC reference map from manifest {key}"))?; + for segment in manifest.segments { + referenced.insert(segment.digest); + } + } + let mut freed = 0u64; + let held = self + .inflight_segments + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .keys() + .cloned() + .collect::>(); + for entry in fs::read_dir(self.root.join(SEGMENT_DIR))? { + let entry = entry?; + let path = entry.path(); + if path.extension().is_none_or(|extension| extension != "seg") { + continue; + } + let Some(stem) = path + .file_stem() + .map(|stem| stem.to_string_lossy().into_owned()) + else { + continue; + }; + if !referenced.contains(&stem) && !held.contains(&stem) { + freed = freed.saturating_add(entry.metadata()?.len()); + fs::remove_file(&path) + .with_context(|| format!("failed to collect segment {stem}"))?; + } + } + let manifests = manifest_keys + .into_iter() + .collect::>(); + freed = freed.saturating_add(self.packed.remove_orphan_indexes(&manifests)?); + freed = freed.saturating_add(self.packed.remove_orphan_packs(&referenced, || { + self.inflight_segments + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .keys() + .cloned() + .collect() + })?); + Ok(freed) + } +} + +/// Refuse a manifest whose payload codec this build cannot decode, before any +/// segment is read or reassembled. Keeps unknown codecs from being committed +/// (nothing unassemblable is ever persisted) and from being assembled (an +/// unknown codec is a miss, never a silent misinterpretation of the bytes). A +/// missing identity is treated as unsupported, not defaulted to raw. +fn reject_unsupported_codec(manifest: &HandoffManifest) -> Result<()> { + match manifest.codec.as_ref() { + Some(codec) if codec.is_supported() => Ok(()), + Some(codec) => bail!( + "manifest {} uses unsupported codec {}/{}; this build assembles only {}/{}", + manifest.payload_digest, + codec.name, + codec.version, + CODEC_RAW, + CODEC_RAW_VERSION + ), + None => bail!( + "manifest {} has no codec identity; this build requires an explicit codec", + manifest.payload_digest + ), + } +} + +/// Versions this build can read and write: the current per-segment format +/// and the two legacy formats (pre-codec and payload-level-codec). +fn manifest_version_is_supported(version: u32) -> bool { + version == MANIFEST_VERSION + || version == LEGACY_MANIFEST_VERSION + || version == LEGACY_PAYLOAD_CODEC_MANIFEST_VERSION +} + +/// Reject a manifest whose segment codec identities this build cannot decode, +/// before any segment is read or reassembled. Also enforces the per-segment +/// negotiation contract: every identity must be internally consistent with +/// the stored bytes it covers, and a mixed manifest must name the payload +/// codec that describes the assembled whole. Positional errors name the +/// segment index so the offending ref is identifiable in the message. +fn reject_unsupported_segment_codecs(manifest: &HandoffManifest) -> Result<()> { + let has_native_kv = manifest.uses_native_kv_passthrough(); + let has_runtime_descriptor = manifest + .kv_desc_json + .as_deref() + .is_some_and(|descriptor| serde_json::from_str::(descriptor).is_ok()); + if has_native_kv + && (manifest.version != MANIFEST_VERSION + || manifest.payload_kind != "kv-recurrent" + || manifest.kv_bytes == 0 + || !has_runtime_descriptor + || manifest.kv_bytes.checked_add(manifest.recurrent_bytes) + != Some(manifest.total_bytes)) + { + bail!( + "manifest {} names native KV segments without a current, complete kv-recurrent payload and runtime page descriptor", + manifest.payload_digest + ); + } + for segment in &manifest.segments { + let Some(identity) = segment.codec_identity.as_ref() else { + // Legacy formats carry no per-segment identity by construction; + // the payload gate covers what their segments decode as. + continue; + }; + if !identity.is_supported() { + bail!( + "manifest {} segment {} uses unsupported codec {}/{}; this build assembles only {}/{} and {}/{}", + manifest.payload_digest, + segment.index, + identity.name, + identity.version, + CODEC_RAW, + CODEC_RAW_VERSION, + CODEC_NATIVE_KV_PAGE, + CODEC_NATIVE_KV_PAGE_VERSION + ); + } + if !identity.is_self_consistent(segment.bytes) { + bail!( + "manifest {} segment {} codec identity is inconsistent with its stored bytes ({}/{} class, decoded_len {}, encoded_len {})", + manifest.payload_digest, + segment.index, + identity.name, + identity.version, + identity.decoded_len, + segment.bytes + ); + } + if has_native_kv { + let end = segment + .offset + .checked_add(segment.bytes) + .context("segment range overflows")?; + let covers_kv = segment.offset < manifest.kv_bytes; + if covers_kv && end > manifest.kv_bytes { + bail!( + "manifest {} segment {} crosses the native KV boundary at byte {}", + manifest.payload_digest, + segment.index, + manifest.kv_bytes + ); + } + if covers_kv != identity.is_native_kv_page() { + bail!( + "manifest {} segment {} representation disagrees with the native KV boundary at byte {}", + manifest.payload_digest, + segment.index, + manifest.kv_bytes + ); + } + } + // The payload codec describes the assembled envelope. Exact native KV + // segments remain byte-for-byte members of the raw envelope; a future + // codec that transforms stored bytes must upgrade that envelope rather + // than leave it claiming raw. + } + Ok(()) +} + +/// The current format requires explicit per-segment codec identity. Shared +/// by `decode_manifest` (on-disk reads) and +/// `validate_manifest_compatibility` (commit and assembly), so neither path +/// can persist or partially read a v4 manifest whose segment identity was +/// stripped — there is no fallback to the payload codec or to raw. +fn require_v4_segment_identities(manifest: &HandoffManifest) -> Result<()> { + for segment in &manifest.segments { + if segment.codec_identity.is_none() { + bail!( + "manifest {digest} at version {MANIFEST_VERSION} is missing per-segment codec identity on segment {index}", + digest = manifest.payload_digest, + index = segment.index + ); + } + } + Ok(()) +} + +/// Validate that an in-memory manifest is one this build can both persist and +/// assemble: a supported format version *and* supported codecs at both the +/// payload and every segment. `try_commit` and `assemble` call this before +/// writing or reading any segment, so a manifest that `decode_manifest` would +/// later reject — an unknown version (e.g. a future version carrying a raw +/// codec) or an unsupported codec at either level — cannot be committed to +/// disk or partially assembled from segments. +fn validate_manifest_compatibility(manifest: &HandoffManifest) -> Result<()> { + if !manifest_version_is_supported(manifest.version) { + bail!( + "manifest {} has version {} but this build reads {MANIFEST_VERSION} or legacy {LEGACY_MANIFEST_VERSION}/{LEGACY_PAYLOAD_CODEC_MANIFEST_VERSION}", + manifest.payload_digest, + manifest.version + ); + } + if manifest.version == MANIFEST_VERSION { + require_v4_segment_identities(manifest)?; + } + reject_unsupported_codec(manifest)?; + reject_unsupported_segment_codecs(manifest) +} + +fn decode_manifest(payload_digest: &str, bytes: &[u8]) -> Result { + let mut manifest: HandoffManifest = + serde_json::from_slice(bytes).context("malformed manifest")?; + match manifest.version { + MANIFEST_VERSION => { + // The current format must stamp both the payload codec and every + // per-segment codec identity explicitly. A missing field means it + // was stripped to force a reinterpretation of bytes this build + // may not decode — reject, never default to raw or to the payload + // codec. + if manifest.codec.is_none() { + bail!( + "manifest {payload_digest} at version {MANIFEST_VERSION} is missing its required codec identity" + ); + } + require_v4_segment_identities(&manifest)?; + } + LEGACY_PAYLOAD_CODEC_MANIFEST_VERSION => { + // #1750 format: codec identity is carried once at the payload + // level and every segment implicitly is that codec. The payload + // codec is required there (a stripped one is rejected, same as + // current), and no per-segment identity exists to normalize — + // the payload gate below covers what the segments decode as. + if manifest.codec.is_none() { + bail!( + "manifest {payload_digest} at version {LEGACY_PAYLOAD_CODEC_MANIFEST_VERSION} is missing its required codec identity" + ); + } + } + LEGACY_MANIFEST_VERSION => { + // Pre-codec-identity manifests are raw by construction. Accept an + // absent codec (normalize to raw) or an explicit raw codec; a + // legacy manifest cannot legitimately name a non-raw codec. + match manifest.codec { + None => manifest.codec = Some(PayloadCodec::raw()), + Some(ref codec) if *codec == PayloadCodec::raw() => {} + Some(ref codec) => bail!( + "legacy manifest {payload_digest} declares non-raw codec {}/{}", + codec.name, + codec.version + ), + } + } + other => bail!( + "manifest {payload_digest} has version {other} but this build reads {MANIFEST_VERSION} or legacy {LEGACY_MANIFEST_VERSION}/{LEGACY_PAYLOAD_CODEC_MANIFEST_VERSION}" + ), + } + // Central capability gate: no load entry point returns a manifest whose + // codecs this build cannot decode. Startup reconciliation quarantines it, + // manifest_for_prefix prunes the link and falls back to a shorter prefix, + // and a direct load fails. + reject_unsupported_codec(&manifest)?; + reject_unsupported_segment_codecs(&manifest)?; + if manifest.payload_digest != payload_digest { + bail!( + "manifest key {payload_digest} disagrees with payload digest {}", + manifest.payload_digest + ); + } + Ok(manifest) +} + +fn files_recursive(directory: &Path) -> Result> { + let entries = match fs::read_dir(directory) { + Ok(entries) => entries, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(error) => return Err(error.into()), + }; + let mut files = Vec::new(); + for entry in entries { + let entry = entry?; + let file_type = entry.file_type()?; + if file_type.is_file() { + files.push(entry.path()); + } else if file_type.is_dir() { + files.extend(files_recursive(&entry.path())?); + } + } + Ok(files) +} + +fn remove_temporary_files(directory: &Path) -> Result { + let mut removed = 0u64; + for path in files_recursive(directory)? { + let is_temporary = path + .file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| name.starts_with(".tmp-")); + if !is_temporary { + continue; + } + match fs::remove_file(&path) { + Ok(()) => removed += 1, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => { + return Err(error) + .with_context(|| format!("failed to remove stale temp {}", path.display())); + } + } + } + Ok(removed) +} + +/// Bytes held by the files directly inside `directory`. Missing directories +/// count as empty so a partially built root reports rather than fails. +fn directory_bytes(directory: &Path) -> Result { + let entries = match fs::read_dir(directory) { + Ok(entries) => entries, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(0), + Err(error) => { + return Err(error).with_context(|| format!("failed to read {}", directory.display())); + } + }; + let mut total = 0u64; + for entry in entries { + let entry = entry?; + if entry.file_type()?.is_file() { + total = total.saturating_add(entry.metadata()?.len()); + } + } + Ok(total) +} + +/// As [`directory_bytes`], following one level of namespace subdirectories: +/// the prefix index is stored per namespace. +fn directory_bytes_recursive(directory: &Path) -> Result { + let entries = match fs::read_dir(directory) { + Ok(entries) => entries, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(0), + Err(error) => { + return Err(error).with_context(|| format!("failed to read {}", directory.display())); + } + }; + let mut total = 0u64; + for entry in entries { + let entry = entry?; + let file_type = entry.file_type()?; + if file_type.is_file() { + total = total.saturating_add(entry.metadata()?.len()); + } else if file_type.is_dir() { + total = total.saturating_add(directory_bytes_recursive(&entry.path())?); + } + } + Ok(total) +} + +fn write_atomically(path: &Path, bytes: &[u8]) -> Result<()> { + let directory = path.parent().context("path has no parent directory")?; + let (temp_path, mut temp_file) = tempfile_in(directory)?; + let write_result = temp_file + .write_all(bytes) + .with_context(|| format!("failed to write {}", temp_path.display())) + .and_then(|()| { + temp_file + .sync_all() + .with_context(|| format!("failed to sync {}", temp_path.display())) + }); + drop(temp_file); + let publish_result = write_result.and_then(|()| fsinfo::replace_file(&temp_path, path)); + if let Err(error) = publish_result { + return match fs::remove_file(&temp_path) { + Ok(()) => Err(error), + Err(cleanup) if cleanup.kind() == std::io::ErrorKind::NotFound => Err(error), + Err(cleanup) => Err(error.context(format!( + "also failed to remove temporary file {}: {cleanup}", + temp_path.display() + ))), + }; + } + Ok(()) +} + +fn tempfile_in(directory: &Path) -> Result<(PathBuf, fs::File)> { + // Distinct per-writer temp names without a clock or RNG dependency: + // process id + a process-local counter. + use std::sync::atomic::{AtomicU64, Ordering}; + static COUNTER: AtomicU64 = AtomicU64::new(0); + let unique = COUNTER.fetch_add(1, Ordering::Relaxed); + let path = directory.join(format!(".tmp-{}-{unique}", std::process::id())); + let file = fs::File::create(&path) + .with_context(|| format!("failed to create temp file {}", path.display()))?; + Ok((path, file)) +} + +#[cfg(test)] +mod tests; diff --git a/crates/skippy-cache/src/l3/packed.rs b/crates/skippy-cache/src/l3/packed.rs new file mode 100644 index 0000000000..4d92d2fd87 --- /dev/null +++ b/crates/skippy-cache/src/l3/packed.rs @@ -0,0 +1,570 @@ +//! Append-only physical storage for logical L3 segments. +//! +//! One spill publishes at most one immutable pack. A local sidecar maps the +//! manifest's portable segment digests to pack offsets, keeping the handoff +//! manifest independent of this node's physical layout. + +use std::{ + collections::{HashMap, HashSet}, + fs::{self, File}, + io::{Read, Seek, SeekFrom, Write}, + path::{Path, PathBuf}, + sync::{Mutex, RwLock}, +}; + +use anyhow::{Context, Result, bail}; +use serde::{Deserialize, Serialize}; + +use super::{segment_digest, tempfile_in, write_atomically}; +use crate::fsinfo; + +pub(super) const PACK_DIR: &str = "packs"; +pub(super) const PACK_INDEX_DIR: &str = "pack-indexes"; +const PACK_INDEX_VERSION: u32 = 1; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)] +pub(super) struct PackedSegmentLocation { + pub pack_digest: String, + pub offset: u64, + pub bytes: u64, +} + +#[derive(Debug, Serialize, Deserialize)] +struct PackedManifestIndex { + version: u32, + payload_digest: String, + segments: Vec>, +} + +#[derive(Debug)] +pub(super) struct PackedStoredSegment { + pub new: bool, +} + +#[derive(Debug)] +pub(super) struct PackedReadError { + pub pack_digest: String, + pub error: anyhow::Error, +} + +#[derive(Debug)] +pub(super) struct PackedReadRequest<'a> { + pub digest: &'a str, + pub bytes: u64, + pub location: &'a PackedSegmentLocation, + pub output_offset: u64, +} + +#[derive(Debug)] +pub(super) struct PackedSegmentStore { + directory: PathBuf, + index_directory: PathBuf, + locations: RwLock>, + /// Serializes immutable pack publication with orphan collection. + mutation: Mutex<()>, +} + +impl PackedSegmentStore { + pub(super) fn open(root: &Path) -> Result { + let directory = root.join(PACK_DIR); + let index_directory = root.join(PACK_INDEX_DIR); + for path in [&directory, &index_directory] { + fsinfo::refuse_symlinked_descendant(root, path)?; + fs::create_dir_all(path) + .with_context(|| format!("failed to create {}", path.display()))?; + fsinfo::restrict_to_owner(path, 0o700)?; + } + Ok(Self { + directory, + index_directory, + locations: RwLock::new(HashMap::new()), + mutation: Mutex::new(()), + }) + } + + pub(super) fn pack_path(&self, pack_digest: &str) -> PathBuf { + self.directory.join(format!("{pack_digest}.pack")) + } + + pub(super) fn index_path(&self, payload_digest: &str) -> PathBuf { + self.index_directory.join(format!("{payload_digest}.json")) + } + + pub(super) fn estimated_new_bytes(&self, segments: &[(&str, &[u8])]) -> u64 { + let locations = self + .locations + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut unique = HashSet::new(); + segments + .iter() + .filter(|(digest, _)| unique.insert(*digest)) + .filter(|(digest, _)| { + locations + .get(*digest) + .is_none_or(|location| !self.location_is_present(location)) + }) + .map(|(_, bytes)| bytes.len() as u64) + .fold(0u64, u64::saturating_add) + } + + /// Publish all missing logical segments as one immutable pack. + pub(super) fn write_batch( + &self, + segments: &[(&str, &[u8])], + ) -> Result<(Vec, u64)> { + let _mutation = self + .mutation + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let existing = self + .locations + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); + let mut planned = HashMap::::new(); + let mut pack_segments = Vec::new(); + let mut pack_hasher = blake3::Hasher::new(); + let mut candidate_bytes = 0u64; + let mut new_digests = HashSet::new(); + + for (digest, bytes) in segments { + if existing + .get(*digest) + .is_some_and(|location| self.location_is_present(location)) + || planned.contains_key(*digest) + { + continue; + } + let offset = candidate_bytes; + candidate_bytes = candidate_bytes + .checked_add(bytes.len() as u64) + .context("packed object size overflows u64")?; + pack_hasher.update(bytes); + pack_segments.push(*bytes); + planned.insert( + (*digest).to_string(), + PackedSegmentLocation { + pack_digest: String::new(), + offset, + bytes: bytes.len() as u64, + }, + ); + new_digests.insert((*digest).to_string()); + } + + let mut new_bytes = 0u64; + if !pack_segments.is_empty() { + let pack_digest = pack_hasher.finalize().to_hex().to_string(); + let path = self.pack_path(&pack_digest); + if path.exists() { + let actual = fs::metadata(&path)?.len(); + if actual != candidate_bytes { + bail!( + "packed object {pack_digest} has {actual} bytes but the write has {candidate_bytes}" + ); + } + } else { + write_pack_atomically(&path, &pack_segments)?; + fsinfo::restrict_to_owner(&path, 0o600)?; + new_bytes = candidate_bytes; + } + for location in planned.values_mut() { + location.pack_digest.clone_from(&pack_digest); + } + } + + let mut locations = self + .locations + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + for (digest, location) in &planned { + locations.insert(digest.clone(), location.clone()); + } + let published_new_pack = new_bytes > 0; + let mut reported_new = HashSet::new(); + let stored = segments + .iter() + .map(|(digest, _)| PackedStoredSegment { + new: published_new_pack + && new_digests.contains(*digest) + && reported_new.insert(*digest), + }) + .collect(); + Ok((stored, new_bytes)) + } + + pub(super) fn location(&self, digest: &str) -> Option { + self.locations + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get(digest) + .filter(|location| self.location_is_present(location)) + .cloned() + } + + pub(super) fn encode_manifest_index( + &self, + payload_digest: &str, + segment_digests: &[String], + ) -> Result>> { + let segments = segment_digests + .iter() + .map(|digest| self.location(digest)) + .collect::>(); + if segments.iter().all(Option::is_none) { + return Ok(None); + } + serde_json::to_vec(&PackedManifestIndex { + version: PACK_INDEX_VERSION, + payload_digest: payload_digest.to_string(), + segments, + }) + .context("failed to serialize packed manifest index") + .map(Some) + } + + pub(super) fn publish_manifest_index( + &self, + payload_digest: &str, + encoded: &[u8], + ) -> Result<()> { + let path = self.index_path(payload_digest); + write_atomically(&path, encoded)?; + fsinfo::restrict_to_owner(&path, 0o600) + } + + pub(super) fn load_manifest_index( + &self, + payload_digest: &str, + segment_count: usize, + ) -> Result>> { + let path = self.index_path(payload_digest); + let bytes = match fs::read(&path) { + Ok(bytes) => bytes, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => { + return Ok(vec![None; segment_count]); + } + Err(error) => return Err(error.into()), + }; + let index: PackedManifestIndex = + serde_json::from_slice(&bytes).context("malformed packed manifest index")?; + if index.version != PACK_INDEX_VERSION + || index.payload_digest != payload_digest + || index.segments.len() != segment_count + { + bail!("packed manifest index does not match manifest {payload_digest}"); + } + for location in index.segments.iter().flatten() { + if !is_digest(&location.pack_digest) || location.bytes == 0 { + bail!("packed manifest index contains an invalid location"); + } + location + .offset + .checked_add(location.bytes) + .context("packed manifest index range overflows")?; + } + Ok(index.segments) + } + + pub(super) fn remove_manifest_index(&self, payload_digest: &str) -> Result { + let path = self.index_path(payload_digest); + let bytes = fs::metadata(&path) + .map(|metadata| metadata.len()) + .unwrap_or(0); + match fs::remove_file(&path) { + Ok(()) => Ok(bytes), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(0), + Err(error) => Err(error.into()), + } + } + + pub(super) fn validate(&self, location: &PackedSegmentLocation, bytes: u64) -> Result<()> { + if location.bytes != bytes { + bail!( + "packed index records {} bytes but manifest records {bytes}", + location.bytes + ); + } + let pack_bytes = fs::metadata(self.pack_path(&location.pack_digest)) + .with_context(|| format!("missing pack {}", location.pack_digest))? + .len(); + let end = location + .offset + .checked_add(bytes) + .context("packed segment range overflows")?; + if end > pack_bytes { + bail!( + "pack {} has {pack_bytes} bytes but segment range ends at {end}", + location.pack_digest + ); + } + Ok(()) + } + + /// Append requests directly into their final payload ranges. + /// + /// Consecutive logical segments that are also consecutive in one pack are + /// issued as one read. Digest verification still happens per logical + /// segment, preserving the manifest contract without allocating one + /// temporary `Vec` for every segment and then copying it into the payload. + pub(super) fn append_many( + &self, + requests: &[PackedReadRequest<'_>], + output: &mut Vec, + mut payload_hasher: Option<&mut blake3::Hasher>, + ) -> Result<(), PackedReadError> { + let mut files = HashMap::::new(); + let mut first = 0usize; + while first < requests.len() { + let first_request = &requests[first]; + let mut last = first + 1; + let mut physical_end = first_request + .location + .offset + .checked_add(first_request.bytes) + .context("packed read range overflows") + .map_err(|error| PackedReadError { + pack_digest: first_request.location.pack_digest.clone(), + error, + })?; + let mut output_end = first_request + .output_offset + .checked_add(first_request.bytes) + .context("packed output range overflows") + .map_err(|error| PackedReadError { + pack_digest: first_request.location.pack_digest.clone(), + error, + })?; + while let Some(next) = requests.get(last) { + if next.location.pack_digest != first_request.location.pack_digest + || next.location.offset != physical_end + || next.output_offset != output_end + { + break; + } + physical_end = + physical_end + .checked_add(next.bytes) + .ok_or_else(|| PackedReadError { + pack_digest: first_request.location.pack_digest.clone(), + error: anyhow::anyhow!("packed read range overflows"), + })?; + output_end = output_end + .checked_add(next.bytes) + .ok_or_else(|| PackedReadError { + pack_digest: first_request.location.pack_digest.clone(), + error: anyhow::anyhow!("packed output range overflows"), + })?; + last += 1; + } + + let result = (|| -> Result<()> { + if usize::try_from(first_request.output_offset) + .context("packed output offset exceeds usize")? + != output.len() + { + bail!("packed output range is not contiguous with the payload"); + } + let file = match files.entry(first_request.location.pack_digest.clone()) { + std::collections::hash_map::Entry::Occupied(entry) => entry.into_mut(), + std::collections::hash_map::Entry::Vacant(entry) => { + let file = File::open(self.pack_path(&first_request.location.pack_digest)) + .with_context(|| { + format!( + "failed to open pack {}", + first_request.location.pack_digest + ) + })?; + entry.insert(file) + } + }; + let output_end = + usize::try_from(output_end).context("packed output end exceeds usize")?; + file.seek(SeekFrom::Start(first_request.location.offset))?; + let run_bytes = output_end + .checked_sub(output.len()) + .context("packed output range precedes payload")?; + let read = file + .take(run_bytes as u64) + .read_to_end(output) + .context("failed to read packed payload range")?; + if read != run_bytes { + bail!("packed payload range ended after {read} of {run_bytes} bytes"); + } + + for request in &requests[first..last] { + let start = usize::try_from(request.output_offset) + .context("segment output offset exceeds usize")?; + let end = usize::try_from( + request + .output_offset + .checked_add(request.bytes) + .context("segment output range overflows")?, + ) + .context("segment output end exceeds usize")?; + let bytes = output + .get(start..end) + .context("segment output range exceeds payload")?; + if segment_digest(bytes) != request.digest { + bail!( + "packed segment {} failed digest verification", + request.digest + ); + } + } + if let Some(hasher) = payload_hasher.as_deref_mut() { + hasher.update(&output[output_end - run_bytes..output_end]); + } + Ok(()) + })(); + match result { + Ok(()) => {} + Err(error) => { + return Err(PackedReadError { + pack_digest: first_request.location.pack_digest.clone(), + error, + }); + } + } + first = last; + } + Ok(()) + } + + pub(super) fn rebuild( + &self, + entries: impl IntoIterator, + ) { + let mut locations = self + .locations + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + locations.clear(); + locations.extend(entries); + } + + pub(super) fn remove_orphan_packs( + &self, + referenced_segments: &HashSet, + held_segments: impl FnOnce() -> HashSet, + ) -> Result { + let _mutation = self + .mutation + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + // Snapshot in-flight publications only after taking the same lock that + // serializes pack publication, so collection cannot race a completed + // pack into existence before its manifest is committed. + let held_segments = held_segments(); + let locations = self + .locations + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let protected = locations + .iter() + .filter(|(digest, _)| { + referenced_segments.contains(*digest) || held_segments.contains(*digest) + }) + .map(|(_, location)| location.pack_digest.clone()) + .collect::>(); + let mut freed = 0u64; + for entry in fs::read_dir(&self.directory)? { + let entry = entry?; + let path = entry.path(); + if path.extension().is_none_or(|extension| extension != "pack") { + continue; + } + let Some(pack_digest) = path.file_stem().and_then(|stem| stem.to_str()) else { + continue; + }; + if protected.contains(pack_digest) { + continue; + } + freed = freed.saturating_add(entry.metadata()?.len()); + fs::remove_file(&path) + .with_context(|| format!("failed to collect pack {pack_digest}"))?; + } + Ok(freed) + } + + pub(super) fn remove_orphan_indexes(&self, manifests: &HashSet) -> Result { + let mut freed = 0u64; + for entry in fs::read_dir(&self.index_directory)? { + let entry = entry?; + let path = entry.path(); + if path.extension().is_none_or(|extension| extension != "json") { + continue; + } + let Some(payload_digest) = path.file_stem().and_then(|stem| stem.to_str()) else { + continue; + }; + if manifests.contains(payload_digest) { + continue; + } + freed = freed.saturating_add(entry.metadata()?.len()); + fs::remove_file(&path)?; + } + Ok(freed) + } + + pub(super) fn footprint_bytes(&self) -> Result { + directory_bytes(&self.directory) + } + + pub(super) fn index_footprint_bytes(&self) -> Result { + directory_bytes(&self.index_directory) + } + + fn location_is_present(&self, location: &PackedSegmentLocation) -> bool { + self.pack_path(&location.pack_digest) + .metadata() + .is_ok_and(|metadata| location.offset < metadata.len()) + } +} + +fn directory_bytes(directory: &Path) -> Result { + let mut total = 0u64; + for entry in fs::read_dir(directory)? { + let entry = entry?; + if entry.file_type()?.is_file() { + total = total.saturating_add(entry.metadata()?.len()); + } + } + Ok(total) +} + +fn write_pack_atomically(path: &Path, segments: &[&[u8]]) -> Result<()> { + let directory = path.parent().context("pack path has no parent directory")?; + let (temp_path, mut temp_file) = tempfile_in(directory)?; + let write_result = (|| -> Result<()> { + for segment in segments { + temp_file + .write_all(segment) + .with_context(|| format!("failed to write {}", temp_path.display()))?; + } + temp_file + .sync_all() + .with_context(|| format!("failed to sync {}", temp_path.display())) + })(); + drop(temp_file); + let publish_result = write_result.and_then(|()| fsinfo::replace_file(&temp_path, path)); + if let Err(error) = publish_result { + return match fs::remove_file(&temp_path) { + Ok(()) => Err(error), + Err(cleanup) if cleanup.kind() == std::io::ErrorKind::NotFound => Err(error), + Err(cleanup) => Err(error.context(format!( + "also failed to remove temporary pack {}: {cleanup}", + temp_path.display() + ))), + }; + } + Ok(()) +} + +fn is_digest(value: &str) -> bool { + value.len() == 64 + && value + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) +} diff --git a/crates/skippy-cache/src/l3/tests.rs b/crates/skippy-cache/src/l3/tests.rs new file mode 100644 index 0000000000..12828f800a --- /dev/null +++ b/crates/skippy-cache/src/l3/tests.rs @@ -0,0 +1,1263 @@ +//! Tests for the L3 segment store. +//! +//! Split out of `l3.rs` to keep that file under the 2,000-line limit the +//! coding guidelines set. + +use super::*; + +fn store(root: &Path, budget: u64) -> HandoffSegmentStore { + HandoffSegmentStore::open(root, budget).expect("open store") +} + +/// Builds a manifest and returns the write-side holds alongside it: a +/// caller that puts segments and commits later must keep them alive, or an +/// eviction in between collects the segments it is about to reference. +/// This is the contract `L3Tier::spill` follows in production. +fn manifest_for<'store>( + store: &'store HandoffSegmentStore, + payload: &[u8], + segment_bytes: usize, +) -> (HandoffManifest, Vec>) { + let mut manifest = HandoffManifest::new("blake3:test".to_string(), "full-state".into()); + let mut held = Vec::new(); + for (index, chunk) in payload.chunks(segment_bytes).enumerate() { + let stored = store.put_segment(chunk).expect("put segment"); + manifest.segments.push(HandoffSegmentRef { + index: index as u32, + offset: (index * segment_bytes) as u64, + bytes: chunk.len() as u64, + digest: stored.digest.clone(), + codec_identity: Some(SegmentCodecIdentity::raw(chunk.len() as u64)), + meta_json: None, + }); + held.push(stored); + } + manifest.total_bytes = payload.len() as u64; + manifest.payload_digest = segment_digest(payload); + (manifest, held) +} + +/// Put a payload's segments and commit the manifest that binds them, +/// releasing the write-side holds afterwards. The shape production uses: +/// hold across the commit, then let eviction have them. +fn commit_payload( + store: &HandoffSegmentStore, + payload: &[u8], + segment_bytes: usize, +) -> HandoffManifest { + let (manifest, held) = manifest_for(store, payload, segment_bytes); + store.commit(&manifest).expect("commit"); + drop(held); + manifest +} + +fn commit_packed_payload( + store: &HandoffSegmentStore, + payload: &[u8], + segment_bytes: usize, +) -> HandoffManifest { + let chunks = payload.chunks(segment_bytes).collect::>(); + let held = store + .try_put_segments(&chunks) + .expect("packed put") + .expect("packed put admitted"); + let mut manifest = HandoffManifest::new("blake3:test".to_string(), "full-state".into()); + let mut offset = 0u64; + for (index, stored) in held.iter().enumerate() { + let bytes = chunks[index].len() as u64; + manifest.segments.push(HandoffSegmentRef { + index: index as u32, + offset, + bytes, + digest: stored.digest.clone(), + codec_identity: Some(SegmentCodecIdentity::raw(bytes)), + meta_json: None, + }); + offset += bytes; + } + manifest.total_bytes = payload.len() as u64; + manifest.payload_digest = segment_digest(payload); + store.commit(&manifest).expect("packed commit"); + drop(held); + manifest +} + +fn temp_root(name: &str) -> PathBuf { + let root = std::env::temp_dir() + .join("skippy-l3-tests") + .join(format!("{name}-{}", std::process::id())); + let _ = fs::remove_dir_all(&root); + root +} + +#[test] +fn roundtrip_assembles_identical_payload() { + let root = temp_root("roundtrip"); + let store = store(&root, 0); + let payload: Vec = (0..100_000u32).map(|value| value as u8).collect(); + let manifest = commit_payload(&store, &payload, 4096); + let loaded = store + .load_manifest(&manifest.payload_digest) + .expect("load manifest"); + assert_eq!(store.assemble(&loaded).expect("assemble"), payload); +} + +#[test] +fn packed_roundtrip_uses_one_physical_file_and_survives_reopen() { + let root = temp_root("packed-roundtrip"); + let payload: Vec = (0..100_000u32).map(|value| value as u8).collect(); + let manifest = { + let store = store(&root, 0); + let manifest = commit_packed_payload(&store, &payload, 4096); + assert_eq!(fs::read_dir(root.join(PACK_DIR)).unwrap().count(), 1); + assert_eq!(fs::read_dir(root.join(SEGMENT_DIR)).unwrap().count(), 0); + let manifest_json = fs::read_to_string(store.manifest_path(&manifest.payload_digest)) + .expect("read portable manifest"); + assert!(!manifest_json.contains("pack_digest")); + assert_eq!(store.assemble(&manifest).expect("assemble"), payload); + manifest + }; + + // Direct callers receive a fully reconciled store, including the packed + // location map needed to read manifests from the previous process. + let reopened = store(&root, 0); + let loaded = reopened + .load_manifest(&manifest.payload_digest) + .expect("load packed manifest after restart"); + assert_eq!(reopened.assemble(&loaded).expect("assemble"), payload); +} + +#[test] +fn corrupt_pack_is_quarantined_and_never_served() { + let root = temp_root("packed-corruption"); + let store = store(&root, 0); + let payload: Vec = (0..32_000u32).map(|value| value as u8).collect(); + let manifest = commit_packed_payload(&store, &payload, 4096); + let pack = fs::read_dir(root.join(PACK_DIR)) + .unwrap() + .next() + .unwrap() + .unwrap() + .path(); + let mut bytes = fs::read(&pack).unwrap(); + bytes[0] ^= 0xff; + fs::write(&pack, bytes).unwrap(); + + assert!(store.assemble(&manifest).is_err()); + assert!(!pack.exists()); + assert!(root.join(QUARANTINE_DIR).exists()); +} + +#[test] +fn uncommitted_pack_is_collected_after_holds_release() { + let root = temp_root("packed-orphan"); + let store = store(&root, 0); + let payload = (0..16_000) + .map(|index| (index / 1024) as u8) + .collect::>(); + let chunks = payload.chunks(1024).collect::>(); + let held = store + .try_put_segments(&chunks) + .unwrap() + .expect("packed put admitted"); + assert_eq!(store.collect_unreferenced_segments().unwrap(), 0); + drop(held); + assert_eq!( + store.collect_unreferenced_segments().unwrap(), + payload.len() as u64 + ); +} + +#[test] +fn cached_usage_never_diverges_from_a_full_scan() { + // The incremental total exists to keep `reserve` off an O(files) scan + // per segment put. It is only safe while it agrees with the disk, so + // check it after every kind of mutation the store performs. + let root = temp_root("usage-drift"); + let store = store(&root, 0); + let reserved = || store.reserved_inflight.load(Ordering::Acquire); + let scanned = + |store: &HandoffSegmentStore| store.rescan_usage_bytes().expect("rescan") + reserved(); + + let assert_agrees = |store: &HandoffSegmentStore, stage: &str| { + let cached = store.managed_usage_bytes().expect("cached usage"); + let truth = scanned(store); + assert_eq!(cached, truth, "cached usage diverged after {stage}"); + }; + + assert_agrees(&store, "open"); + + // Segment puts: the one path that adjusts the total incrementally. + let payload: Vec = (0..50_000u32).map(|value| value as u8).collect(); + let manifest = commit_payload(&store, &payload, 4096); + assert_agrees(&store, "put and commit"); + + // A prefix link is a new file under the index tree. + store + .link_prefix("namespace", 2, "prefix-key", &manifest.payload_digest) + .expect("link prefix"); + assert_agrees(&store, "link_prefix"); + + // A hit touches metadata only. + let _ = store.manifest_for_prefix("namespace", 2, "prefix-key"); + assert_agrees(&store, "prefix hit"); + + // Bulk removal. + store.clear().expect("clear"); + assert_agrees(&store, "clear"); + + let _ = fs::remove_dir_all(&root); +} + +#[test] +fn puts_are_idempotent_and_deduplicated() { + let root = temp_root("idempotent"); + let store = store(&root, 0); + let first = store.put_segment(b"same bytes").expect("first put"); + let first_digest = first.digest.clone(); + let second = store.put_segment(b"same bytes").expect("second put"); + let second_digest = second.digest.clone(); + assert_eq!(first_digest, second_digest); + assert!(first.put.new); + assert!(!second.put.new); + assert_eq!(store.segment_footprint_bytes().expect("footprint"), 10); + drop(first); + assert_eq!( + store.collect_unreferenced_segments().unwrap(), + 0, + "one writer released a segment still held by another writer" + ); + drop(second); + assert_eq!(store.collect_unreferenced_segments().unwrap(), 10); +} + +#[test] +fn commit_rejects_missing_segments_and_bad_tiling() { + let root = temp_root("completeness"); + let store = store(&root, 0); + let payload = vec![7u8; 10_000]; + let (mut manifest, _held) = manifest_for(&store, &payload, 4096); + + let mut missing = manifest.clone(); + missing.segments[1].digest = segment_digest(b"never stored"); + assert!(store.commit(&missing).is_err()); + + manifest.segments[2].offset += 1; + assert!(store.commit(&manifest).is_err()); +} + +#[test] +fn corrupted_segment_fails_verification_on_read() { + let root = temp_root("corruption"); + let store = store(&root, 0); + let payload = vec![42u8; 8192]; + let manifest = commit_payload(&store, &payload, 4096); + + let victim = store.segment_path(&manifest.segments[0].digest); + let mut bytes = fs::read(&victim).expect("read segment file"); + bytes[0] ^= 0xFF; + fs::write(&victim, bytes).expect("corrupt segment file"); + + assert!(store.assemble(&manifest).is_err()); +} + +#[test] +fn budget_evicts_the_least_recently_used_manifest() { + let root = temp_root("budget"); + // Budget fits one payload but not two. + let store = store(&root, 12_000); + let old_payload = vec![1u8; 8_000]; + let new_payload = vec![2u8; 8_000]; + let old_manifest = commit_payload(&store, &old_payload, 4096); + // Ensure a later mtime for the second manifest. + std::thread::sleep(std::time::Duration::from_millis(20)); + let new_manifest = commit_payload(&store, &new_payload, 4096); + + let manifests = store.list_manifests().expect("list"); + assert!( + !manifests.contains(&old_manifest.payload_digest), + "the older entry survived eviction: {manifests:?}" + ); + assert_eq!(manifests, vec![new_manifest.payload_digest.clone()]); + assert!(store.assemble(&new_manifest).is_ok()); + assert!(store.segment_footprint_bytes().expect("footprint") <= 12_000); +} + +#[test] +fn eviction_follows_last_use_not_last_write() { + let root = temp_root("lru-by-use"); + // Fits two payloads plus bookkeeping, not three. + let store = store(&root, 20_000); + let first = commit_payload(&store, &vec![1u8; 8_000], 4096); + std::thread::sleep(std::time::Duration::from_millis(20)); + let second = commit_payload(&store, &vec![2u8; 8_000], 4096); + + // The older entry is the one being read, so it is the one that should + // survive. Under least-recently-written it would be evicted first. + std::thread::sleep(std::time::Duration::from_millis(20)); + store.touch_manifest(&first.payload_digest); + + std::thread::sleep(std::time::Duration::from_millis(20)); + // A third entry the budget cannot hold: something must go. + commit_payload(&store, &vec![3u8; 8_000], 4096); + + let manifests = store.list_manifests().expect("list"); + assert!( + manifests.contains(&first.payload_digest), + "the recently used entry was evicted: {manifests:?}" + ); + assert!( + !manifests.contains(&second.payload_digest), + "the least recently used entry survived: {manifests:?}" + ); +} + +#[test] +fn a_segment_larger_than_the_budget_is_refused() { + let root = temp_root("oversize-segment"); + let store = store(&root, 4_000); + let refusal = store + .try_put_segment(&vec![7u8; 16_000]) + .expect("put") + .expect_err("a segment larger than the whole budget was stored"); + assert_eq!(refusal, WriteRefusal::SkippedOversize); + assert_eq!(refusal.reason(), "skipped_oversize"); + assert_eq!(store.segment_footprint_bytes().expect("footprint"), 0); +} + +#[test] +fn an_entry_larger_than_a_shrunken_budget_is_refused_at_commit() { + // A live budget update can make an in-flight entry larger than the cap: + // it was admissible when its segments were written and is not any more. + let root = temp_root("oversize-commit"); + let store = store(&root, 0); + let (manifest, held) = manifest_for(&store, &vec![7u8; 16_000], 4_000); + store + .update_limits(StoreLimits::new(8_000, 0)) + .expect("shrink limits"); + let error = store + .commit(&manifest) + .expect_err("an entry larger than the budget was committed"); + drop(held); + assert!( + format!("{error:#}").contains("skipped_oversize"), + "refusal did not carry the reason code: {error:#}" + ); + assert!( + store.list_manifests().expect("list").is_empty(), + "the refused entry was left loadable" + ); +} + +#[test] +fn managed_usage_counts_more_than_segments() { + let root = temp_root("usage"); + let store = store(&root, 0); + let manifest = commit_payload(&store, &vec![5u8; 4096], 4096); + store.commit(&manifest).expect("commit"); + store + .link_prefix("namespace", 128, "prefix", &manifest.payload_digest) + .expect("link prefix"); + + let segments = store.segment_footprint_bytes().expect("footprint"); + let managed = store.managed_usage_bytes().expect("usage"); + assert!( + managed > segments, + "managed usage {managed} ignored manifests and index files (segments {segments})" + ); +} + +#[test] +fn prefix_links_obey_the_hard_budget_before_creating_the_index_tree() { + let root = temp_root("prefix-budget"); + let store = store(&root, 0); + let manifest = commit_payload(&store, &vec![5u8; 4096], 4096); + let pin = store.pin(&manifest.payload_digest); + let used = store.managed_usage_bytes().expect("usage before link"); + store + .update_limits(StoreLimits::new(used, 0)) + .expect("set exact hard cap"); + + let error = store + .link_prefix("namespace", 128, "prefix", &manifest.payload_digest) + .expect_err("prefix link exceeded the hard budget"); + assert!( + format!("{error:#}").contains("insufficient_space"), + "unexpected refusal: {error:#}" + ); + assert!( + !store.namespace_dir("namespace").exists(), + "a refused prefix link created index directories" + ); + assert_eq!(store.managed_usage_bytes().expect("usage after link"), used); + drop(pin); +} + +#[test] +fn atomic_publish_removes_temporary_file_after_rename_failure() { + let root = temp_root("atomic-cleanup"); + fs::create_dir_all(&root).expect("create root"); + let destination = root.join("destination"); + fs::create_dir(&destination).expect("create blocking directory"); + + write_atomically(&destination, b"partial bytes") + .expect_err("publishing a file over a directory succeeded"); + + let leftovers: Vec<_> = fs::read_dir(&root) + .expect("read root") + .filter_map(Result::ok) + .filter(|entry| entry.file_name().to_string_lossy().starts_with(".tmp-")) + .collect(); + assert!( + leftovers.is_empty(), + "temporary files survived: {leftovers:?}" + ); +} + +#[test] +fn a_pinned_manifest_outranks_a_new_write() { + // Under pressure the store refuses the incoming write rather than + // pulling state out from under an operation still using it. The new + // entry is a miss; the pinned one stays loadable. + let root = temp_root("pinned"); + let store = store(&root, 12_000); + let pinned = commit_payload(&store, &vec![1u8; 8_000], 4096); + store.commit(&pinned).expect("commit pinned"); + let guard = store.pin(&pinned.payload_digest); + + std::thread::sleep(std::time::Duration::from_millis(20)); + let refusal = store + .try_put_segment(&vec![2u8; 8_000]) + .expect("put") + .expect_err("the pinned entry was evicted to admit a new write"); + assert_eq!(refusal, WriteRefusal::InsufficientSpace); + + let manifests = store.list_manifests().expect("list"); + assert!( + manifests.contains(&pinned.payload_digest), + "a pinned manifest was evicted: {manifests:?}" + ); + + // Once nothing is using it, the same write is admitted. + drop(guard); + store + .try_put_segment(&vec![2u8; 8_000]) + .expect("put") + .expect("the write stayed refused after the pin was released"); +} + +#[test] +fn the_free_space_reserve_refuses_writes() { + let root = temp_root("reserve"); + // A reserve no filesystem can satisfy, rather than one derived from + // live free space: another test freeing a few MiB mid-run must not + // decide whether this one passes. + let store = HandoffSegmentStore::open_with_limits(&root, StoreLimits::new(0, u64::MAX)) + .expect("open store"); + let refusal = store + .try_put_segment(b"bytes that do not fit the reserve") + .expect("put") + .expect_err("write was admitted below the reserve"); + assert_eq!(refusal, WriteRefusal::ReadOnlyLowSpace); + assert_eq!(refusal.reason(), "read_only_low_space"); +} + +#[test] +fn reservations_are_released_after_the_write() { + let root = temp_root("reservation"); + let store = store(&root, 1_000_000); + store.put_segment(b"some bytes").expect("put"); + assert_eq!( + store.usage().expect("usage").reserved_inflight_bytes, + 0, + "a completed write left capacity reserved" + ); +} + +#[test] +fn clear_removes_every_unpinned_entry() { + let root = temp_root("clear"); + let store = store(&root, 0); + let linked = commit_payload(&store, &vec![1u8; 4096], 4096); + store + .link_prefix("namespace", 128, "prefix", &linked.payload_digest) + .unwrap(); + commit_payload(&store, &vec![2u8; 4096], 4096); + + let freed = store.clear().expect("clear"); + assert!(freed > 0, "clear freed nothing"); + assert!(store.list_manifests().expect("list").is_empty()); + assert_eq!(store.segment_footprint_bytes().expect("footprint"), 0); + assert!( + store + .recorded_prefix_lengths("namespace") + .unwrap() + .is_empty(), + "clear left a dangling prefix link" + ); +} + +#[test] +fn prune_frees_down_to_the_target() { + let root = temp_root("prune"); + let store = store(&root, 0); + for fill in 1u8..=3 { + let manifest = commit_payload(&store, &vec![fill; 8_000], 4096); + store.commit(&manifest).expect("commit"); + std::thread::sleep(std::time::Duration::from_millis(20)); + } + let before = store.managed_usage_bytes().expect("usage"); + store.prune_to(before / 2).expect("prune"); + let after = store.managed_usage_bytes().expect("usage"); + assert!(after < before, "prune freed nothing ({before} -> {after})"); +} + +#[test] +fn live_limit_update_changes_the_pair_and_prunes_inactive_entries() { + let root = temp_root("live-limits"); + let store = store(&root, 1_000_000); + for fill in 1u8..=3 { + let manifest = commit_payload(&store, &vec![fill; 8_000], 4096); + store.commit(&manifest).expect("commit"); + } + let before = store.managed_usage_bytes().expect("usage"); + let next = StoreLimits::new(before / 2, 4096); + let previous = store.update_limits(next).expect("update limits"); + assert_eq!(previous, StoreLimits::new(1_000_000, 0)); + assert_eq!(store.limits(), next); + assert!( + store.managed_usage_bytes().expect("usage after") <= next.budget_bytes, + "live shrink did not prune to the new cap" + ); +} + +#[test] +fn a_corrupt_segment_is_quarantined_not_left_in_place() { + let root = temp_root("quarantine"); + let store = store(&root, 0); + let digest = store.put_segment(b"segment bytes").expect("put").digest; + fs::write( + root.join("segments").join(format!("{digest}.seg")), + b"tampered", + ) + .expect("tamper with the segment"); + + let error = store + .read_segment(&digest) + .expect_err("a tampered segment was served"); + assert!( + format!("{error:#}").contains("quarantined"), + "corrupt segment was not quarantined: {error:#}" + ); + assert!( + !store.has_segment(&digest), + "the corrupt segment is still in the managed tree" + ); + assert!( + root.join("quarantine").exists(), + "nothing was moved to quarantine" + ); +} + +/// Not a pass/fail assertion: a stopwatch on the cost that smaller windows +/// buy. Eviction parses every manifest to build its reference map, and a +/// 64-row window turns a 19K-token entry into ~9.5k segment refs. Run with +/// `cargo test -p skippy-cache --lib eviction_cost -- --ignored --nocapture`. +#[test] +#[ignore = "measurement, not a check; takes tens of seconds"] +fn eviction_cost_at_realistic_segment_counts() { + const SEGMENTS_PER_MANIFEST: usize = 9_504; // 16 layers x 2 x ceil(19000/64) + const MANIFESTS: usize = 20; + let root = temp_root("eviction-cost"); + let store = store(&root, 0); + + // One physical segment shared by every ref: this measures manifest + // parsing and reference mapping, not filesystem write throughput. + let bytes = vec![7u8; 65_536]; + let digest = store.put_segment(&bytes).expect("put").digest; + let build = std::time::Instant::now(); + for manifest_index in 0..MANIFESTS { + let mut manifest = HandoffManifest::new("blake3:cost".to_string(), "full-state".into()); + for index in 0..SEGMENTS_PER_MANIFEST { + manifest.segments.push(HandoffSegmentRef { + index: index as u32, + offset: (index * bytes.len()) as u64, + bytes: bytes.len() as u64, + digest: digest.clone(), + codec_identity: Some(SegmentCodecIdentity::raw(bytes.len() as u64)), + meta_json: Some(format!("k:{}:0:{}", index % 32, index / 32)), + }); + } + manifest.total_bytes = (SEGMENTS_PER_MANIFEST * bytes.len()) as u64; + manifest.payload_digest = format!("blake3:manifest-{manifest_index}"); + store.commit(&manifest).expect("commit"); + } + let build_ms = build.elapsed().as_millis(); + + let manifest_bytes = directory_bytes(&root.join(MANIFEST_DIR)).expect("manifest bytes"); + let usage = store.managed_usage_bytes().expect("usage"); + let evict = std::time::Instant::now(); + store.enforce_budget_to(usage / 2).expect("enforce"); + let evict_ms = evict.elapsed().as_millis(); + + println!( + "eviction cost: {MANIFESTS} manifests x {SEGMENTS_PER_MANIFEST} refs, \ + manifest bytes {manifest_bytes} ({} KiB each), build {build_ms} ms, \ + enforce_budget {evict_ms} ms", + manifest_bytes / MANIFESTS as u64 / 1024 + ); + let _ = fs::remove_dir_all(&root); +} + +#[test] +fn eviction_leaves_headroom_so_a_full_cache_is_not_repriced_per_commit() { + let root = temp_root("low-water"); + let store = store(&root, 40_000); + let mut eviction_triggered = false; + for fill in 1u8..=16 { + let manifest = commit_payload(&store, &vec![fill; 8_000], 4096); + store.commit(&manifest).expect("commit"); + if store.usage().unwrap().evicted_manifests > 0 { + eviction_triggered = true; + break; + } + std::thread::sleep(std::time::Duration::from_millis(5)); + } + assert!(eviction_triggered, "fixture never crossed the budget"); + let usage = store.managed_usage_bytes().expect("usage"); + assert!( + usage < 40_000, + "eviction left no headroom below the cap: {usage}" + ); + assert_eq!( + store.enforce_budget().expect("second pass"), + 0, + "a store already under the cap must not pay for another pass" + ); +} + +#[test] +fn unreferenced_segments_are_collected() { + let root = temp_root("gc"); + let store = store(&root, 0); + store.put_segment(b"orphan bytes").expect("orphan put"); + let payload = vec![9u8; 4096]; + let manifest = commit_payload(&store, &payload, 4096); + + let freed = store.collect_unreferenced_segments().expect("collect"); + assert_eq!(freed, 12); + assert!(store.assemble(&manifest).is_ok()); +} + +#[test] +fn manifest_stamps_explicit_raw_codec_and_round_trips() { + let root = temp_root("codec-raw-roundtrip"); + let store = store(&root, 0); + let payload: Vec = (0..50_000u32).map(|value| value as u8).collect(); + let manifest = commit_payload(&store, &payload, 4096); + assert_eq!(manifest.codec, Some(PayloadCodec::raw())); + + // The codec identity is written explicitly, not inferred at read time. + let manifest_json = fs::read_to_string(store.manifest_path(&manifest.payload_digest)) + .expect("read manifest json"); + let value: serde_json::Value = + serde_json::from_str(&manifest_json).expect("parse manifest json"); + assert_eq!(value["codec"]["name"], CODEC_RAW); + assert_eq!(value["codec"]["version"], CODEC_RAW_VERSION); + + let loaded = store + .load_manifest(&manifest.payload_digest) + .expect("load manifest"); + assert_eq!(loaded.codec, Some(PayloadCodec::raw())); + assert_eq!(store.assemble(&loaded).expect("assemble"), payload); +} + +#[test] +fn legacy_manifest_without_codec_field_reads_and_assembles_as_raw() { + let root = temp_root("codec-legacy-read"); + let store = store(&root, 0); + let payload: Vec = (0..40_000u32).map(|value| value as u8).collect(); + let manifest = commit_payload(&store, &payload, 4096); + + // Rewrite the on-disk manifest as an older build wrote it: the legacy + // format version and no codec field. + let path = store.manifest_path(&manifest.payload_digest); + let mut value: serde_json::Value = + serde_json::from_slice(&fs::read(&path).expect("read manifest")).expect("parse manifest"); + { + let object = value.as_object_mut().expect("manifest object"); + object.insert( + "version".to_string(), + serde_json::json!(LEGACY_MANIFEST_VERSION), + ); + object.remove("codec"); + } + assert!(value.get("codec").is_none(), "legacy manifest has no codec"); + fs::write(&path, serde_json::to_vec(&value).expect("serialize legacy")) + .expect("write legacy manifest"); + + let loaded = store + .load_manifest(&manifest.payload_digest) + .expect("load legacy manifest"); + assert_eq!( + loaded.codec, + Some(PayloadCodec::raw()), + "a manifest without a codec field is the legacy raw format" + ); + assert_eq!(store.assemble(&loaded).expect("assemble legacy"), payload); +} + +#[test] +fn unknown_codec_is_refused_at_commit_and_leaves_no_manifest() { + let root = temp_root("codec-unknown-commit"); + let store = store(&root, 0); + let payload = vec![7u8; 8192]; + let (mut manifest, held) = manifest_for(&store, &payload, 4096); + manifest.codec = Some(PayloadCodec { + name: "zstd".to_string(), + version: 1, + }); + + let error = store + .commit(&manifest) + .expect_err("unknown codec must not commit"); + assert!( + error.to_string().contains("codec"), + "commit error should name the codec: {error}" + ); + drop(held); + + // Fallback contract: nothing unassemblable was persisted, so a later + // restore is a clean miss rather than a broken entry. + assert!( + store.load_manifest(&manifest.payload_digest).is_err(), + "no manifest should exist after a refused unknown-codec commit" + ); +} + +#[test] +fn unknown_codec_is_rejected_before_assembly() { + let root = temp_root("codec-unknown-assemble"); + let store = store(&root, 0); + let payload = vec![3u8; 8192]; + let manifest = commit_payload(&store, &payload, 4096); + let mut loaded = store + .load_manifest(&manifest.payload_digest) + .expect("load manifest"); + loaded.codec = Some(PayloadCodec { + name: "lz4".to_string(), + version: 1, + }); + let error = store + .assemble(&loaded) + .expect_err("unknown codec must not assemble"); + assert!( + error.to_string().contains("codec"), + "assemble error should name the codec: {error}" + ); +} + +#[test] +fn unknown_raw_codec_version_is_rejected_at_commit_and_assembly() { + let root = temp_root("codec-unknown-version"); + let store = store(&root, 0); + let payload = vec![5u8; 8192]; + + // Writer side: a future raw version is refused, never migrated. + let (mut manifest, held) = manifest_for(&store, &payload, 4096); + manifest.codec = Some(PayloadCodec { + name: CODEC_RAW.to_string(), + version: CODEC_RAW_VERSION + 1, + }); + assert!( + store.commit(&manifest).is_err(), + "a future raw codec version must not commit" + ); + drop(held); + + // Reader side: a valid raw entry re-tagged to a future version is refused. + let good = commit_payload(&store, &payload, 4096); + let mut loaded = store + .load_manifest(&good.payload_digest) + .expect("load manifest"); + loaded.codec.as_mut().expect("codec present").version = CODEC_RAW_VERSION + 1; + assert!( + store.assemble(&loaded).is_err(), + "a future raw codec version must not assemble" + ); +} + +#[test] +fn codec_gate_does_not_mask_payload_corruption() { + let root = temp_root("codec-corruption"); + let store = store(&root, 0); + let payload = vec![1u8; 8192]; + let manifest = commit_packed_payload(&store, &payload, 4096); + + // Supported codec, but the underlying bytes are tampered: the codec check + // passes and digest verification still catches the corruption. + let loaded = store + .load_manifest(&manifest.payload_digest) + .expect("load manifest"); + assert_eq!(loaded.codec, Some(PayloadCodec::raw())); + let pack = fs::read_dir(root.join(PACK_DIR)) + .expect("read packs") + .next() + .expect("one pack") + .expect("pack entry") + .path(); + let mut bytes = fs::read(&pack).expect("read pack"); + bytes[0] ^= 0xFF; + fs::write(&pack, &bytes).expect("corrupt pack"); + + assert!( + store.assemble(&loaded).is_err(), + "corruption under a supported codec must still fail" + ); +} + +/// Overwrite the `codec` object of an on-disk manifest, simulating a +/// future/remote entry this build cannot decode. +fn rewrite_on_disk_codec(store: &HandoffSegmentStore, digest: &str, name: &str, version: u32) { + let path = store.manifest_path(digest); + let mut value: serde_json::Value = + serde_json::from_slice(&fs::read(&path).expect("read manifest")).expect("parse manifest"); + value["codec"] = serde_json::json!({ "name": name, "version": version }); + fs::write(&path, serde_json::to_vec(&value).expect("serialize")).expect("write manifest"); +} + +#[test] +fn stripping_codec_from_current_version_rejects_but_legacy_v2_reads_as_raw() { + let root = temp_root("codec-downgrade"); + let store = store(&root, 0); + let payload = vec![2u8; 8192]; + let manifest = commit_payload(&store, &payload, 4096); + let path = store.manifest_path(&manifest.payload_digest); + let original: serde_json::Value = + serde_json::from_slice(&fs::read(&path).expect("read manifest")).expect("parse manifest"); + assert_eq!(original["version"], MANIFEST_VERSION); + + // Current version with the codec stripped: reject, never default to raw. + let mut stripped = original.clone(); + stripped.as_object_mut().expect("object").remove("codec"); + fs::write(&path, serde_json::to_vec(&stripped).expect("serialize")).expect("write stripped"); + assert!( + store.load_manifest(&manifest.payload_digest).is_err(), + "a current-version manifest with codec removed must be rejected" + ); + + // A genuine legacy v2 manifest without a codec still reads/assembles as raw. + let mut legacy = original; + { + let object = legacy.as_object_mut().expect("object"); + object.insert( + "version".to_string(), + serde_json::json!(LEGACY_MANIFEST_VERSION), + ); + object.remove("codec"); + } + fs::write(&path, serde_json::to_vec(&legacy).expect("serialize")).expect("write legacy"); + let loaded = store + .load_manifest(&manifest.payload_digest) + .expect("legacy v2 load"); + assert_eq!(loaded.codec, Some(PayloadCodec::raw())); + assert_eq!(store.assemble(&loaded).expect("assemble legacy"), payload); +} + +#[test] +fn on_disk_unsupported_codec_fails_direct_load() { + let root = temp_root("codec-load-reject"); + let store = store(&root, 0); + let payload = vec![4u8; 8192]; + let manifest = commit_payload(&store, &payload, 4096); + rewrite_on_disk_codec(&store, &manifest.payload_digest, "zstd", 1); + let error = store + .load_manifest(&manifest.payload_digest) + .expect_err("unsupported codec must not load"); + assert!( + error.to_string().contains("codec"), + "load error should name the codec: {error}" + ); +} + +#[test] +fn startup_reconciliation_quarantines_unsupported_codec_manifest() { + let root = temp_root("codec-reconcile"); + let payload = vec![6u8; 8192]; + let digest = { + let store = store(&root, 0); + let manifest = commit_packed_payload(&store, &payload, 4096); + rewrite_on_disk_codec(&store, &manifest.payload_digest, "future", 9); + manifest.payload_digest + }; + + let reopened = + HandoffSegmentStore::open_unreconciled_with_limits(&root, StoreLimits::new(0, 0)) + .expect("open unreconciled store"); + let report = reopened.reconcile_startup().expect("reconcile"); + assert_eq!( + report.quarantined_manifests, 1, + "an unsupported-codec manifest is not a valid committed entry" + ); + assert!( + reopened.load_manifest(&digest).is_err(), + "the quarantined manifest is gone from the live set" + ); + assert!(root.join(QUARANTINE_DIR).exists()); +} + +#[test] +fn future_version_raw_manifest_is_refused_at_commit_and_leaves_no_manifest() { + let root = temp_root("codec-future-version-commit"); + let store = store(&root, 0); + let payload = vec![8u8; 8192]; + let (mut manifest, held) = manifest_for(&store, &payload, 4096); + // A supported (raw) codec but an unknown future version: commit must refuse + // it, or it would persist a manifest load_manifest immediately rejects. + manifest.version = MANIFEST_VERSION + 1; + let error = store + .commit(&manifest) + .expect_err("a future manifest version must not commit"); + assert!( + error.to_string().contains("version"), + "commit error should name the version: {error}" + ); + drop(held); + assert!( + store.load_manifest(&manifest.payload_digest).is_err(), + "no manifest should exist after a refused future-version commit" + ); +} + +#[test] +fn future_version_raw_manifest_is_refused_before_assembly() { + let root = temp_root("codec-future-version-assemble"); + let store = store(&root, 0); + let payload = vec![9u8; 8192]; + let manifest = commit_payload(&store, &payload, 4096); + let mut loaded = store + .load_manifest(&manifest.payload_digest) + .expect("load manifest"); + loaded.version = MANIFEST_VERSION + 1; + let error = store + .assemble(&loaded) + .expect_err("a future manifest version must not assemble"); + assert!( + error.to_string().contains("version"), + "assemble error should name the version: {error}" + ); +} + +/// Rewrites the on-disk manifest as a v3 (#1750) build wrote it: payload-level +/// codec identity only, no per-segment identity. +fn rewrite_on_disk_as_v3(store: &HandoffSegmentStore, digest: &str) { + let path = store.manifest_path(digest); + let mut value: serde_json::Value = + serde_json::from_slice(&fs::read(&path).expect("read manifest")).expect("parse manifest"); + let object = value.as_object_mut().expect("manifest object"); + object.insert( + "version".to_string(), + serde_json::json!(LEGACY_PAYLOAD_CODEC_MANIFEST_VERSION), + ); + for segment in object + .get_mut("segments") + .expect("segments") + .as_array_mut() + .expect("segment array") + { + segment + .as_object_mut() + .expect("segment object") + .remove("codec_identity"); + } + fs::write(&path, serde_json::to_vec(&value).expect("serialize")).expect("write v3 manifest"); +} + +#[test] +fn v4_manifest_stamps_per_segment_identity_and_round_trips() { + let root = temp_root("codec-v4-roundtrip"); + let store = store(&root, 0); + let payload: Vec = (0..50_000u32).map(|value| value as u8).collect(); + let manifest = commit_payload(&store, &payload, 4096); + assert_eq!(manifest.version, MANIFEST_VERSION); + assert!( + manifest.segments.iter().all(|segment| segment.codec_identity + == Some(SegmentCodecIdentity::raw(segment.bytes))), + "every v4 segment carries its raw identity" + ); + + // The identity is written explicitly per segment, not inferred at read. + let manifest_json = fs::read_to_string(store.manifest_path(&manifest.payload_digest)) + .expect("read manifest json"); + let value: serde_json::Value = + serde_json::from_str(&manifest_json).expect("parse manifest json"); + for segment in value["segments"].as_array().expect("segment array") { + assert_eq!(segment["codec_identity"]["name"], CODEC_RAW); + assert_eq!(segment["codec_identity"]["version"], CODEC_RAW_VERSION); + assert_eq!(segment["codec_identity"]["class"], "exact"); + assert_eq!(segment["codec_identity"]["decoded_len"], segment["bytes"]); + assert!( + segment["codec_identity"] + .get("calibration_digest") + .is_none() + ); + } + + let loaded = store + .load_manifest(&manifest.payload_digest) + .expect("load manifest"); + assert_eq!(store.assemble(&loaded).expect("assemble"), payload); +} + +#[test] +fn v3_manifest_reads_and_assembles_through_payload_codec() { + let root = temp_root("codec-v3-read"); + let store = store(&root, 0); + let payload = vec![11u8; 8192]; + let manifest = commit_payload(&store, &payload, 4096); + rewrite_on_disk_as_v3(&store, &manifest.payload_digest); + + let loaded = store + .load_manifest(&manifest.payload_digest) + .expect("load v3 manifest"); + assert_eq!( + loaded.version, LEGACY_PAYLOAD_CODEC_MANIFEST_VERSION, + "v3 stays v3: identity is normalized per segment, not rewritten" + ); + assert!( + loaded + .segments + .iter() + .all(|segment| segment.codec_identity.is_none()), + "v3 segments carry no per-segment identity" + ); + assert_eq!(store.assemble(&loaded).expect("assemble v3"), payload); +} + +#[test] +fn stripping_identity_from_a_v4_segment_rejects_everywhere() { + let root = temp_root("codec-v4-stripped"); + let store = store(&root, 0); + let payload = vec![12u8; 8192]; + let manifest = commit_payload(&store, &payload, 4096); + let path = store.manifest_path(&manifest.payload_digest); + let original: serde_json::Value = + serde_json::from_slice(&fs::read(&path).expect("read manifest")).expect("parse manifest"); + assert_eq!(original["version"], MANIFEST_VERSION); + + // On-disk: a v4 manifest with one segment's identity stripped must never + // load — it cannot fall back to the payload codec or to raw. + let mut stripped = original.clone(); + stripped["segments"][1] + .as_object_mut() + .expect("segment object") + .remove("codec_identity"); + fs::write(&path, serde_json::to_vec(&stripped).expect("serialize")).expect("write stripped"); + let error = store + .load_manifest(&manifest.payload_digest) + .expect_err("a stripped v4 segment identity must not load"); + assert!( + error.to_string().contains("per-segment codec identity"), + "load error should name the missing per-segment identity: {error}" + ); + + // In-memory: the same shape must not commit (nothing unassemblable is + // ever persisted) and must not assemble (an unloadable entry is a miss). + let mut memory = manifest.clone(); + memory.segments[0].codec_identity = None; + assert!( + store.commit(&memory).is_err(), + "a v4 commit with a stripped segment identity must be refused" + ); + assert!( + store.assemble(&memory).is_err(), + "a v4 assembly with a stripped segment identity must be refused" + ); + + // Restore the on-disk manifest and strip ALL identities: still rejected, + // proving no aggregate fallback to the payload codec exists. + fs::write(&path, serde_json::to_vec(&original).expect("serialize")).expect("write original"); + let mut all_stripped: serde_json::Value = + serde_json::from_slice(&fs::read(&path).expect("read manifest")).expect("parse"); + for segment in all_stripped["segments"] + .as_array_mut() + .expect("segment array") + { + segment + .as_object_mut() + .expect("segment object") + .remove("codec_identity"); + } + fs::write(&path, serde_json::to_vec(&all_stripped).expect("serialize")) + .expect("write all-stripped"); + assert!( + store.load_manifest(&manifest.payload_digest).is_err(), + "stripping every v4 segment identity must not enable a payload-codec fallback" + ); +} + +#[test] +fn v4_rejects_unsupported_segment_codecs_naming_the_segment() { + let root = temp_root("codec-v4-unsupported"); + let store = store(&root, 0); + let payload = vec![13u8; 8192]; + let manifest = commit_payload(&store, &payload, 4096); + + // In-memory: a single lossy-coded segment makes the whole manifest + // unassemblable here, and the refusal names the offending segment. + let mut tampered = manifest.clone(); + tampered.segments[0].codec_identity = Some(SegmentCodecIdentity { + name: "cachegen".to_string(), + version: 1, + class: CodecClass::Lossy, + decoded_len: tampered.segments[0].bytes, + calibration_digest: Some("blake3:calibration".to_string()), + }); + let error = store + .commit(&tampered) + .expect_err("an unsupported segment codec must not commit"); + let message = error.to_string(); + assert!( + message.contains("segment 0") && message.contains("cachegen"), + "commit error should name the segment and codec: {message}" + ); + drop(store); + + // On-disk: the same shape cannot load on a fresh store. Tamper only the + // segment's identity (payload codec stays raw) so the per-segment gate, + // not the payload gate, is what rejects it. + let reopened = HandoffSegmentStore::open(&root, 0).expect("reopen store"); + let path = reopened.manifest_path(&manifest.payload_digest); + let mut value: serde_json::Value = + serde_json::from_slice(&fs::read(&path).expect("read manifest")).expect("parse manifest"); + value["segments"][0]["codec_identity"] = serde_json::json!({ + "name": "cachegen", + "version": 1, + "class": "lossy", + "decoded_len": manifest.segments[0].bytes, + "calibration_digest": "blake3:calibration", + }); + fs::write(&path, serde_json::to_vec(&value).expect("serialize")).expect("write tampered"); + let error = reopened + .load_manifest(&manifest.payload_digest) + .expect_err("an on-disk unsupported segment codec must not load"); + assert!( + error.to_string().contains("segment"), + "load error should name the segment: {error}" + ); +} + +#[test] +fn v4_rejects_unknown_native_kv_version_before_assembly() { + let root = temp_root("codec-v4-native-future"); + let store = store(&root, 0); + let payload = vec![15u8; 8192]; + let manifest = commit_payload(&store, &payload, 4096); + let mut future = manifest.clone(); + future.payload_kind = "kv-recurrent".to_string(); + future.kv_bytes = future.total_bytes; + future.kv_desc_json = Some("{\"runtime\":\"native\"}".to_string()); + for segment in &mut future.segments { + let mut identity = SegmentCodecIdentity::native_kv_page(segment.bytes); + identity.version = CODEC_NATIVE_KV_PAGE_VERSION + 1; + segment.codec_identity = Some(identity); + } + + let error = store + .assemble(&future) + .expect_err("a future native KV representation must not assemble"); + let message = error.to_string(); + assert!( + message.contains(CODEC_NATIVE_KV_PAGE) && message.contains("unsupported codec"), + "error should name the unsupported native representation: {message}" + ); + assert!( + store.commit(&future).is_err(), + "a future native KV representation must not commit" + ); +} + +#[test] +fn v4_native_kv_segments_must_tile_exactly_to_the_kv_boundary() { + let root = temp_root("codec-v4-native-boundary"); + let store = store(&root, 0); + let payload = vec![16u8; 8192]; + let manifest = commit_payload(&store, &payload, 4096); + let mut mixed = manifest.clone(); + mixed.payload_kind = "kv-recurrent".to_string(); + mixed.kv_bytes = 4096; + mixed.recurrent_bytes = 4096; + mixed.kv_desc_json = Some("{\"runtime\":\"native\"}".to_string()); + mixed.segments[0].codec_identity = Some(SegmentCodecIdentity::native_kv_page(4096)); + assert_eq!( + store.assemble(&mixed).expect("valid mixed manifest"), + payload + ); + + let mut crossing = mixed.clone(); + crossing.kv_bytes = 5000; + crossing.recurrent_bytes = 3192; + let error = store + .assemble(&crossing) + .expect_err("a segment crossing the KV boundary must not assemble"); + assert!(error.to_string().contains("crosses the native KV boundary")); + + let mut native_auxiliary = mixed; + native_auxiliary.segments[1].codec_identity = Some(SegmentCodecIdentity::native_kv_page(4096)); + let error = store + .assemble(&native_auxiliary) + .expect_err("auxiliary bytes cannot claim the native KV representation"); + assert!(error.to_string().contains("representation disagrees")); +} + +#[test] +fn v4_segment_identity_length_mismatch_is_refused() { + let root = temp_root("codec-v4-length"); + let store = store(&root, 0); + let payload = vec![14u8; 8192]; + let manifest = commit_payload(&store, &payload, 4096); + + // An exact identity whose decoded_len disagrees with the stored bytes is + // corruption, not a hint: refuse before any segment is read. + let mut tampered = manifest.clone(); + let bytes = tampered.segments[0].bytes; + tampered.segments[0].codec_identity = Some(SegmentCodecIdentity::raw(bytes + 1)); + assert!( + store.assemble(&tampered).is_err(), + "an exact identity with a decoded_len mismatch must not assemble" + ); + assert!( + store.commit(&tampered).is_err(), + "an exact identity with a decoded_len mismatch must not commit" + ); +} + +#[test] +fn gc_aborts_when_a_manifest_is_unreadable() { + let root = temp_root("gc-corrupt-manifest"); + let store = store(&root, 0); + let manifest = commit_payload(&store, b"referenced bytes", 4096); + let segment = store.segment_path(&manifest.segments[0].digest); + fs::write(store.manifest_path(&manifest.payload_digest), b"not json") + .expect("corrupt manifest"); + + assert!(store.collect_unreferenced_segments().is_err()); + assert!( + segment.is_file(), + "GC deleted state behind an unreadable manifest" + ); +} + +#[test] +fn gc_ignores_atomic_publish_temporary_files() { + let root = temp_root("gc-temp-file"); + let store = store(&root, 0); + let temporary = root.join(SEGMENT_DIR).join(".tmp-test-writer"); + fs::write(&temporary, b"in flight").expect("write temporary file"); + + assert_eq!(store.collect_unreferenced_segments().expect("collect"), 0); + assert!( + temporary.is_file(), + "GC removed an in-flight temporary file" + ); +} diff --git a/crates/skippy-cache/src/l3_remote.rs b/crates/skippy-cache/src/l3_remote.rs new file mode 100644 index 0000000000..3440b7d3ec --- /dev/null +++ b/crates/skippy-cache/src/l3_remote.rs @@ -0,0 +1,479 @@ +//! `skippy-kv/1` — the network backend of the L3 stream contract. +//! +//! A peer serves its `HandoffSegmentStore` over a framed byte stream; +//! clients pull manifests and content-addressed segments by digest. Pulls +//! are idempotent: segments already present locally are skipped, and every +//! fetched segment is digest-verified before it lands in the local store, +//! so a corrupt or malicious peer cannot poison it. Ordering and +//! completeness come from the manifest, exactly as on disk. +//! +//! Transport is any `Read + Write` pair. In the mesh this rides an iroh +//! QUIC stream under the `skippy-kv/1` ALPN (`skippy_protocol::KV_ALPN_V1`) +//! bridged to a local socket, the same pattern the stage transport uses; +//! the harness drives it over plain TCP. +//! +//! **The server side has no authentication**: any process that can reach +//! the port can enumerate and drain the store. Digest verification protects +//! the *client* from a malicious peer, not this server from disclosure. +//! Plain-TCP serving is for the lab harness on trusted networks only — do +//! not wire it to a mesh-reachable listener; mesh exposure goes through the +//! `skippy-kv/1` iroh ALPN with mesh-membership auth. + +use std::{ + io::{BufReader, BufWriter, Read, Write}, + net::{TcpListener, TcpStream}, + time::Duration, +}; + +use anyhow::{Context, Result, bail}; +use serde::{Deserialize, Serialize, de::DeserializeOwned}; + +use crate::l3::{HandoffManifest, HandoffSegmentStore, segment_digest}; + +const MAX_HEADER_BYTES: u64 = 16 * 1024 * 1024; +const MAX_SEGMENT_BYTES: u64 = 256 * 1024 * 1024; +const STREAM_BUFFER_BYTES: usize = 1024 * 1024; +const DEFAULT_IO_TIMEOUT: Duration = Duration::from_secs(600); + +fn configure_stream(stream: &TcpStream, timeout: Duration) -> Result<()> { + stream.set_nodelay(true).ok(); + stream + .set_read_timeout(Some(timeout)) + .context("failed to set skippy-kv read timeout")?; + stream + .set_write_timeout(Some(timeout)) + .context("failed to set skippy-kv write timeout")?; + Ok(()) +} + +mod frame_kind { + pub const GET_MANIFEST: u8 = 1; + pub const MANIFEST: u8 = 2; + pub const GET_SEGMENT: u8 = 3; + pub const SEGMENT: u8 = 4; + pub const LIST_MANIFESTS: u8 = 5; + pub const MANIFEST_LIST: u8 = 6; +} + +#[derive(Serialize, Deserialize)] +struct GetManifestHeader { + /// Manifest key (payload digest); `None` asks for the newest. + key: Option, +} + +#[derive(Serialize, Deserialize)] +struct ManifestHeader { + found: bool, + manifest: Option, +} + +#[derive(Serialize, Deserialize)] +struct GetSegmentHeader { + digest: String, +} + +#[derive(Serialize, Deserialize)] +struct SegmentReplyHeader { + found: bool, + digest: String, +} + +#[derive(Serialize, Deserialize)] +struct ManifestListHeader { + keys: Vec, +} + +/// Statistics for one `fetch_into_store` pull. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct FetchStats { + pub segments_fetched: usize, + pub segments_skipped: usize, + pub bytes_fetched: u64, +} + +/// Serve the store to one connected peer until it disconnects. +pub fn serve_connection(store: &HandoffSegmentStore, stream: TcpStream) -> Result<()> { + serve_connection_with_timeout(store, stream, DEFAULT_IO_TIMEOUT) +} + +fn serve_connection_with_timeout( + store: &HandoffSegmentStore, + stream: TcpStream, + timeout: Duration, +) -> Result<()> { + configure_stream(&stream, timeout)?; + let mut reader = BufReader::with_capacity(STREAM_BUFFER_BYTES, stream.try_clone()?); + let mut writer = BufWriter::with_capacity(STREAM_BUFFER_BYTES, stream); + loop { + let (kind, header, _) = match read_frame(&mut reader) { + Ok(frame) => frame, + // A closed connection between requests is a normal end. + Err(_) => return Ok(()), + }; + match kind { + frame_kind::GET_MANIFEST => { + let request: GetManifestHeader = + serde_json::from_value(header).context("malformed manifest request")?; + let key = match request.key { + Some(key) => Some(key), + None => store.list_manifests()?.into_iter().next(), + }; + let manifest = key.and_then(|key| store.load_manifest(&key).ok()); + write_frame( + &mut writer, + frame_kind::MANIFEST, + &ManifestHeader { + found: manifest.is_some(), + manifest, + }, + &[], + )?; + } + frame_kind::GET_SEGMENT => { + let request: GetSegmentHeader = + serde_json::from_value(header).context("malformed segment request")?; + match store.read_segment(&request.digest) { + Ok(bytes) => write_frame( + &mut writer, + frame_kind::SEGMENT, + &SegmentReplyHeader { + found: true, + digest: request.digest, + }, + &bytes, + )?, + Err(_) => write_frame( + &mut writer, + frame_kind::SEGMENT, + &SegmentReplyHeader { + found: false, + digest: request.digest, + }, + &[], + )?, + } + } + frame_kind::LIST_MANIFESTS => { + write_frame( + &mut writer, + frame_kind::MANIFEST_LIST, + &ManifestListHeader { + keys: store.list_manifests()?, + }, + &[], + )?; + } + other => bail!("unexpected skippy-kv frame kind {other}"), + } + writer.flush().context("failed to flush skippy-kv reply")?; + } +} + +/// Serve the store on a listener for `accept_count` connections +/// (0 = until the process dies). +pub fn serve_store( + store: &HandoffSegmentStore, + listener: &TcpListener, + accept_count: usize, +) -> Result<()> { + serve_store_with_timeout(store, listener, accept_count, DEFAULT_IO_TIMEOUT) +} + +/// Serve the store with an explicit per-connection I/O deadline. +pub fn serve_store_with_timeout( + store: &HandoffSegmentStore, + listener: &TcpListener, + accept_count: usize, + timeout: Duration, +) -> Result<()> { + std::thread::scope(|scope| -> Result<()> { + let mut served = 0usize; + loop { + let (stream, _) = listener.accept().context("skippy-kv accept failed")?; + served += 1; + scope.spawn(move || { + // A malformed, stalled, or disconnected peer is isolated to + // its worker. The listener remains available to other peers. + let _ = serve_connection_with_timeout(store, stream, timeout); + }); + if accept_count != 0 && served >= accept_count { + return Ok(()); + } + } + }) +} + +/// A client connection to a peer's store. +pub struct KvFetchClient { + reader: BufReader, + writer: BufWriter, +} + +impl KvFetchClient { + pub fn connect(peer: &str) -> Result { + Self::connect_with_timeout(peer, DEFAULT_IO_TIMEOUT) + } + + pub fn connect_with_timeout(peer: &str, timeout: Duration) -> Result { + let stream = TcpStream::connect(peer) + .with_context(|| format!("failed to connect to skippy-kv peer {peer}"))?; + configure_stream(&stream, timeout)?; + Ok(Self { + reader: BufReader::with_capacity(STREAM_BUFFER_BYTES, stream.try_clone()?), + writer: BufWriter::with_capacity(STREAM_BUFFER_BYTES, stream), + }) + } + + pub fn list_manifests(&mut self) -> Result> { + write_frame( + &mut self.writer, + frame_kind::LIST_MANIFESTS, + &serde_json::json!({}), + &[], + )?; + self.writer.flush()?; + let (header, _) = + read_frame_expect::(&mut self.reader, frame_kind::MANIFEST_LIST)?; + Ok(header.keys) + } + + pub fn fetch_manifest(&mut self, key: Option<&str>) -> Result { + write_frame( + &mut self.writer, + frame_kind::GET_MANIFEST, + &GetManifestHeader { + key: key.map(str::to_string), + }, + &[], + )?; + self.writer.flush()?; + let (header, _) = + read_frame_expect::(&mut self.reader, frame_kind::MANIFEST)?; + header.manifest.context("peer has no matching manifest") + } + + pub fn fetch_segment(&mut self, digest: &str) -> Result> { + write_frame( + &mut self.writer, + frame_kind::GET_SEGMENT, + &GetSegmentHeader { + digest: digest.to_string(), + }, + &[], + )?; + self.writer.flush()?; + let (header, bytes) = + read_frame_expect::(&mut self.reader, frame_kind::SEGMENT)?; + if !header.found { + bail!("peer does not hold segment {digest}"); + } + if segment_digest(&bytes) != digest { + bail!("segment {digest} from peer failed digest verification"); + } + Ok(bytes) + } + + /// Pull one manifest and every segment the local store is missing, then + /// commit the manifest locally. Content addressing makes this + /// idempotent: re-fetching an already-held manifest transfers nothing. + pub fn fetch_into_store( + &mut self, + key: Option<&str>, + store: &HandoffSegmentStore, + ) -> Result<(HandoffManifest, FetchStats)> { + let manifest = self.fetch_manifest(key)?; + let mut stats = FetchStats::default(); + for segment in &manifest.segments { + if store.has_segment(&segment.digest) { + stats.segments_skipped += 1; + continue; + } + let bytes = self.fetch_segment(&segment.digest)?; + store.put_segment(&bytes)?; + stats.segments_fetched += 1; + stats.bytes_fetched += bytes.len() as u64; + } + store + .commit(&manifest) + .context("failed to commit fetched manifest")?; + Ok((manifest, stats)) + } +} + +fn write_frame( + writer: &mut impl Write, + kind: u8, + header: &impl Serialize, + payload: &[u8], +) -> Result<()> { + let header_bytes = serde_json::to_vec(header).context("failed to encode frame header")?; + if header_bytes.len() as u64 > MAX_HEADER_BYTES { + bail!("frame header of {} bytes exceeds limit", header_bytes.len()); + } + writer.write_all(&[kind])?; + writer.write_all(&(header_bytes.len() as u32).to_le_bytes())?; + writer.write_all(&header_bytes)?; + writer.write_all(&(payload.len() as u64).to_le_bytes())?; + writer.write_all(payload)?; + Ok(()) +} + +fn read_frame(reader: &mut impl Read) -> Result<(u8, serde_json::Value, Vec)> { + let mut kind = [0u8; 1]; + reader + .read_exact(&mut kind) + .context("skippy-kv stream closed")?; + let mut header_len = [0u8; 4]; + reader.read_exact(&mut header_len)?; + let header_len = u32::from_le_bytes(header_len) as u64; + if header_len > MAX_HEADER_BYTES { + bail!("frame header of {header_len} bytes exceeds limit"); + } + let mut header_bytes = vec![0u8; header_len as usize]; + reader.read_exact(&mut header_bytes)?; + let header = serde_json::from_slice(&header_bytes).context("malformed frame header")?; + let mut payload_len = [0u8; 8]; + reader.read_exact(&mut payload_len)?; + let payload_len = u64::from_le_bytes(payload_len); + if payload_len > MAX_SEGMENT_BYTES { + bail!("frame payload of {payload_len} bytes exceeds limit"); + } + let mut payload = vec![0u8; payload_len as usize]; + reader.read_exact(&mut payload)?; + Ok((kind[0], header, payload)) +} + +fn read_frame_expect( + reader: &mut impl Read, + expected_kind: u8, +) -> Result<(T, Vec)> { + let (kind, header, payload) = read_frame(reader)?; + if kind != expected_kind { + bail!("expected skippy-kv frame kind {expected_kind}, got {kind}"); + } + Ok(( + serde_json::from_value(header).context("malformed frame header for expected kind")?, + payload, + )) +} + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + + use super::*; + use crate::l3::{HandoffManifest, HandoffSegmentRef, SegmentCodecIdentity}; + + fn temp_root(name: &str) -> PathBuf { + let root = std::env::temp_dir() + .join("skippy-kv-tests") + .join(format!("{name}-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&root); + root + } + + fn seeded_store(root: &PathBuf, payload: &[u8]) -> (HandoffSegmentStore, HandoffManifest) { + let store = HandoffSegmentStore::open(root, 0).expect("open store"); + let mut manifest = + HandoffManifest::new("blake3:test-identity".to_string(), "full-state".into()); + for (index, chunk) in payload.chunks(1024).enumerate() { + let digest = store.put_segment(chunk).expect("put").digest; + manifest.segments.push(HandoffSegmentRef { + index: index as u32, + offset: (index * 1024) as u64, + bytes: chunk.len() as u64, + digest, + codec_identity: Some(SegmentCodecIdentity::raw(chunk.len() as u64)), + meta_json: None, + }); + } + manifest.total_bytes = payload.len() as u64; + manifest.payload_digest = segment_digest(payload); + store.commit(&manifest).expect("commit"); + (store, manifest) + } + + #[test] + fn fetch_into_store_pulls_verifies_and_is_idempotent() { + let payload: Vec = (0..10_000u32).map(|value| (value % 251) as u8).collect(); + let server_root = temp_root("server"); + let client_root = temp_root("client"); + let (server_store, manifest) = seeded_store(&server_root, &payload); + let client_store = HandoffSegmentStore::open(&client_root, 0).expect("open client"); + + let listener = TcpListener::bind("127.0.0.1:0").expect("bind"); + let address = listener.local_addr().expect("addr").to_string(); + let server = std::thread::spawn(move || { + serve_store(&server_store, &listener, 2).expect("serve"); + }); + + let mut client = KvFetchClient::connect(&address).expect("connect"); + assert_eq!( + client.list_manifests().expect("list"), + vec![manifest.payload_digest.clone()] + ); + let (fetched, stats) = client + .fetch_into_store(None, &client_store) + .expect("first fetch"); + assert_eq!(fetched.payload_digest, manifest.payload_digest); + assert_eq!(stats.segments_fetched, manifest.segments.len()); + assert_eq!(stats.segments_skipped, 0); + assert_eq!(client_store.assemble(&fetched).expect("assemble"), payload); + drop(client); + + // Second pull on a fresh connection: everything present, nothing moves. + let mut client = KvFetchClient::connect(&address).expect("reconnect"); + let (_, stats) = client + .fetch_into_store(Some(&manifest.payload_digest), &client_store) + .expect("second fetch"); + assert_eq!(stats.segments_fetched, 0); + assert_eq!(stats.bytes_fetched, 0); + assert_eq!(stats.segments_skipped, manifest.segments.len()); + drop(client); + server.join().expect("server thread"); + } + + #[test] + fn missing_segments_and_manifests_are_reported() { + let server_root = temp_root("missing"); + let client_root = temp_root("missing-client"); + let payload = vec![7u8; 2048]; + let (server_store, _) = seeded_store(&server_root, &payload); + let _client_store = HandoffSegmentStore::open(&client_root, 0).expect("open client"); + + let listener = TcpListener::bind("127.0.0.1:0").expect("bind"); + let address = listener.local_addr().expect("addr").to_string(); + let server = std::thread::spawn(move || { + serve_store(&server_store, &listener, 1).expect("serve"); + }); + + let mut client = KvFetchClient::connect(&address).expect("connect"); + assert!(client.fetch_manifest(Some("no-such-key")).is_err()); + assert!(client.fetch_segment(&segment_digest(b"absent")).is_err()); + drop(client); + server.join().expect("server thread"); + } + + #[test] + fn stalled_peer_does_not_block_the_listener() { + let server_root = temp_root("concurrent-listener"); + let (server_store, manifest) = seeded_store(&server_root, b"payload"); + let listener = TcpListener::bind("127.0.0.1:0").expect("bind"); + let address = listener.local_addr().expect("addr").to_string(); + let server = std::thread::spawn(move || { + serve_store_with_timeout(&server_store, &listener, 2, Duration::from_secs(2)) + .expect("serve"); + }); + + let stalled = TcpStream::connect(&address).expect("connect stalled peer"); + let mut client = KvFetchClient::connect_with_timeout(&address, Duration::from_secs(2)) + .expect("connect active peer"); + assert_eq!( + client.list_manifests().expect("list while peer is stalled"), + vec![manifest.payload_digest] + ); + drop(client); + drop(stalled); + server.join().expect("server thread"); + } +} diff --git a/crates/skippy-cache/src/lib.rs b/crates/skippy-cache/src/lib.rs index 05ae62427f..6da4340630 100644 --- a/crates/skippy-cache/src/lib.rs +++ b/crates/skippy-cache/src/lib.rs @@ -1,14 +1,42 @@ +pub mod cachegen; pub mod config; +pub mod fsinfo; pub mod identity; +pub mod l2; +pub mod l3; +pub mod l3_remote; +pub mod manager; pub mod payload; +pub mod policy; pub mod radix; pub mod resident; +pub mod source; +pub mod tier; pub use config::{ResidentCacheConfig, SparseCheckpointPolicy}; pub use identity::{ - NATIVE_KV_DTYPE, NATIVE_KV_RUNTIME_ABI_VERSION, PrefixIdentity, activation_page_id, - prefix_hash, prefix_hash_with_namespace, prefix_identity, prefix_identity_with_namespace, - prefix_namespace_hash, + ExactStateIdentityParams, NATIVE_KV_DTYPE, NATIVE_KV_RUNTIME_ABI_VERSION, PrefixIdentity, + activation_page_id, exact_state_identity, exact_state_identity_for_stage, + numerical_model_identity_for_stage, prefix_hash, prefix_hash_with_namespace, prefix_identity, + prefix_identity_with_namespace, prefix_namespace_hash, +}; +pub use l2::{ + ExactStatePayloadMirror, L2Eviction, L2Hit, L2InsertRefusal, L2Origin, L2Peek, L2Stats, L2Tier, + l2_cache_key, +}; +pub use l3::{ + CODEC_NATIVE_KV_PAGE, CODEC_NATIVE_KV_PAGE_VERSION, CODEC_RAW, CODEC_RAW_VERSION, CodecClass, + GeometryBlock, GeometryKind, HandoffManifest, HandoffSegmentRef, HandoffSegmentStore, + LEGACY_MANIFEST_VERSION, LEGACY_PAYLOAD_CODEC_MANIFEST_VERSION, MANIFEST_VERSION, ManifestPin, + PayloadCodec, PayloadGeometry, Reservation, SegmentCodecIdentity, SegmentHold, SegmentPut, + StoreLimits, StoreReconciliation, StoreUsage, StoredSegment, WriteRefusal, segment_digest, +}; +pub use l3_remote::{ + FetchStats, KvFetchClient, serve_connection, serve_store, serve_store_with_timeout, +}; +pub use manager::{ + L3ActivitySnapshot, L3CacheManager, L3EffectiveState, L3EffectiveStatus, L3InventoryEntry, + L3StateReason, L3StateTransition, }; pub use payload::{ CacheBlobStore, CacheBytes, CacheBytesReconstructStats, CacheDedupeStats, ExactStatePayload, @@ -21,6 +49,9 @@ pub use resident::{ ResidentActivationCache, ResidentActivationLookup, ResidentActivationRecordOutcome, ResidentActivationStats, }; +pub use source::{ManifestSource, SegmentSource}; + +pub use tier::{L3Fill, L3Location, L3Status, L3Tier, l3_namespace_key, l3_prefix_key}; /// llama.cpp's hard sequence-id capacity for one context. pub const LLAMA_MAX_SEQ: i32 = 256; diff --git a/crates/skippy-cache/src/manager.rs b/crates/skippy-cache/src/manager.rs new file mode 100644 index 0000000000..5a6a844695 --- /dev/null +++ b/crates/skippy-cache/src/manager.rs @@ -0,0 +1,717 @@ +//! Node-scoped ownership for the durable L3 cache root. + +use std::{ + collections::{BTreeMap, VecDeque}, + fs, + path::{Path, PathBuf}, + sync::{ + Arc, LazyLock, Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard, Weak, + atomic::{AtomicU64, Ordering}, + }, +}; + +use anyhow::{Context, Result, bail}; +use serde::Serialize; + +use crate::{ + l3::{HandoffSegmentStore, StoreLimits, StoreReconciliation, StoreUsage, WriteRefusal}, + tier::L3Tier, +}; + +static ROOT_MANAGERS: LazyLock>>> = + LazyLock::new(|| Mutex::new(BTreeMap::new())); +const MAX_PENDING_STATE_TRANSITIONS: usize = 64; + +/// What every stage attached to the node's L3 root has done since open. +#[derive(Debug, Default)] +pub(crate) struct L3Activity { + pub(crate) fills: AtomicU64, + pub(crate) hits: AtomicU64, + pub(crate) misses: AtomicU64, + pub(crate) writes: AtomicU64, + pub(crate) geometry_rejected: AtomicU64, + pub(crate) bytes_read: AtomicU64, + pub(crate) bytes_written: AtomicU64, + last_error: Mutex>, +} + +impl L3Activity { + pub(crate) fn record_error(&self, error: &anyhow::Error) { + *self + .last_error + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(format!("{error:#}")); + } + + fn snapshot(&self, usage: Option<&StoreUsage>) -> L3ActivitySnapshot { + L3ActivitySnapshot { + fills: self.fills.load(Ordering::Relaxed), + hits: self.hits.load(Ordering::Relaxed), + misses: self.misses.load(Ordering::Relaxed), + writes: self.writes.load(Ordering::Relaxed), + evictions: usage.map_or(0, |usage| usage.evicted_manifests), + corrupt_entries: usage.map_or(0, |usage| usage.quarantined_objects), + bytes_read: self.bytes_read.load(Ordering::Relaxed), + bytes_written: self.bytes_written.load(Ordering::Relaxed), + geometry_rejected: self.geometry_rejected.load(Ordering::Relaxed), + last_error: self + .last_error + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(), + } + } +} + +/// Point-in-time activity across every stage attached to one root manager. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct L3ActivitySnapshot { + pub fills: u64, + pub hits: u64, + pub misses: u64, + pub writes: u64, + pub evictions: u64, + pub corrupt_entries: u64, + pub bytes_read: u64, + pub bytes_written: u64, + pub geometry_rejected: u64, + pub last_error: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct L3InventoryEntry { + pub model_identity: String, + pub state_identity: String, + pub payload_kind: String, + pub token_count: u64, + pub total_bytes: u64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum L3EffectiveState { + Active, + ReadOnlyLowSpace, + Degraded, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum L3StateReason { + ReadOnlyLowSpace, + InsufficientSpace, + StorageError, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +pub struct L3EffectiveStatus { + pub state: L3EffectiveState, + pub reason: Option, +} + +impl Default for L3EffectiveStatus { + fn default() -> Self { + Self { + state: L3EffectiveState::Active, + reason: None, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +pub struct L3StateTransition { + pub previous: L3EffectiveStatus, + pub current: L3EffectiveStatus, +} + +#[derive(Debug)] +struct L3ManagerInner { + store: Arc, + activity: Arc, + fill_claims: Arc>>, + record_claims: Mutex>>>>, + effective: Mutex, + transitions: Mutex>, + operations: RwLock<()>, + reconciliation: StoreReconciliation, +} + +/// The single physical owner of a node-local L3 root. +/// +/// Clones are cheap stage handles into the same reservation, pin, lifecycle, +/// activity, and filesystem-lock domain. +#[derive(Clone, Debug)] +pub struct L3CacheManager { + inner: Arc, +} + +impl L3CacheManager { + /// Acquire the manager for `root`, reusing the live node owner when one + /// exists. A second process is rejected by the store's root lock. + pub fn acquire(root: impl AsRef, limits: StoreLimits) -> Result { + let root = canonical_cache_root(root.as_ref())?; + let mut managers = ROOT_MANAGERS + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + // An entry whose strong count has reached zero is not necessarily gone: + // the owning thread decrements the count before dropping the inner + // value, and the store's root `flock` is only released by that drop. + let expiring_owner = managers + .get(&root) + .is_some_and(|manager| manager.strong_count() == 0); + managers.retain(|_, manager| manager.strong_count() > 0); + if let Some(inner) = managers.get(&root).and_then(Weak::upgrade) { + if inner.store.limits() != limits { + bail!( + "cache root {} is already open with different limits", + root.display() + ); + } + return Ok(Self { inner }); + } + + // Taking the root lock while the previous in-process owner is still + // unwinding would report the root as owned by another manager, which + // is a false answer: no other process holds it. Give that drop a + // bounded moment to finish rather than failing the caller. + let store = Arc::new(open_store_for_acquire(&root, limits, expiring_owner)?); + let reconciliation = store.reconcile_startup()?; + let inner = Arc::new(L3ManagerInner { + store, + activity: Arc::new(L3Activity::default()), + fill_claims: Arc::new(Mutex::new(std::collections::BTreeSet::new())), + record_claims: Mutex::new(BTreeMap::new()), + effective: Mutex::new(L3EffectiveStatus::default()), + transitions: Mutex::new(VecDeque::new()), + operations: RwLock::new(()), + reconciliation, + }); + managers.insert(root, Arc::downgrade(&inner)); + Ok(Self { inner }) + } + + pub fn tier(&self, state_identity: String, segment_bytes: usize) -> L3Tier { + self.tier_for_model(state_identity.clone(), state_identity, segment_bytes) + } + + pub fn tier_for_model( + &self, + model_identity: String, + state_identity: String, + segment_bytes: usize, + ) -> L3Tier { + L3Tier::from_manager(self.clone(), model_identity, state_identity, segment_bytes) + } + + pub fn root(&self) -> &Path { + self.inner.store.root() + } + + pub fn limits(&self) -> StoreLimits { + self.inner.store.limits() + } + + pub fn update_limits(&self, limits: StoreLimits) -> Result { + let _lifecycle = self.lifecycle_guard(); + self.inner.store.update_limits(limits) + } + + pub fn reconciliation(&self) -> StoreReconciliation { + self.inner.reconciliation + } + + pub fn usage(&self) -> Result { + self.inner.store.usage() + } + + pub fn activity(&self) -> Result { + let usage = self.usage()?; + Ok(self.inner.activity.snapshot(Some(&usage))) + } + + pub fn inventory(&self) -> Result> { + let mut inventory = self + .inner + .store + .list_manifests()? + .into_iter() + .filter_map(|key| self.inner.store.load_manifest(&key).ok()) + .map(|manifest| L3InventoryEntry { + model_identity: manifest.model_identity, + state_identity: manifest.state_identity, + payload_kind: manifest.payload_kind, + token_count: manifest.token_count, + total_bytes: manifest.total_bytes, + }) + .collect::>(); + inventory.sort_by(|left, right| { + left.model_identity + .cmp(&right.model_identity) + .then_with(|| left.state_identity.cmp(&right.state_identity)) + .then_with(|| right.token_count.cmp(&left.token_count)) + }); + Ok(inventory) + } + + pub fn effective_status(&self) -> L3EffectiveStatus { + *self + .inner + .effective + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + pub fn take_state_transitions(&self) -> Vec { + self.inner + .transitions + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .drain(..) + .collect() + } + + pub fn prune_to(&self, target_bytes: u64) -> Result { + let _lifecycle = self.lifecycle_guard(); + self.inner.store.prune_to(target_bytes) + } + + pub fn prune_model_to(&self, model_identity: &str, target_bytes: u64) -> Result { + let _lifecycle = self.lifecycle_guard(); + self.inner + .store + .prune_model_to(model_identity, target_bytes) + } + + pub fn clear(&self) -> Result { + let _lifecycle = self.lifecycle_guard(); + self.inner.store.clear() + } + + pub fn clear_model(&self, model_identity: &str) -> Result { + let _lifecycle = self.lifecycle_guard(); + self.inner.store.clear_model(model_identity) + } + + pub fn shares_root_with(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.inner, &other.inner) + } + + /// Manifest keys currently being filled from this root. Keeping claims + /// here makes single-flight node-wide instead of duplicating physical + /// reads when placement-equivalent stages miss at the same time. + pub fn fill_claims(&self) -> Arc>> { + self.inner.fill_claims.clone() + } + + /// Record claims are shared only by stages with the same full numerical + /// state identity. This prevents duplicate exports and temporary writes + /// across placement replicas without suppressing a different payload or + /// stage layout that happens to use the same radix page id. + pub fn record_claims( + &self, + state_identity: &str, + ) -> Arc>> { + self.inner + .record_claims + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .entry(state_identity.to_string()) + .or_default() + .clone() + } + + pub(crate) fn store(&self) -> &HandoffSegmentStore { + &self.inner.store + } + + pub(crate) fn activity_counters(&self) -> &L3Activity { + &self.inner.activity + } + + pub(crate) fn activity_snapshot(&self) -> L3ActivitySnapshot { + let usage = self.usage().ok(); + self.inner.activity.snapshot(usage.as_ref()) + } + + pub(crate) fn record_write_refusal(&self, refusal: WriteRefusal) { + let next = match refusal { + WriteRefusal::SkippedOversize => return, + WriteRefusal::ReadOnlyLowSpace => L3EffectiveStatus { + state: L3EffectiveState::ReadOnlyLowSpace, + reason: Some(L3StateReason::ReadOnlyLowSpace), + }, + WriteRefusal::InsufficientSpace => L3EffectiveStatus { + state: L3EffectiveState::Degraded, + reason: Some(L3StateReason::InsufficientSpace), + }, + }; + self.transition_to(next); + } + + pub(crate) fn record_successful_write(&self) { + let current = self.effective_status(); + if matches!( + current.reason, + Some(L3StateReason::ReadOnlyLowSpace | L3StateReason::InsufficientSpace) + ) { + self.transition_to(L3EffectiveStatus::default()); + } + } + + pub(crate) fn record_storage_error(&self) { + self.transition_to(L3EffectiveStatus { + state: L3EffectiveState::Degraded, + reason: Some(L3StateReason::StorageError), + }); + } + + fn transition_to(&self, next: L3EffectiveStatus) { + let mut current = self + .inner + .effective + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if *current == next { + return; + } + let previous = *current; + *current = next; + drop(current); + let mut transitions = self + .inner + .transitions + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if transitions.len() == MAX_PENDING_STATE_TRANSITIONS { + transitions.pop_front(); + } + transitions.push_back(L3StateTransition { + previous, + current: next, + }); + } + + pub(crate) fn operation_guard(&self) -> RwLockReadGuard<'_, ()> { + self.inner + .operations + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + pub(crate) fn try_operation_guard(&self) -> Option> { + match self.inner.operations.try_read() { + Ok(guard) => Some(guard), + Err(std::sync::TryLockError::Poisoned(error)) => Some(error.into_inner()), + Err(std::sync::TryLockError::WouldBlock) => None, + } + } + + fn lifecycle_guard(&self) -> RwLockWriteGuard<'_, ()> { + self.inner + .operations + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } +} + +/// Open the store, retrying briefly when the previous in-process owner for +/// this root is still being dropped. +/// +/// Only the handoff window is retried. A root genuinely held by another +/// process fails with the same error it always did, one short delay later. +fn open_store_for_acquire( + root: &Path, + limits: StoreLimits, + expiring_owner: bool, +) -> Result { + const HANDOFF_ATTEMPTS: u32 = 20; + const HANDOFF_BACKOFF: std::time::Duration = std::time::Duration::from_millis(5); + + let attempts = if expiring_owner { HANDOFF_ATTEMPTS } else { 1 }; + let mut last = None; + for attempt in 0..attempts { + match HandoffSegmentStore::open_unreconciled_with_limits(root, limits) { + Ok(store) => return Ok(store), + Err(error) => { + last = Some(error); + if attempt + 1 < attempts { + std::thread::sleep(HANDOFF_BACKOFF); + } + } + } + } + Err(last.expect("at least one attempt was made")) +} + +fn canonical_cache_root(root: &Path) -> Result { + if !root.is_absolute() { + bail!("cache root must be absolute: {}", root.display()); + } + crate::fsinfo::refuse_symlink(root)?; + fs::create_dir_all(root) + .with_context(|| format!("failed to create cache root {}", root.display()))?; + fs::canonicalize(root) + .with_context(|| format!("failed to resolve cache root {}", root.display())) +} + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Barrier}; + + use super::*; + use crate::ExactStatePayload; + + fn temp_root(name: &str) -> PathBuf { + static NEXT: AtomicU64 = AtomicU64::new(0); + let root = std::env::temp_dir() + .join("skippy-l3-manager-tests") + .join(format!( + "{name}-{}-{}", + std::process::id(), + NEXT.fetch_add(1, Ordering::Relaxed) + )); + let _ = fs::remove_dir_all(&root); + root + } + + #[test] + fn one_live_manager_owns_each_root() { + let root = temp_root("shared-owner"); + let limits = StoreLimits::new(1_000_000, 0); + let first = L3CacheManager::acquire(&root, limits).expect("first manager"); + let second = L3CacheManager::acquire(&root, limits).expect("shared manager"); + + assert!(first.shares_root_with(&second)); + assert_eq!(first.root(), second.root()); + assert!( + L3CacheManager::acquire(&root, StoreLimits::new(2_000_000, 0)).is_err(), + "one root accepted contradictory budgets" + ); + } + + #[test] + fn concurrent_stages_cannot_double_reserve_the_budget() { + let root = temp_root("atomic-reservation"); + let manager = L3CacheManager::acquire(&root, StoreLimits::new(10_000, 0)).unwrap(); + let barrier = Arc::new(Barrier::new(3)); + let mut tasks = Vec::new(); + for _ in 0..2 { + let manager = manager.clone(); + let barrier = barrier.clone(); + tasks.push(std::thread::spawn(move || { + let reservation = manager.store().reserve(8_000).unwrap(); + let admitted = reservation.is_ok(); + barrier.wait(); + admitted + })); + } + barrier.wait(); + let admitted = tasks + .into_iter() + .map(|task| task.join().unwrap()) + .filter(|admitted| *admitted) + .count(); + + assert_eq!(admitted, 1, "two stages reserved the same bytes"); + assert_eq!(manager.usage().unwrap().reserved_inflight_bytes, 0); + } + + #[test] + fn activity_and_store_accounting_are_node_wide() { + let root = temp_root("node-status"); + let manager = L3CacheManager::acquire(&root, StoreLimits::new(1_000_000, 0)).unwrap(); + let stage_a = manager.tier("state-a".to_string(), 4); + let stage_b = manager.tier("state-b".to_string(), 4); + + stage_a + .spill( + "namespace-a", + &[1, 2, 3], + &ExactStatePayload::full_state(b"stage-a-state".to_vec()), + None, + None, + ) + .unwrap(); + + assert_eq!(stage_b.activity().writes, 1); + let stage_a_usage = stage_a.status().unwrap().usage; + let stage_b_usage = stage_b.status().unwrap().usage; + // Free filesystem capacity is sampled per status call and may change + // between these reads; compare only the manager-owned accounting. + assert_eq!(stage_a_usage.budget_bytes, stage_b_usage.budget_bytes); + assert_eq!(stage_a_usage.used_bytes, stage_b_usage.used_bytes); + assert_eq!( + stage_a_usage.reserved_inflight_bytes, + stage_b_usage.reserved_inflight_bytes + ); + assert_eq!( + stage_a_usage.minimum_free_bytes, + stage_b_usage.minimum_free_bytes + ); + assert_eq!(stage_a_usage.manifests, stage_b_usage.manifests); + assert_eq!(stage_a_usage.unique_segments, stage_b_usage.unique_segments); + assert_eq!( + stage_a_usage.evicted_manifests, + stage_b_usage.evicted_manifests + ); + assert_eq!( + stage_a_usage.quarantined_objects, + stage_b_usage.quarantined_objects + ); + assert_eq!(stage_a.status().unwrap().restorable_manifests, 1); + assert_eq!(stage_b.status().unwrap().restorable_manifests, 0); + } + + #[test] + fn one_stage_cannot_evict_another_stages_pinned_manifest() { + let root = temp_root("cross-stage-pin"); + let manager = L3CacheManager::acquire(&root, StoreLimits::new(1_000_000, 0)).unwrap(); + let stage_a = manager.tier("state-a".to_string(), 4); + let stage_b = manager.tier("state-b".to_string(), 4); + let key = stage_a + .spill( + "namespace-a", + &[1, 2, 3], + &ExactStatePayload::full_state(b"stage-a-state".to_vec()), + None, + None, + ) + .unwrap(); + + let pin = stage_a.store().pin(&key); + assert_eq!(stage_b.manager().prune_to(0).unwrap(), 0); + assert!(stage_a.store().load_manifest(&key).is_ok()); + + drop(pin); + assert!(stage_b.manager().prune_to(0).unwrap() > 0); + assert!(stage_a.store().load_manifest(&key).is_err()); + } + + #[test] + fn root_lock_rejects_an_independent_store_owner() { + let root = temp_root("root-lock"); + let limits = StoreLimits::new(1_000_000, 0); + let _manager = L3CacheManager::acquire(&root, limits).unwrap(); + + let error = HandoffSegmentStore::open_with_limits(&root, limits) + .expect_err("a second physical root owner acquired the lock"); + assert!(error.to_string().contains("already owned")); + } + + #[test] + fn startup_reconciles_temps_corrupt_manifests_links_and_orphans() { + let root = temp_root("reconcile"); + let limits = StoreLimits::new(1_000_000, 0); + let manager = L3CacheManager::acquire(&root, limits).unwrap(); + fs::write(root.join("segments/.tmp-dead"), b"partial").unwrap(); + fs::write(root.join("segments/orphan.seg"), b"orphan").unwrap(); + fs::write(root.join("manifests/corrupt.json"), b"{").unwrap(); + let mut incompatible = + crate::HandoffManifest::new("state".to_string(), "full-state".to_string()); + incompatible.version = crate::MANIFEST_VERSION + 1; + incompatible.payload_digest = "incompatible".to_string(); + fs::write( + root.join("manifests/incompatible.json"), + serde_json::to_vec(&incompatible).unwrap(), + ) + .unwrap(); + let mut incomplete = + crate::HandoffManifest::new("state".to_string(), "full-state".to_string()); + incomplete.payload_digest = "incomplete".to_string(); + incomplete.total_bytes = 4; + incomplete.segments.push(crate::HandoffSegmentRef { + index: 0, + offset: 0, + bytes: 4, + digest: "missing-segment".to_string(), + codec_identity: Some(crate::SegmentCodecIdentity::raw(4)), + meta_json: None, + }); + fs::write( + root.join("manifests/incomplete.json"), + serde_json::to_vec(&incomplete).unwrap(), + ) + .unwrap(); + let namespace = root.join("prefixes/namespace"); + fs::create_dir_all(&namespace).unwrap(); + fs::write(namespace.join("000000000001-prefix.key"), b"missing").unwrap(); + drop(manager); + + let reopened = L3CacheManager::acquire(&root, limits).unwrap(); + let report = reopened.reconciliation(); + assert_eq!(report.removed_temporary_files, 1); + assert_eq!(report.quarantined_manifests, 3); + assert_eq!(report.removed_prefix_links, 1); + assert_eq!(report.removed_orphan_bytes, 6); + assert!(!root.join("segments/orphan.seg").exists()); + assert!(root.join("quarantine/corrupt.json").exists()); + assert!(root.join("quarantine/incompatible.json").exists()); + assert!(root.join("quarantine/incomplete.json").exists()); + } + + #[test] + fn low_space_state_transitions_once_and_recovers_after_a_write() { + let root = temp_root("state-transitions"); + let manager = L3CacheManager::acquire(&root, StoreLimits::new(1_000_000, 0)).unwrap(); + + manager.record_write_refusal(WriteRefusal::ReadOnlyLowSpace); + manager.record_write_refusal(WriteRefusal::ReadOnlyLowSpace); + assert_eq!( + manager.effective_status(), + L3EffectiveStatus { + state: L3EffectiveState::ReadOnlyLowSpace, + reason: Some(L3StateReason::ReadOnlyLowSpace), + } + ); + assert_eq!(manager.take_state_transitions().len(), 1); + + manager.record_successful_write(); + assert_eq!(manager.effective_status(), L3EffectiveStatus::default()); + let recovery = manager.take_state_transitions(); + assert_eq!(recovery.len(), 1); + assert_eq!(recovery[0].current, L3EffectiveStatus::default()); + } + + #[test] + fn request_path_can_skip_cache_while_a_lifecycle_operation_drains() { + let root = temp_root("lifecycle-fallback"); + let manager = L3CacheManager::acquire(&root, StoreLimits::new(1_000_000, 0)).unwrap(); + + let lifecycle = manager.lifecycle_guard(); + assert!(manager.try_operation_guard().is_none()); + drop(lifecycle); + assert!(manager.try_operation_guard().is_some()); + } + + #[test] + fn model_clear_matches_internal_identity_exactly() { + let root = temp_root("model-clear"); + let manager = L3CacheManager::acquire(&root, StoreLimits::new(1_000_000, 0)).unwrap(); + let model_a = manager.tier_for_model("model-a".to_string(), "state-a".to_string(), 4); + let model_b = manager.tier_for_model("model-b".to_string(), "state-b".to_string(), 4); + let key_a = model_a + .spill( + "namespace-a", + &[1, 2, 3], + &ExactStatePayload::full_state(b"stage-a-state".to_vec()), + None, + None, + ) + .unwrap(); + let key_b = model_b + .spill( + "namespace-b", + &[1, 2, 3], + &ExactStatePayload::full_state(b"stage-b-state".to_vec()), + None, + None, + ) + .unwrap(); + + assert!(manager.clear_model("model-a").unwrap() > 0); + assert!(manager.store().load_manifest(&key_a).is_err()); + assert!(manager.store().load_manifest(&key_b).is_ok()); + assert_eq!(model_a.status().unwrap().restorable_manifests, 0); + assert_eq!(model_b.status().unwrap().restorable_manifests, 1); + } +} diff --git a/crates/skippy-cache/src/payload/blob_store.rs b/crates/skippy-cache/src/payload/blob_store.rs index 6222071bbd..35044494b9 100644 --- a/crates/skippy-cache/src/payload/blob_store.rs +++ b/crates/skippy-cache/src/payload/blob_store.rs @@ -593,7 +593,7 @@ mod tests { } assert_eq!( blobs.physical_bytes(), - expected.values().map(|(bytes, _)| bytes).sum(), + expected.values().map(|(bytes, _)| bytes).sum::(), "seed={seed:#x} step={step}" ); } @@ -797,12 +797,15 @@ mod tests { assert_storage_accounting(&blobs, &owners, seed, step); assert!( blobs.physical_bytes() - >= expected_blocks.keys().map(|block| block.len() as u64).sum(), + >= expected_blocks + .keys() + .map(|block| block.len() as u64) + .sum::(), "seed={seed:#x} step={step}" ); assert_eq!( blobs.logical_ref_count(), - expected_blocks.values().sum(), + expected_blocks.values().sum::(), "seed={seed:#x} step={step}" ); } diff --git a/crates/skippy-cache/src/payload/bytes.rs b/crates/skippy-cache/src/payload/bytes.rs index 063684bc11..a3201c1b32 100644 --- a/crates/skippy-cache/src/payload/bytes.rs +++ b/crates/skippy-cache/src/payload/bytes.rs @@ -25,7 +25,7 @@ pub(super) enum CacheBytesRepr { } #[derive(Debug, Clone)] -pub(super) struct CacheBlockRef { +pub(crate) struct CacheBlockRef { pub(super) hash: String, /// Shared indirection lets eviction materialize a surviving deduped block /// before releasing its former contiguous backing allocation. @@ -71,6 +71,45 @@ impl CacheBytes { } } + /// Crate-internal: build a block-backed view over shared immutable + /// storages without copying bytes. Each item is `(hash, storage, range)`; + /// `hash` is bookkeeping identity for the block (the L2 tier uses the + /// segment digest). When exactly one item covers its whole storage the + /// view borrows it contiguously; otherwise reads reconstruct in block + /// order. + pub(crate) fn from_shared_blocks( + len: u64, + blocks: impl IntoIterator>, Range)>, + ) -> Self { + let refs: Vec = blocks + .into_iter() + .map(|(hash, storage, range)| { + CacheBlockRef::new( + hash, + Arc::new(RwLock::new(CacheBlockBytes::new(storage, range))), + ) + }) + .collect(); + let contiguous = match refs.as_slice() { + [single] => { + let bytes = single + .bytes + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + (bytes.range.start == 0 && bytes.range.end == bytes.storage.len()) + .then(|| Arc::clone(&bytes.storage)) + } + _ => None, + }; + Self { + len, + repr: CacheBytesRepr::Blocks { + blocks: refs.into(), + contiguous, + }, + } + } + pub(super) fn blocks( len: u64, blocks: Vec, diff --git a/crates/skippy-cache/src/policy/accounting.rs b/crates/skippy-cache/src/policy/accounting.rs new file mode 100644 index 0000000000..0c79d0f59c --- /dev/null +++ b/crates/skippy-cache/src/policy/accounting.rs @@ -0,0 +1,121 @@ +//! Shared-segment accounting: fractional credit so physical bytes are never +//! double-counted across entries (#1650 first slice). + +use std::collections::{BTreeMap, BTreeSet}; + +use crate::policy::EntryKey; + +/// Opaque shared-segment identity. +pub type SegmentId = u64; + +/// Ledger of shared physical segments and their referencing entries. +/// +/// A segment referenced by N entries contributes `size / N` bytes to each +/// entry's effective footprint; releasing an entry drops its reference, and a +/// segment with no references disappears entirely. +#[derive(Debug, Default, Clone)] +pub struct SharedSegmentLedger { + segments: BTreeMap, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct SegmentRecord { + pub size: u64, + pub references: BTreeSet, +} + +impl SharedSegmentLedger { + /// Register a segment of `size` bytes referenced by `entry`. Adding the + /// same reference twice is a no-op. + pub fn add(&mut self, segment: SegmentId, size: u64, entry: EntryKey) { + let record = self.segments.entry(segment).or_insert(SegmentRecord { + size, + references: BTreeSet::new(), + }); + record.references.insert(entry); + } + + /// Release `entry`'s reference to `segment`; drop the segment when the + /// last reference goes. + pub fn release(&mut self, segments: &[SegmentId], entry: EntryKey) { + for segment in segments { + if let Some(record) = self.segments.get_mut(segment) { + record.references.remove(&entry); + if record.references.is_empty() { + self.segments.remove(segment); + } + } + } + } + + /// Fractional bytes charged to `entry` for its shared segments. + pub fn fractional_bytes(&self, entry: EntryKey, segments: &[SegmentId]) -> f64 { + segments + .iter() + .filter_map(|s| self.segments.get(s)) + .filter(|r| r.references.contains(&entry)) + .map(|r| r.size as f64 / r.references.len() as f64) + .sum() + } + + /// Read access for marginal-release computation. + pub fn segment_record(&self, segment: SegmentId) -> Option<&SegmentRecord> { + self.segments.get(&segment) + } + + /// Physical bytes that would actually be released if `entry` were + /// removed: its exclusive bytes are caller-side, so this covers only + /// segments where this entry holds the last reference — the marginal + /// physical release, not the fractional credit. + pub fn marginal_physical_bytes(&self, entry: EntryKey, segments: &[SegmentId]) -> u64 { + segments + .iter() + .filter_map(|s| self.segments.get(s)) + .filter(|r| r.references.len() == 1 && r.references.contains(&entry)) + .map(|r| r.size) + .sum() + } + + /// Total physical bytes held in the ledger. + pub fn total_bytes(&self) -> u64 { + self.segments.values().map(|r| r.size).sum() + } + + pub fn segment_count(&self) -> usize { + self.segments.len() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn fractional_credit_splits_bytes_without_double_counting() { + let mut ledger = SharedSegmentLedger::default(); + ledger.add(1, 300, 10); + ledger.add(1, 300, 11); + ledger.add(1, 300, 12); + assert_eq!(ledger.total_bytes(), 300); + for entry in [10u64, 11, 12] { + assert_eq!(ledger.fractional_bytes(entry, &[1]), 100.0); + } + } + + #[test] + fn releasing_last_reference_drops_segment() { + let mut ledger = SharedSegmentLedger::default(); + ledger.add(7, 128, 1); + ledger.release(&[7], 1); + assert_eq!(ledger.total_bytes(), 0); + assert_eq!(ledger.segment_count(), 0); + } + + #[test] + fn duplicate_reference_is_idempotent() { + let mut ledger = SharedSegmentLedger::default(); + ledger.add(7, 128, 1); + ledger.add(7, 128, 1); + assert_eq!(ledger.fractional_bytes(1, &[7]), 128.0); + } +} diff --git a/crates/skippy-cache/src/policy/admission.rs b/crates/skippy-cache/src/policy/admission.rs new file mode 100644 index 0000000000..a30f5326a5 --- /dev/null +++ b/crates/skippy-cache/src/policy/admission.rs @@ -0,0 +1,295 @@ +//! Admission, probation, promotion, and eviction selection (#1650 first +//! slice). Pure decision logic over `BenefitPolicy` state. + +use serde::Serialize; + +use crate::policy::{ + BenefitPolicy, CostSample, EntryKey, EvictionVerdict, GhostStats, PolicyEntry, SegmentId, +}; + +/// What the policy decided to do with a candidate or admitted entry. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +pub enum AdmissionDecisionKind { + AdmitProbation, + AdmitPersist, + Promote, + Reject, +} + +impl AdmissionDecision { + /// A rejected offer whose measured cost was invalid. Constructed without + /// touching policy state. + pub fn rejected_invalid_cost() -> Self { + Self::rejected("invalid-cost-sample") + } + + /// A structurally rejected offer. Constructed without touching policy + /// state. + pub fn rejected(reason: &str) -> Self { + Self { + kind: AdmissionDecisionKind::Reject, + verdict: AdmissionVerdict::Reject, + reasons: vec![reason.into()], + probation_cap_repair: crate::policy::CapRepair::satisfied(), + } + } +} + +/// The verdict plus opaque reasons for logging. +#[derive(Debug, Clone, PartialEq)] +pub struct AdmissionDecision { + pub kind: AdmissionDecisionKind, + pub verdict: AdmissionVerdict, + pub reasons: Vec, + /// Cap-repair plan after this admission (pins never selected; a + /// `Deferred` variant reports the shortfall). The policy does not + /// remove victims itself: committed removal stays with + /// `BenefitPolicy::remove`. + pub probation_cap_repair: crate::policy::CapRepair, +} + +/// Lifecycle state of a policy entry. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PolicyEntryState { + /// Resident in memory only; not persisted to disk. + Probation, + /// Persist-eligible: survived the hit threshold. + Admitted, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AdmissionVerdict { + Admit, + Reject, +} + +pub(crate) fn consider( + policy: &mut BenefitPolicy, + key: EntryKey, + exclusive_bytes: u64, + shared: Vec<(SegmentId, u64)>, + cost: CostSample, +) -> AdmissionDecision { + let reasons: Vec; + let kind; + + let mut seen = std::collections::BTreeSet::new(); + let duplicate_segment = shared.iter().any(|(s, _)| !seen.insert(*s)); + // Prevalidate every known-size conflict across the whole list before any + // mutation (ghost removal or ledger reference) so a rejected offer + // leaves policy state untouched. + let size_conflict = shared.iter().any(|(segment, size)| { + policy + .segments + .segment_record(*segment) + .is_some_and(|record| !record.references.contains(&key) && record.size != *size) + }); + if !cost.is_valid() { + // Invalid measured costs (NaN/infinite/negative) are rejected before + // any mutation: they must never enter entry state. + kind = AdmissionDecisionKind::Reject; + reasons = vec!["invalid-cost-sample".into()]; + } else if duplicate_segment { + kind = AdmissionDecisionKind::Reject; + reasons = vec!["duplicate-segment-reference".into()]; + } else if size_conflict { + kind = AdmissionDecisionKind::Reject; + reasons = vec!["segment-size-conflict".into()]; + } else if cost.net_benefit() <= 0.0 { + kind = AdmissionDecisionKind::Reject; + reasons = vec!["no-net-benefit".into()]; + } else if exclusive_bytes == 0 && shared.is_empty() { + // Zero-footprint entries are free to keep. + kind = AdmissionDecisionKind::AdmitProbation; + reasons = vec!["zero-exclusive-bytes".into()]; + } else { + kind = AdmissionDecisionKind::AdmitProbation; + reasons = vec!["probation-new-entry".into()]; + } + + if matches!(kind, AdmissionDecisionKind::Reject) { + return AdmissionDecision { + kind, + verdict: AdmissionVerdict::Reject, + reasons, + probation_cap_repair: crate::policy::CapRepair::satisfied(), + }; + } + // Carry ghost history in: a recurring entry re-enters with its past + // reuse signal, and an entry whose history already clears the hit + // threshold admits straight to the admitted class (equivalent value + // signal). + let had_ghost = policy.ghosts.contains_key(&key); + let mut ghost = policy.ghosts.remove(&key).unwrap_or(GhostStats { + hits: 0, + reuse_weight: 0.0, + observation_weight: 0.0, + last_observation: 0, + }); + if had_ghost { + // This offer *is* a recurrence: count it as a reuse observation, the + // same way record_hit would, so eviction cannot erase the value + // signal the recurrence just demonstrated. + ghost.hits += 1; + ghost.reuse_weight = ghost.reuse_weight * policy.config.decay.factor + 1.0; + ghost.observation_weight = ghost.observation_weight * policy.config.decay.factor + 1.0; + } + let ghost_promoted = ghost.hits >= policy.config.persistence_hit_threshold as u64; + let state = if ghost_promoted { + PolicyEntryState::Admitted + } else { + PolicyEntryState::Probation + }; + // The decision kind must match the resulting entry state so a + // store-facing caller persists exactly what the policy admitted. + let kind = if ghost_promoted { + AdmissionDecisionKind::AdmitPersist + } else { + kind + }; + let segment_ids: Vec = shared.iter().map(|(s, _)| *s).collect(); + for (segment, size) in &shared { + policy.segments.add(*segment, *size, key); + } + policy.entries.insert( + key, + PolicyEntry { + state, + hits: ghost.hits, + misses: 0, + reuse_weight: ghost.reuse_weight, + observation_weight: ghost.observation_weight, + last_observation: policy.clock, + last_cost: Some(cost), + exclusive_bytes, + segments: segment_ids, + }, + ); + AdmissionDecision { + kind, + verdict: AdmissionVerdict::Admit, + reasons, + probation_cap_repair: crate::policy::CapRepair::satisfied(), + } +} + +pub(crate) fn record_hit( + policy: &mut BenefitPolicy, + key: EntryKey, + cost: CostSample, +) -> Option { + if !cost.is_valid() { + // Invalid measured costs never mutate entry state. + return None; + } + let entry = policy.entries.get_mut(&key)?; + entry.hits += 1; + entry.last_observation = policy.clock; + entry.last_cost = Some(cost); + entry.reuse_weight = entry.reuse_weight * policy.config.decay.factor + 1.0; + entry.observation_weight = entry.observation_weight * policy.config.decay.factor + 1.0; + + if entry.state == PolicyEntryState::Probation + && entry.hits >= policy.config.persistence_hit_threshold as u64 + { + entry.state = PolicyEntryState::Admitted; + return Some(AdmissionDecision { + kind: AdmissionDecisionKind::Promote, + verdict: AdmissionVerdict::Admit, + reasons: vec!["probation-second-hit".into()], + probation_cap_repair: crate::policy::CapRepair::satisfied(), + }); + } + None +} + +pub(crate) fn choose_victims( + policy: &mut BenefitPolicy, + bytes_to_free: u64, + pinned: &[EntryKey], +) -> Vec<(EntryKey, EvictionVerdict)> { + let grace = policy.config.grace_observations; + let clock = policy.clock; + let in_grace = |e: &PolicyEntry| clock.saturating_sub(e.last_observation) < grace; + + // NaN-free scores only (score::compute rejects invalid costs/config), so + // a total ordering is safe: finite values ordered normally, keys + // tie-break. We still avoid partial_cmp().unwrap() defensively. + let mut candidates: Vec<(f64, EntryKey)> = policy + .entries + .iter() + .filter(|(k, _)| !pinned.contains(k)) + .filter_map(|(k, e)| { + let score = super::score::compute(&policy.config, *k, e, &policy.segments)?; + Some((score.value, *k)) + }) + .collect(); + candidates.sort_by(|a, b| a.0.total_cmp(&b.0).then_with(|| a.1.cmp(&b.1))); + + // Marginal physical release for `key` given `selected` victims already + // chosen: exclusive bytes plus segments whose remaining references are + // all inside the selected victim set (or the entry itself). + let marginal_release = + |key: EntryKey, selected: &std::collections::BTreeSet| -> u64 { + let entry = policy.entries.get(&key).expect("candidate present"); + let mut release = entry.exclusive_bytes; + for segment in &entry.segments { + if let Some(record) = policy.segments.segment_record(*segment) { + let others_kept = record + .references + .iter() + .any(|r| *r != key && !selected.contains(r)); + if !others_kept { + release += record.size; + } + } + } + release + }; + + // Eviction order: no-hit probationers first (they have not earned + // bytes), oldest observation first — LRU within the probation class so a + // recently offered candidate is not the automatic first victim. Then + // ascending benefit score. Grace waives under a hard capacity request: + // the byte budget always wins. + let mut order: Vec<(u64, EntryKey)> = policy + .entries + .iter() + .filter(|(k, e)| { + !pinned.contains(k) && e.state == PolicyEntryState::Probation && e.hits == 0 + }) + .map(|(k, e)| (e.last_observation, *k)) + .collect(); + order.sort(); // (oldest observation, key): deterministic + let mut ordered_keys: Vec = order.into_iter().map(|(_, k)| k).collect(); + ordered_keys.extend(candidates.into_iter().map(|(_, key)| key)); + let order = ordered_keys; + + let mut freed: u64 = 0; + let mut victims = Vec::new(); + let mut selected: std::collections::BTreeSet = Default::default(); + // Pass 1: grace honored. Pass 2 (only if the hard budget is still unmet): + // grace waived — the hard byte budget always wins under pressure. + for waive_grace in [false, true] { + for key in &order { + if freed >= bytes_to_free { + break; + } + if selected.contains(key) || !policy.entries.contains_key(key) { + continue; // dedup: probation-first and scored paths overlap + } + let evictable = waive_grace || policy.entries.get(key).is_some_and(|e| !in_grace(e)); + if !evictable { + continue; + } + let release = marginal_release(*key, &selected); + selected.insert(*key); + victims.push((*key, EvictionVerdict::Evict)); + freed += release; + } + if freed >= bytes_to_free { + break; + } + } + victims +} diff --git a/crates/skippy-cache/src/policy/decay.rs b/crates/skippy-cache/src/policy/decay.rs new file mode 100644 index 0000000000..e862af8a0c --- /dev/null +++ b/crates/skippy-cache/src/policy/decay.rs @@ -0,0 +1,25 @@ +//! Exponential decay of reuse statistics (#1650 first slice). + +/// Decay tunables. `factor` is the per-observation retention weight applied to +/// past history: 0.9 keeps 90% of each entry's accumulated ratio weight per +/// new observation. Higher `pressure_decay_sensitivity` (via +/// `BenefitPolicy::observe_pressure`) pulls retention toward 0 faster under +/// churn. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct DecayConfig { + /// Per-observation retention factor in `(0, 1)`. + pub factor: f64, +} + +impl DecayConfig { + /// The retention factor must be finite and strictly inside `(0, 1)`. + pub fn is_valid(&self) -> bool { + self.factor.is_finite() && self.factor > 0.0 && self.factor < 1.0 + } +} + +impl Default for DecayConfig { + fn default() -> Self { + Self { factor: 0.9 } + } +} diff --git a/crates/skippy-cache/src/policy/lru_baseline.rs b/crates/skippy-cache/src/policy/lru_baseline.rs new file mode 100644 index 0000000000..4cbaeab97e --- /dev/null +++ b/crates/skippy-cache/src/policy/lru_baseline.rs @@ -0,0 +1,60 @@ +//! Matched same-capacity LRU baseline for policy comparison (#1650 first +//! slice). Capacity is in exclusive bytes so the comparison is apples-to- +//! apples against `BenefitPolicy` under a hard byte budget. + +use crate::policy::traces::TraceAccess; + +pub struct LruCache { + capacity_bytes: u64, + used_bytes: u64, + /// Most-recent first. + order: Vec, + sizes: std::collections::HashMap, + pub hits: u64, + pub misses: u64, + /// Saved cold-prefill cost (higher is better) over the trace. + pub saved_cost: f64, + /// Bytes written into the cache (lower is better). + pub bytes_written: u64, +} + +impl LruCache { + pub fn new(capacity_bytes: u64) -> Self { + Self { + capacity_bytes, + used_bytes: 0, + order: Vec::new(), + sizes: std::collections::HashMap::new(), + hits: 0, + misses: 0, + saved_cost: 0.0, + bytes_written: 0, + } + } + + pub fn access(&mut self, access: &TraceAccess) -> bool { + if self.order.contains(&access.entry) { + self.hits += 1; + self.saved_cost += (access.cold_prefill_cost - access.restore_cost).max(0.0); + self.order.retain(|e| *e != access.entry); + self.order.insert(0, access.entry); + return true; + } + self.misses += 1; + while self.used_bytes + access.exclusive_bytes > self.capacity_bytes { + let Some(victim) = self.order.pop() else { + break; + }; + if let Some(size) = self.sizes.remove(&victim) { + self.used_bytes = self.used_bytes.saturating_sub(size); + } + } + if access.exclusive_bytes <= self.capacity_bytes { + self.order.insert(0, access.entry); + self.sizes.insert(access.entry, access.exclusive_bytes); + self.used_bytes += access.exclusive_bytes; + self.bytes_written += access.exclusive_bytes; + } + false + } +} diff --git a/crates/skippy-cache/src/policy/mod.rs b/crates/skippy-cache/src/policy/mod.rs new file mode 100644 index 0000000000..0ce74a9bee --- /dev/null +++ b/crates/skippy-cache/src/policy/mod.rs @@ -0,0 +1,615 @@ +//! Benefit-per-exclusive-byte admission, probation, and eviction policy +//! (issue #1650, first slice). +//! +//! This module is deliberately a *pure policy*: it consumes observed events +//! (candidate offers, hits, misses, cost samples) and produces decisions +//! (admit / probation / persist / evict) with opaque reasons. It performs no +//! I/O, holds no locks on the restore path, and never sees prompt content — +//! entries are addressed by an opaque `EntryKey` the caller assigns. +//! +//! The score is the one the issue prescribes: +//! +//! ```text +//! reuse_probability * max(cold_prefill_cost - restore_cost, 0) +//! ----------------------------------------------------------- +//! exclusive_physical_bytes +//! ``` +//! +//! Shared segments are credited fractionally: a physical byte referenced by +//! N entries counts as `bytes / N` against each of them, so total accounted +//! bytes never double-count a segment. +//! +//! Determinism: every ordering falls back to `(score, entry_key)` so two runs +//! over the same trace make identical decisions. + +mod accounting; +mod admission; +mod decay; +#[cfg(test)] +mod lru_baseline; +mod score; +#[cfg(test)] +mod tests; +#[cfg(test)] +mod traces; + +pub use accounting::{SegmentId, SharedSegmentLedger}; +pub use admission::{AdmissionDecision, AdmissionDecisionKind, AdmissionVerdict, PolicyEntryState}; +pub use decay::DecayConfig; +pub use score::{BenefitScore, ScoreInputs}; + +use std::collections::BTreeMap; + +use serde::Serialize; + +/// Opaque, caller-assigned entry identity. Ordered so tie-breaks are +/// deterministic; content-free so policy logs leak nothing about prompts. +pub type EntryKey = u64; + +/// Opaque shared-segment identity. +pub type SegmentRef = SegmentId; + +/// Observed costs, in caller-defined units (the policy only compares them). +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct CostSample { + /// Cold prefill cost for the entry's token range (e.g. ms or tokens). + pub cold_prefill_cost: f64, + /// Measured `queue + restore + suffix-prefill` cost for the same range. + pub restore_cost: f64, +} + +impl CostSample { + /// Net benefit of a restore hit over recomputing cold. Never negative. + pub fn net_benefit(&self) -> f64 { + (self.cold_prefill_cost - self.restore_cost).max(0.0) + } + + /// A usable sample must be finite and nonnegative; measured costs that + /// arrive NaN/infinite (or negative) are rejected so scores and orderings + /// stay total and panic-free. + pub fn is_valid(&self) -> bool { + self.cold_prefill_cost.is_finite() + && self.restore_cost.is_finite() + && self.cold_prefill_cost >= 0.0 + && self.restore_cost >= 0.0 + } +} + +/// Policy decision log line: what was decided, and why, without content. +#[derive(Debug, Clone, PartialEq, Serialize)] +pub struct DecisionReason { + pub entry: EntryKey, + pub decision: AdmissionDecisionKind, + /// Machine-readable reason tokens, e.g. `probation-second-hit`. + pub reasons: Vec, + pub score: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +pub enum EvictionVerdict { + Keep, + Evict, +} + +/// Tunables. Defaults follow the issue's guidance; every field is `Copy` and +/// plain so config files can carry it verbatim later. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct PolicyConfig { + /// Bytes probation entries may collectively charge before the policy + /// must start dropping the least valuable probationers. + pub probation_byte_budget: u64, + /// Hits required to leave probation and become persist-eligible. + /// The issue names second-hit admission: `2`. + pub persistence_hit_threshold: u32, + /// Minimum reuse probability the estimator may report (floor so a single + /// hit still admits under pressure, and division stays sane). + pub min_reuse_probability: f64, + pub decay: DecayConfig, + /// Observations of grace after admission during which an entry cannot be + /// chosen as an eviction victim: probation must get a fair chance to land + /// its second hit before pressure can reclaim its bytes. + pub grace_observations: u64, + /// Maximum retained ghost records. One-shot keys must not create + /// permanent metadata. + pub ghost_capacity: usize, + /// Ghosts older than this many observations are expired so stale + /// popularity cannot revive indefinitely. + pub ghost_max_age_observations: u64, +} + +impl Default for PolicyConfig { + fn default() -> Self { + Self { + probation_byte_budget: 256 << 20, + persistence_hit_threshold: 2, + min_reuse_probability: 0.01, + decay: DecayConfig::default(), + grace_observations: 32, + ghost_capacity: 4096, + ghost_max_age_observations: 1024, + } + } +} + +impl PolicyConfig { + /// Config bounds: NaN/empty decay or an out-of-range reuse floor would + /// poison every score. `probation_byte_budget` may be 0 (probation off). + pub fn is_valid(&self) -> bool { + self.min_reuse_probability.is_finite() + && (0.0..=1.0).contains(&self.min_reuse_probability) + && self.decay.is_valid() + && self.persistence_hit_threshold >= 1 + } +} + +/// Per-entry policy state and statistics. +#[derive(Debug, Clone, PartialEq)] +pub struct PolicyEntry { + pub state: PolicyEntryState, + pub hits: u64, + pub misses: u64, + /// Decayed reuse-estimator numerator/denominator inputs. + pub reuse_weight: f64, + pub observation_weight: f64, + pub last_cost: Option, + /// Exclusive (non-shared) physical bytes charged to this entry. + pub exclusive_bytes: u64, + /// Segments this entry references; fractional credit lives in the ledger. + pub segments: Vec, + /// Clock value at this entry's last admission or hit; drives grace. + pub last_observation: u64, +} + +impl PolicyEntry { + /// Estimated reuse probability under the configured decay window: a + /// smoothed hit ratio in `[0, 1]`. + pub fn reuse_probability(&self) -> f64 { + if self.observation_weight <= 0.0 { + return 0.0; + } + (self.reuse_weight / self.observation_weight).clamp(0.0, 1.0) + } +} + +/// Cap-repair selection result. `Deferred` means pins/holds (or candidate +/// exhaustion) make the hard probation cap temporarily unsatisfiable: the +/// listed victims should still be committed, and the shortfall reported — +/// a pin is never selected to close it. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum CapRepair { + /// The class is at or under cap (possibly after committing `victims`). + Satisfied { victims: Vec }, + /// Committing `victims` still leaves the class `shortfall_bytes` over + /// cap; repair is deferred until pins release. + Deferred { + victims: Vec, + shortfall_bytes: u64, + }, +} + +impl CapRepair { + fn satisfied() -> Self { + CapRepair::Satisfied { + victims: Vec::new(), + } + } + + fn satisfied_with(victims: Vec) -> Self { + CapRepair::Satisfied { victims } + } + + fn deferred(victims: Vec, shortfall_bytes: u64) -> Self { + CapRepair::Deferred { + victims, + shortfall_bytes, + } + } + + /// Victims to commit regardless of variant. + pub fn victims(&self) -> &[EntryKey] { + match self { + CapRepair::Satisfied { victims } | CapRepair::Deferred { victims, .. } => victims, + } + } + + /// Over-cap bytes that cannot be repaired while pins are held. + pub fn shortfall_bytes(&self) -> u64 { + match self { + CapRepair::Satisfied { .. } => 0, + CapRepair::Deferred { + shortfall_bytes, .. + } => *shortfall_bytes, + } + } + + /// True when no further repair is possible right now. + pub fn is_deferred(&self) -> bool { + matches!(self, CapRepair::Deferred { .. }) + } +} + +/// Result of a committed removal. +#[derive(Debug, Clone, PartialEq)] +pub struct RemovalOutcome { + pub entry: PolicyEntry, + /// Cap-repair plan after this removal (shares may have risen). Pins are + /// never selected; a `Deferred` result reports the shortfall. + pub probation_cap_repair: CapRepair, +} + +/// The policy engine. Owns per-entry statistics and the shared-segment ledger; +/// the caller drives it from cache events. +pub struct BenefitPolicy { + pub(crate) config: PolicyConfig, + pub(crate) entries: BTreeMap, + pub(crate) segments: SharedSegmentLedger, + /// Monotonic observation counter driving the probation grace window. + pub(crate) clock: u64, + /// Reuse statistics that outlive eviction ("ghosts"): an entry that + /// recurs after eviction carries its history back in, so the second-hit + /// value signal survives cache pressure. Bounded by count and age. + pub(crate) ghosts: BTreeMap, +} + +impl BenefitPolicy { + /// Insert a ghost, evicting the oldest ghost when the count bound is + /// exceeded. `ghost_capacity = 0` disables ghost retention entirely. + fn insert_ghost(&mut self, key: EntryKey, stats: GhostStats) { + if self.config.ghost_capacity == 0 { + return; + } + while self.ghosts.len() >= self.config.ghost_capacity { + let oldest = self + .ghosts + .iter() + .min_by_key(|(k, g)| (g.last_observation, **k)) + .map(|(k, _)| *k); + match oldest { + Some(k) => { + self.ghosts.remove(&k); + } + None => break, + } + } + self.ghosts.insert(key, stats); + } + + /// Drop ghosts older than the configured age bound. Called from the + /// observation-driven entry points so expiry is deterministic over a + /// trace without a background timer. + fn expire_ghosts(&mut self) { + let horizon = self + .clock + .saturating_sub(self.config.ghost_max_age_observations); + self.ghosts.retain(|_, g| g.last_observation >= horizon); + } +} + +/// Surviving statistics for an evicted entry. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct GhostStats { + pub hits: u64, + pub reuse_weight: f64, + pub observation_weight: f64, + /// Clock value when the ghost was created; drives age expiry. + pub last_observation: u64, +} + +impl BenefitPolicy { + /// Panics on invalid config so misconfiguration fails at startup + /// rather than producing NaN scores later. + pub fn new(config: PolicyConfig) -> Self { + assert!(config.is_valid(), "invalid PolicyConfig: {:?}", config); + Self { + config, + entries: BTreeMap::new(), + segments: SharedSegmentLedger::default(), + clock: 0, + ghosts: BTreeMap::new(), + } + } + + pub fn config(&self) -> &PolicyConfig { + &self.config + } + + pub fn ghost(&self, key: EntryKey) -> Option<&GhostStats> { + self.ghosts.get(&key) + } + + pub fn ghost_count(&self) -> usize { + self.ghosts.len() + } + + /// Current observation clock (test-only access). + #[cfg(test)] + pub fn clock_debug(&self) -> u64 { + self.clock + } + + pub fn entry(&self, key: EntryKey) -> Option<&PolicyEntry> { + self.entries.get(&key) + } + + pub fn len(&self) -> usize { + self.entries.len() + } + + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + /// Offer a new candidate for admission. `exclusive_bytes` are bytes only + /// this entry would reference; `shared` lists `(segment, size)` segments + /// it would join (size is ignored if the segment is already registered). + /// `pinned` is the caller's active pin/hold set: pinned entries are + /// never selected for probation cap repair (#1650). + pub fn consider_admission( + &mut self, + key: EntryKey, + exclusive_bytes: u64, + shared: Vec<(SegmentId, u64)>, + cost: CostSample, + pinned: &[EntryKey], + ) -> AdmissionDecision { + // Full structural prevalidation before any policy mutation: an + // invalid cost, duplicate segment IDs, a size conflict, or an + // already-resident key must not advance the clock (aging grace) or + // expire ghosts. + if !cost.is_valid() { + return AdmissionDecision::rejected_invalid_cost(); + } + if self.entries.contains_key(&key) { + return AdmissionDecision::rejected("already-resident-key"); + } + let mut seen = std::collections::BTreeSet::new(); + if shared.iter().any(|(s, _)| !seen.insert(*s)) { + return AdmissionDecision::rejected("duplicate-segment-reference"); + } + if shared.iter().any(|(segment, size)| { + self.segments + .segment_record(*segment) + .is_some_and(|r| !r.references.contains(&key) && r.size != *size) + }) { + return AdmissionDecision::rejected("segment-size-conflict"); + } + self.clock += 1; + self.expire_ghosts(); + let mut decision = admission::consider(self, key, exclusive_bytes, shared, cost); + if decision.verdict == crate::policy::AdmissionVerdict::Admit { + // Hard probation cap is part of admission: the decision carries + // the cap-repair plan (pins excluded). Selection only — committed + // removal stays with `remove` so victims become ghosts. + decision.probation_cap_repair = self.select_probation_cap_victims(pinned); + } + decision + } + + /// Record a restore hit on an admitted entry; may promote out of probation. + pub fn record_hit(&mut self, key: EntryKey, cost: CostSample) -> Option { + // Validate before any mutation (including the clock). + if !cost.is_valid() { + return None; + } + self.clock += 1; + admission::record_hit(self, key, cost) + } + + /// Record a miss/cold recompute for an admitted entry (decays reuse). + /// A miss is a real observation that advances the clock, but unlike a + /// hit it is a negative value signal: it decays reuse and does NOT + /// refresh `last_observation`, so a miss-only stream ages the grace + /// window and the entry becomes evictable — grace cannot hold a + /// never-reused entry indefinitely. + pub fn record_miss(&mut self, key: EntryKey) { + self.clock += 1; + if let Some(entry) = self.entries.get_mut(&key) { + entry.misses += 1; + entry.observation_weight = entry.observation_weight * self.config.decay.factor + 1.0; + } + } + + /// Observe demand pressure (0 = idle, 1 = saturated). Higher pressure + /// decays reuse history faster than the observation base, so stale + /// popularity cannot pin bytes forever. + pub fn observe_pressure(&mut self, pressure: f64) { + // Non-finite pressure never poisons history: NaN is ignored, + // +inf saturates to full pressure, -inf to none. + if pressure.is_nan() { + return; + } + let pressure = if pressure == f64::INFINITY { + 1.0 + } else if pressure == f64::NEG_INFINITY { + 0.0 + } else { + pressure.clamp(0.0, 1.0) + }; + let base = self.config.decay.factor; + let reuse_factor = base * (1.0 - pressure); + for entry in self.entries.values_mut() { + entry.reuse_weight *= reuse_factor; + entry.observation_weight *= base + (1.0 - base) * pressure; + } + for ghost in self.ghosts.values_mut() { + ghost.reuse_weight *= reuse_factor; + ghost.observation_weight *= base + (1.0 - base) * pressure; + } + self.expire_ghosts(); + } + + /// Score an entry under the current statistics. Returns `None` for + /// entries with no cost observation yet. + pub fn score(&self, key: EntryKey) -> Option { + let entry = self.entries.get(&key)?; + score::compute(&self.config, key, entry, &self.segments) + } + + /// Pick eviction victims until `bytes_to_free` exclusive-and-fractional + /// bytes are released. Lowest score first, deterministic `(score, key)` + /// tie-break. Pinned entries are never chosen while an unpinned + /// candidate remains. + pub fn choose_victims( + &mut self, + bytes_to_free: u64, + pinned: &[EntryKey], + ) -> Vec<(EntryKey, EvictionVerdict)> { + admission::choose_victims(self, bytes_to_free, pinned) + } + + /// Bytes charged to the probation class: exclusive bytes plus the + /// fractional shared-segment credit. Shares are summed exactly (f64) + /// across all probation entries before a single rounding, so a small + /// segment shared by many references still charges its physical bytes + /// in aggregate — per-entry truncation cannot zero it out. + pub fn probation_bytes(&self) -> u64 { + let mut exact = 0.0f64; + for (k, e) in &self.entries { + if e.state == PolicyEntryState::Probation { + exact += e.exclusive_bytes as f64 + self.segments.fractional_bytes(*k, &e.segments); + } + } + // Deterministic round-half-up; the class charge is a hard bound, so + // we always round up any fractional residue. + exact.ceil() as u64 + } + + /// Select the no-hit probationers that must be evicted to bring the + /// probation class back under its byte cap, oldest observation first + /// (grace waived: the hard cap always wins). Selection only — this does + /// not mutate policy state; the caller commits each removal via + /// `remove`, which also records the ghost. Keys are deterministic + /// `(last_observation, key)` order. + /// + /// Pinned entries are never selected (#1650: preserve active + /// pins/holds). If pins make the cap temporarily unsatisfiable, the + /// result reports the shortfall instead of selecting a pin. + pub fn select_probation_cap_victims(&self, pinned: &[EntryKey]) -> CapRepair { + let cap = self.config.probation_byte_budget; + if self.probation_bytes() <= cap { + return CapRepair::satisfied(); + } + let mut probationers: Vec<(u64, EntryKey)> = self + .entries + .iter() + .filter(|(k, e)| { + e.state == PolicyEntryState::Probation && e.hits == 0 && !pinned.contains(k) + }) + .map(|(k, e)| (e.last_observation, *k)) + .collect(); + probationers.sort(); + // Fallback: probationers with hits still count against the class + // charge (e.g. their share rose when an admitted co-reference was + // removed), so cap repair must be able to select them too — after + // the zero-hit class, in deterministic age/key order. + let mut fallback: Vec<(u64, EntryKey)> = self + .entries + .iter() + .filter(|(k, e)| { + e.state == PolicyEntryState::Probation && e.hits > 0 && !pinned.contains(k) + }) + .map(|(k, e)| (e.last_observation, *k)) + .collect(); + fallback.sort(); + probationers.extend(fallback); + // Simulate each removal against a scratch ledger: removing a shared + // reference raises the survivors' fractional shares, so the remaining + // class charge must be recomputed, not decremented by stale shares. + let mut scratch = self.segments.clone(); + let mut removed: std::collections::BTreeSet = Default::default(); + let mut victims = Vec::new(); + for (_, key) in probationers { + let Some(entry) = self.entries.get(&key) else { + continue; + }; + scratch.release(&entry.segments, key); + removed.insert(key); + let charge: f64 = self + .entries + .iter() + .filter(|(k, e)| e.state == PolicyEntryState::Probation && !removed.contains(*k)) + .map(|(k, e)| e.exclusive_bytes as f64 + scratch.fractional_bytes(*k, &e.segments)) + .sum(); + victims.push(key); + if charge.ceil() as u64 <= cap { + break; + } + } + // After simulating every unpinned candidate, the class may still be + // over cap because the remainder is pinned (or candidateless). Report + // the shortfall explicitly instead of ever selecting a pin. + let final_charge: u64 = self + .entries + .iter() + .filter(|(k, e)| e.state == PolicyEntryState::Probation && !removed.contains(*k)) + .map(|(k, e)| e.exclusive_bytes as f64 + scratch.fractional_bytes(*k, &e.segments)) + .sum::() + .ceil() as u64; + if final_charge <= cap { + CapRepair::satisfied_with(victims) + } else { + CapRepair::deferred(victims, final_charge - cap) + } + } + + /// Test helper: run `f` with a temporarily different probation budget. + #[cfg(test)] + pub fn with_probation_budget(&mut self, budget: u64, f: impl FnOnce(&Self) -> R) -> R { + let saved = std::mem::replace( + // Safety of construction: same struct, one field changed. + &mut self.config.probation_byte_budget, + budget, + ); + let result = f(self); + self.config.probation_byte_budget = saved; + result + } + + /// Test-only convenience: cap enforcement with an empty pin set for the + /// comparison harness/traces, which model unpinned workloads. Store + /// callers holding pins must call `select_probation_cap_victims` with + /// their pin set and commit via `remove`. + #[cfg(test)] + pub fn enforce_probation_cap(&mut self) -> Vec { + let victims = self.select_probation_cap_victims(&[]).victims().to_vec(); + for key in victims.iter().copied() { + let _ = self.remove(key, &[]); + } + victims + } + + /// Store-facing removal: releases the entry's fractional segment credit + /// and stashes its reuse statistics as a ghost so a recurrence is + /// recognized as a value signal. Removing an admitted co-reference can + /// raise survivors' shares over the probation cap, so the removal + /// response carries the cap-repair plan; commit its victims through + /// further `remove` calls. `pinned` is the caller's active pin/hold + /// set: pins are never selected (#1650). + pub fn remove(&mut self, key: EntryKey, pinned: &[EntryKey]) -> Option { + let entry = self.entries.remove(&key)?; + self.segments.release(&entry.segments, key); + self.insert_ghost( + key, + GhostStats { + hits: entry.hits, + reuse_weight: entry.reuse_weight, + observation_weight: entry.observation_weight, + last_observation: self.clock, + }, + ); + let cap_repair = self.select_probation_cap_victims(pinned); + Some(RemovalOutcome { + entry, + probation_cap_repair: cap_repair, + }) + } + + /// Remove without returning cap-repair victims (test-only harness use + /// where the caller re-selects the cap itself). External callers must + /// use `remove` so the hard probation cap cannot be bypassed. + #[cfg(test)] + pub fn remove_without_cap_repair(&mut self, key: EntryKey) -> Option { + let outcome = self.remove(key, &[])?; + Some(outcome.entry) + } +} diff --git a/crates/skippy-cache/src/policy/score.rs b/crates/skippy-cache/src/policy/score.rs new file mode 100644 index 0000000000..b0c8dbc724 --- /dev/null +++ b/crates/skippy-cache/src/policy/score.rs @@ -0,0 +1,50 @@ +//! The benefit-per-exclusive-byte score (#1650 first slice). + +use crate::policy::{CostSample, EntryKey, PolicyConfig, PolicyEntry, SharedSegmentLedger}; + +/// Raw inputs and the resulting score, for tests and opaque logging. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct ScoreInputs { + pub reuse_probability: f64, + pub net_benefit: f64, + pub exclusive_bytes: f64, +} + +/// A computed score with its inputs. Higher is better; ordering by +/// `(value, entry_key)` is the canonical deterministic order. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct BenefitScore { + pub value: f64, + pub inputs: ScoreInputs, +} + +pub(crate) fn compute( + config: &PolicyConfig, + key: EntryKey, + entry: &PolicyEntry, + ledger: &SharedSegmentLedger, +) -> Option { + let cost: CostSample = entry.last_cost?; + if !cost.is_valid() { + return None; + } + let reuse = entry.reuse_probability().max(config.min_reuse_probability); + let shared = ledger.fractional_bytes(key, &entry.segments); + let exclusive = entry.exclusive_bytes as f64 + shared; + if exclusive <= 0.0 { + return None; + } + let net = cost.net_benefit(); + let value = reuse * net / exclusive; + if !value.is_finite() { + return None; + } + Some(BenefitScore { + value, + inputs: ScoreInputs { + reuse_probability: reuse, + net_benefit: net, + exclusive_bytes: exclusive, + }, + }) +} diff --git a/crates/skippy-cache/src/policy/tests.rs b/crates/skippy-cache/src/policy/tests.rs new file mode 100644 index 0000000000..6a6b03df49 --- /dev/null +++ b/crates/skippy-cache/src/policy/tests.rs @@ -0,0 +1,1156 @@ +//! Unit and comparison tests for the benefit policy (#1650 first slice). + +use super::*; +use crate::policy::admission::AdmissionDecisionKind; +use crate::policy::lru_baseline::LruCache; +use crate::policy::traces; + +fn cost(cold: f64, restore: f64) -> CostSample { + CostSample { + cold_prefill_cost: cold, + restore_cost: restore, + } +} + +#[test] +fn rejects_entries_with_no_net_benefit() { + let mut policy = BenefitPolicy::new(PolicyConfig::default()); + let decision = policy.consider_admission(1, 1 << 20, vec![], cost(100.0, 150.0), &[]); + assert_eq!(decision.kind, AdmissionDecisionKind::Reject); + assert!(policy.is_empty()); +} + +#[test] +fn new_entries_start_in_probation_and_promote_on_second_hit() { + let mut policy = BenefitPolicy::new(PolicyConfig::default()); + policy.consider_admission(1, 1 << 20, vec![], cost(400.0, 120.0), &[]); + assert_eq!(policy.entry(1).unwrap().state, PolicyEntryState::Probation); + + assert!(policy.record_hit(1, cost(400.0, 120.0)).is_none()); // first hit + let promoted = policy.record_hit(1, cost(400.0, 120.0)).unwrap(); // second + assert_eq!(promoted.kind, AdmissionDecisionKind::Promote); + assert_eq!(policy.entry(1).unwrap().state, PolicyEntryState::Admitted); +} + +#[test] +fn score_divides_by_exclusive_bytes() { + let mut policy = BenefitPolicy::new(PolicyConfig::default()); + policy.consider_admission(1, 2 << 20, vec![], cost(300.0, 100.0), &[]); + policy.record_hit(1, cost(300.0, 100.0)); + policy.record_hit(1, cost(300.0, 100.0)); + let score = policy.score(1).unwrap(); + // reuse=1.0, net=200.0, exclusive=2MiB + assert!((score.inputs.net_benefit - 200.0).abs() < 1e-9); + assert!((score.inputs.exclusive_bytes - (2.0 * 1024.0 * 1024.0)).abs() < 1e-9); + assert!((score.value - 200.0 / (2.0 * 1024.0 * 1024.0)).abs() < 1e-9); +} + +#[test] +fn shared_segments_get_fractional_credit_in_score() { + let mut policy = BenefitPolicy::new(PolicyConfig::default()); + let seg = (7u64, 3 << 20); + policy.consider_admission(1, 0, vec![seg], cost(400.0, 100.0), &[]); + policy.consider_admission(2, 0, vec![seg], cost(400.0, 100.0), &[]); + policy.consider_admission(3, 0, vec![seg], cost(400.0, 100.0), &[]); + // Each entry charges 1MiB of the shared 3MiB segment. + assert_eq!(policy.segments.total_bytes(), 3 << 20); + for key in [1u64, 2, 3] { + let score = policy.score(key).unwrap(); + assert!((score.inputs.exclusive_bytes - (1.0 * 1024.0 * 1024.0)).abs() < 1e-9); + } +} + +#[test] +fn eviction_picks_lowest_score_deterministically() { + // Grace off: force the score order. + let mut policy = BenefitPolicy::new(PolicyConfig { + grace_observations: 0, + ..PolicyConfig::default() + }); + // Same value, different footprint: smaller footprint -> higher score. + policy.consider_admission(1, 8 << 20, vec![], cost(400.0, 100.0), &[]); + policy.consider_admission(2, 1 << 20, vec![], cost(400.0, 100.0), &[]); + for key in [1u64, 2] { + policy.record_hit(key, cost(400.0, 100.0)); + policy.record_hit(key, cost(400.0, 100.0)); + } + let victims = policy.choose_victims(8 << 20, &[]); + assert_eq!(victims.first().map(|(k, _)| *k), Some(1)); +} + +#[test] +fn pinned_entries_are_not_victims() { + // Grace off: force the pinned check. + let mut policy = BenefitPolicy::new(PolicyConfig { + grace_observations: 0, + ..PolicyConfig::default() + }); + policy.consider_admission(1, 1 << 20, vec![], cost(400.0, 100.0), &[]); + policy.consider_admission(2, 1 << 20, vec![], cost(400.0, 100.0), &[]); + let victims = policy.choose_victims(1 << 20, &[2]); + assert_eq!(victims.iter().map(|(k, _)| *k).collect::>(), vec![1]); +} + +#[test] +fn decay_shrinks_stale_reuse_probability() { + let mut policy = BenefitPolicy::new(PolicyConfig::default()); + policy.consider_admission(1, 1 << 20, vec![], cost(400.0, 100.0), &[]); + for _ in 0..10 { + policy.record_hit(1, cost(400.0, 100.0)); + } + let hot = policy.entry(1).unwrap().reuse_probability(); + for _ in 0..50 { + policy.observe_pressure(1.0); + } + let stale = policy.entry(1).unwrap().reuse_probability(); + assert!(stale < hot); +} + +#[test] +fn policy_decisions_are_deterministic_across_runs() { + let trace = traces::zipf_hotset_trace(42, 500); + let run = || { + let mut policy = BenefitPolicy::new(PolicyConfig::default()); + let mut evictions = Vec::new(); + for access in &trace { + if policy.entry(access.entry).is_some() { + policy.record_hit( + access.entry, + cost(access.cold_prefill_cost, access.restore_cost), + ); + } else { + policy.consider_admission( + access.entry, + access.exclusive_bytes, + vec![], + cost(access.cold_prefill_cost, access.restore_cost), + &[], + ); + } + for (key, verdict) in policy.choose_victims(1 << 30, &[]) { + if verdict == EvictionVerdict::Evict { + evictions.push(key); + policy.remove(key, &[]); + } + } + } + evictions + }; + assert_eq!(run(), run()); +} + +/// Policy-vs-LRU comparison metrics over a trace at matched capacity. +struct Comparison { + policy_saved_cost: f64, + lru_saved_cost: f64, + policy_bytes_written: u64, + lru_bytes_written: u64, +} + +fn compare(trace: &[traces::TraceAccess], capacity_bytes: u64) -> Comparison { + let mut policy = BenefitPolicy::new(PolicyConfig::default()); + let mut lru = LruCache::new(capacity_bytes); + let mut policy_saved = 0.0f64; + let mut policy_written = 0u64; + + for access in trace { + if policy.entry(access.entry).is_some() { + let decision = policy.record_hit( + access.entry, + cost(access.cold_prefill_cost, access.restore_cost), + ); + policy_saved += (access.cold_prefill_cost - access.restore_cost).max(0.0); + // Promotion is the persistence event for a probational entry: + // the bytes written are the resident entry's retained size at + // admission, not the current access's (re-sampled) size. + if let Some(AdmissionDecision { + kind: AdmissionDecisionKind::Promote, + .. + }) = &decision + { + let persisted = policy + .entry(access.entry) + .map(|e| e.exclusive_bytes) + .expect("promoted entry resident"); + policy_written += persisted; + } + } else { + // Re-offer a rejected/evicted entry each time it recurs; + // probation keeps new candidates in the accounting set until + // pressure forces a choice. + let decision = policy.consider_admission( + access.entry, + access.exclusive_bytes, + vec![], + cost(access.cold_prefill_cost, access.restore_cost), + &[], + ); + // Count bytes only when the decision actually persists. + // AdmitProbation is memory-only accounting; the entry becomes + // written when a later hit promotes it to Admitted (persist). + if decision.kind == AdmissionDecisionKind::AdmitPersist { + policy_written += access.exclusive_bytes; + } + // Hard probation cap is enforced as part of admission: commit the + // selected victims. + for key in decision.probation_cap_repair.victims().iter().copied() { + policy.remove(key, &[]); + } + } + let mut used: u64 = policy.entries.values().map(|e| e.exclusive_bytes).sum(); + while used > capacity_bytes { + let victims = policy.choose_victims(used - capacity_bytes, &[]); + if victims.is_empty() { + break; + } + for (key, verdict) in victims { + if verdict == EvictionVerdict::Evict { + policy.remove(key, &[]); + } + } + used = policy.entries.values().map(|e| e.exclusive_bytes).sum(); + } + lru.access(access); + } + Comparison { + policy_saved_cost: policy_saved, + lru_saved_cost: lru.saved_cost, + policy_bytes_written: policy_written, + lru_bytes_written: lru.bytes_written, + } +} + +/// The acceptance direction the issue requires: on a trace where one-shot +/// large entries pollute an LRU, the benefit policy must STRICTLY save +/// more cold prefill cost or STRICTLY write fewer persisted bytes, with +/// the non-winning dimension held inside an explicit regression +/// tolerance. +#[test] +fn beats_lru_on_one_shot_pollution_trace() { + let trace = traces::one_shot_trace(7, 2_000); + let capacity = 64 << 20; + let comparison = compare(&trace, capacity); + const TOLERANCE: f64 = 0.05; + let wins_saved = comparison.policy_saved_cost > comparison.lru_saved_cost; + let wins_written = comparison.policy_bytes_written < comparison.lru_bytes_written; + // Strict improvement in at least one dimension (#1650 acceptance). + assert!( + wins_saved || wins_written, + "one-shot pollution @ {capacity}: saved policy={:.1} lru={:.1}, written policy={} lru={} — no strict improvement", + comparison.policy_saved_cost, + comparison.lru_saved_cost, + comparison.policy_bytes_written, + comparison.lru_bytes_written + ); + // The non-winning dimension must stay inside an explicit regression + // tolerance rather than degrading unboundedly. + if !wins_saved { + assert!( + comparison.policy_saved_cost >= comparison.lru_saved_cost * (1.0 - TOLERANCE), + "policy saved {} vs lru {} — saved cost regressed beyond {:.0}%", + comparison.policy_saved_cost, + comparison.lru_saved_cost, + TOLERANCE * 100.0 + ); + } + if !wins_written { + assert!( + (comparison.policy_bytes_written as f64) + <= comparison.lru_bytes_written as f64 * (1.0 + TOLERANCE), + "policy wrote {} vs lru {} — writes regressed beyond {:.0}%", + comparison.policy_bytes_written, + comparison.lru_bytes_written, + TOLERANCE * 100.0 + ); + } +} + +#[test] +fn no_regression_on_turn_growth_trace() { + let trace = traces::turn_growth_trace(11, 4, 12); + let capacity = 64 << 20; + let comparison = compare(&trace, capacity); + assert!( + comparison.policy_saved_cost >= comparison.lru_saved_cost * 0.95, + "policy saved {} vs lru {}", + comparison.policy_saved_cost, + comparison.lru_saved_cost + ); +} + +#[test] +fn no_regression_on_zipf_hotset_trace() { + let trace = traces::zipf_hotset_trace(3, 4_000); + let capacity = 64 << 20; + let comparison = compare(&trace, capacity); + assert!( + comparison.policy_saved_cost >= comparison.lru_saved_cost * 0.95, + "policy saved {} vs lru {}", + comparison.policy_saved_cost, + comparison.lru_saved_cost + ); +} + +#[test] +fn no_regression_on_mixed_size_trace() { + let trace = traces::mixed_size_trace(5, 3_000); + let capacity = 256 << 20; + let comparison = compare(&trace, capacity); + assert!( + comparison.policy_saved_cost >= comparison.lru_saved_cost * 0.95, + "policy saved {} vs lru {}", + comparison.policy_saved_cost, + comparison.lru_saved_cost + ); +} + +/// The mixed-size trace must actually exercise mixed resident size +/// classes and large-entry eviction at a capacity smaller than the +/// 256 MiB class. +#[test] +fn mixed_size_trace_has_real_mixed_residency_and_large_eviction() { + let trace = traces::mixed_size_trace(5, 300); + // Distinct keys per class, each consistently sized. + let mut sizes = std::collections::BTreeMap::new(); + let mut classes = std::collections::BTreeSet::new(); + for a in &trace { + let prev = sizes.insert(a.entry, a.exclusive_bytes); + assert!( + prev.is_none_or(|s| s == a.exclusive_bytes), + "key {} changed size class", + a.entry + ); + classes.insert(a.exclusive_bytes); + } + assert_eq!(classes.len(), 3, "all three size classes present"); + // Run the trace at a capacity that cannot hold a 256 MiB entry + // alongside the hot 4 MiB set: large entries must be admitted + // (probation) and evicted, never violating capacity. + let capacity = 96 << 20; + let mut policy = BenefitPolicy::new(PolicyConfig::default()); + let mut resident_classes = std::collections::BTreeSet::new(); + for access in &trace { + if policy.entry(access.entry).is_some() { + policy.record_hit( + access.entry, + cost(access.cold_prefill_cost, access.restore_cost), + ); + } else { + let decision = policy.consider_admission( + access.entry, + access.exclusive_bytes, + vec![], + cost(access.cold_prefill_cost, access.restore_cost), + &[], + ); + for key in decision.probation_cap_repair.victims().iter().copied() { + policy.remove(key, &[]); + } + } + let mut used: u64 = policy.entries.values().map(|e| e.exclusive_bytes).sum(); + while used > capacity { + let victims = policy.choose_victims(used - capacity, &[]); + if victims.is_empty() { + break; + } + for (key, verdict) in victims { + if verdict == EvictionVerdict::Evict { + policy.remove(key, &[]); + } + } + used = policy.entries.values().map(|e| e.exclusive_bytes).sum(); + } + assert!(used <= capacity, "used {} over capacity after access", used); + for e in policy.entries.values() { + resident_classes.insert(e.exclusive_bytes); + } + } + assert!( + resident_classes.len() >= 2, + "mixed resident sizes never observed: {:?}", + resident_classes + ); +} + +#[test] +fn invalid_costs_and_config_are_rejected_not_panicked() { + let mut policy = BenefitPolicy::new(PolicyConfig::default()); + let nan = f64::NAN; + policy.consider_admission( + 1, + 1 << 20, + vec![], + CostSample { + cold_prefill_cost: nan, + restore_cost: 10.0, + }, + &[], + ); + assert!(policy.score(1).is_none()); // NaN never enters the ordering + assert!( + !PolicyConfig { + min_reuse_probability: f64::NAN, + ..PolicyConfig::default() + } + .is_valid() + ); + assert!( + !PolicyConfig { + persistence_hit_threshold: 0, + ..PolicyConfig::default() + } + .is_valid() + ); + assert!( + !PolicyConfig { + decay: DecayConfig { factor: 0.0 }, + ..PolicyConfig::default() + } + .is_valid() + ); + assert!( + !PolicyConfig { + decay: DecayConfig { factor: 1.5 }, + ..PolicyConfig::default() + } + .is_valid() + ); +} + +#[test] +fn choose_victims_never_returns_a_duplicate_key() { + let mut policy = BenefitPolicy::new(PolicyConfig { + grace_observations: 0, + ..PolicyConfig::default() + }); + for key in 1..=8u64 { + policy.consider_admission(key, 1 << 20, vec![], cost(400.0, 100.0), &[]); + } + let victims = policy.choose_victims(u64::MAX, &[]); + let keys: Vec = victims.iter().map(|(k, _)| *k).collect(); + let mut sorted = keys.clone(); + sorted.sort(); + sorted.dedup(); + assert_eq!(keys.len(), sorted.len(), "duplicate victim keys"); + assert_eq!(keys.len(), 8); +} + +#[test] +fn victim_selection_counts_marginal_physical_bytes_of_shared_segments() { + // Three entries share one segment; only evicting the last reference + // physically frees it. Victim selection must keep choosing until the + // requested *physical* bytes are covered. + let mut policy = BenefitPolicy::new(PolicyConfig { + grace_observations: 0, + ..PolicyConfig::default() + }); + let seg = (1u64, 3 << 20); + policy.consider_admission(1, 0, vec![seg], cost(400.0, 100.0), &[]); + policy.consider_admission(2, 0, vec![seg], cost(400.0, 100.0), &[]); + policy.consider_admission(3, 0, vec![seg], cost(400.0, 100.0), &[]); + for key in [1u64, 2, 3] { + policy.record_hit(key, cost(400.0, 100.0)); + policy.record_hit(key, cost(400.0, 100.0)); + } + // Freeing 3 MiB must select all three references, not one (each alone + // releases nothing physical). + let victims = policy.choose_victims(3 << 20, &[]); + assert_eq!(victims.len(), 3); + // Freeing 1 MiB: marginal release of two of the three is 0, the third is + // 3 MiB; the loop must not stop before the target is met. + let victims = policy.choose_victims(1 << 20, &[]); + assert_eq!(victims.len(), 3); +} + +#[test] +fn probation_byte_cap_is_enforced_over_grace() { + // Grace must never hold the policy over the hard probation byte cap. + // Cap fits two small entries only. + let mut policy = BenefitPolicy::new(PolicyConfig { + probation_byte_budget: 2 << 20, + ..PolicyConfig::default() + }); + for key in 1..=10u64 { + policy.consider_admission(key, 1 << 20, vec![], cost(400.0, 100.0), &[]); + } + assert!(policy.probation_bytes() > 2 << 20); + for key in policy.enforce_probation_cap() { + policy.remove(key, &[]); + } + assert!( + policy.probation_bytes() <= 2 << 20, + "probation bytes {} over cap {}", + policy.probation_bytes(), + 2 << 20 + ); +} + +#[test] +fn recurrence_after_eviction_is_a_value_signal() { + // An entry evicted before its second hit must not restart from zero + // history: its recurrence carries ghost statistics and counts as a + // reuse observation (issue: "second-hit or equivalent value signal"). + let mut policy = BenefitPolicy::new(PolicyConfig { + grace_observations: 0, + ..PolicyConfig::default() + }); + policy.consider_admission(1, 1 << 20, vec![], cost(400.0, 100.0), &[]); + assert!(policy.record_hit(1, cost(400.0, 100.0)).is_none()); // first hit + policy.remove(1, &[]); // evicted before the second hit + policy.consider_admission(1, 1 << 20, vec![], cost(400.0, 100.0), &[]); + assert_eq!( + policy.entry(1).unwrap().state, + PolicyEntryState::Admitted, + "recurrence with ghost history must re-admit as a value signal" + ); +} + +#[test] +fn invalid_cost_samples_never_mutate_entry_state() { + let mut policy = BenefitPolicy::new(PolicyConfig::default()); + let nan = f64::NAN; + // Admission with an invalid sample must not insert. + let decision = policy.consider_admission( + 1, + 1 << 20, + vec![], + CostSample { + cold_prefill_cost: nan, + restore_cost: 10.0, + }, + &[], + ); + assert_eq!(decision.verdict, AdmissionVerdict::Reject); + assert!(policy.is_empty()); + // A valid admission followed by an invalid hit must not poison state. + policy.consider_admission(1, 1 << 20, vec![], cost(400.0, 100.0), &[]); + let before = policy.entry(1).unwrap().clone(); + assert!( + policy + .record_hit( + 1, + CostSample { + cold_prefill_cost: 400.0, + restore_cost: nan + } + ) + .is_none() + ); + assert_eq!(policy.entry(1).unwrap().last_cost, before.last_cost); + assert_eq!(policy.entry(1).unwrap().hits, before.hits); + // NaN pressure is ignored entirely. + policy.observe_pressure(nan); + assert_eq!(policy.entry(1).unwrap().reuse_weight, before.reuse_weight); +} + +#[test] +fn probation_cap_counts_shared_charge_and_is_enforced_by_admission() { + // Probation entries backed entirely by shared segments charge fractional + // bytes and must not grow without bound. + let mut policy = BenefitPolicy::new(PolicyConfig { + probation_byte_budget: 1 << 20, // smaller than one shared segment + grace_observations: 0, + ..PolicyConfig::default() + }); + let mut victims = Vec::new(); + for key in 1..=6u64 { + let decision = + policy.consider_admission(key, 0, vec![(1u64, 3 << 20)], cost(400.0, 100.0), &[]); + for v in decision.probation_cap_repair.victims().iter().copied() { + policy.remove(v, &[]); + victims.push(v); + } + } + assert!( + policy.probation_bytes() <= 1 << 20, + "probation bytes {} over cap", + policy.probation_bytes() + ); + assert!(!victims.is_empty(), "shared-only probation must be capped"); +} + +#[test] +fn ghosts_are_bounded_by_count_and_age() { + // Count bound: one-shot keys must not create permanent metadata. + let mut policy = BenefitPolicy::new(PolicyConfig { + ghost_capacity: 8, + ghost_max_age_observations: 100, + grace_observations: 0, + ..PolicyConfig::default() + }); + for key in 0..64u64 { + policy.consider_admission(key, 1 << 20, vec![], cost(400.0, 100.0), &[]); + for (_, verdict) in policy.choose_victims(u64::MAX, &[]) { + if verdict == EvictionVerdict::Evict { + policy.remove(key, &[]); + break; + } + } + } + assert!( + policy.ghost_count() <= 8, + "ghost count {}", + policy.ghost_count() + ); + + // Age bound: stale popularity cannot revive indefinitely. + let mut policy = BenefitPolicy::new(PolicyConfig { + ghost_capacity: 1024, + ghost_max_age_observations: 10, + grace_observations: 0, + ..PolicyConfig::default() + }); + policy.consider_admission(99, 1 << 20, vec![], cost(400.0, 100.0), &[]); + policy.remove(99, &[]); + assert!(policy.ghost(99).is_some()); + for i in 0..50u64 { + // Unique keys: an already-resident key is rejected before the clock + // advances, so age must be driven by fresh observations. + policy.consider_admission(1000 + i, 1 << 20, vec![], cost(400.0, 100.0), &[]); + } + assert!( + policy.ghost(99).is_none(), + "old ghost must expire via age bound" + ); +} + +#[test] +fn probation_cap_selection_is_not_committed_removal() { + // Selection must leave state untouched so callers can commit physically; + // committed removal through `remove` records the ghost. + let mut policy = BenefitPolicy::new(PolicyConfig { + probation_byte_budget: 1 << 20, + grace_observations: 0, + ..PolicyConfig::default() + }); + for key in 1..=4u64 { + policy.consider_admission(key, 1 << 20, vec![], cost(400.0, 100.0), &[]); + } + let victims = policy.select_probation_cap_victims(&[]).victims().to_vec(); + assert!(!victims.is_empty()); + assert_eq!(policy.len(), 4, "selection must not remove entries"); + let decision = policy.consider_admission(5, 1 << 20, vec![], cost(400.0, 100.0), &[]); + assert!(!decision.probation_cap_repair.victims().is_empty()); + for key in decision.probation_cap_repair.victims() { + policy.remove(*key, &[]); + } + // Committed removals become ghosts (bounded). + assert!(policy.ghost_count() > 0); +} + +#[test] +fn invalid_observation_leaves_clock_and_ghosts_untouched() { + let nan = f64::NAN; + let mut policy = BenefitPolicy::new(PolicyConfig { + grace_observations: 2, + ghost_max_age_observations: 1, + ..PolicyConfig::default() + }); + // Seed one ghost. + policy.consider_admission(9, 1 << 20, vec![], cost(400.0, 100.0), &[]); + policy.remove(9, &[]); + assert!(policy.ghost(9).is_some()); + + // Invalid admission: clock must not advance, ghost must survive + // (age bound is 1, so a real observation would have expired it). + let before_clock = policy.clock_debug(); + let decision = policy.consider_admission( + 1, + 1 << 20, + vec![], + CostSample { + cold_prefill_cost: nan, + restore_cost: 10.0, + }, + &[], + ); + assert_eq!(decision.verdict, AdmissionVerdict::Reject); + assert_eq!(policy.clock_debug(), before_clock); + assert!(policy.ghost(9).is_some()); + + // Invalid hit: same invariants. + policy.consider_admission(2, 1 << 20, vec![], cost(400.0, 100.0), &[]); + let before_clock = policy.clock_debug(); + assert!( + policy + .record_hit( + 2, + CostSample { + cold_prefill_cost: nan, + restore_cost: 10.0 + } + ) + .is_none() + ); + assert_eq!(policy.clock_debug(), before_clock); +} + +#[test] +fn many_reference_small_segment_still_charges_probation_bytes() { + // A 4-byte segment referenced by 9 probation entries: per-entry + // truncation would charge 0; the class total must still be 4. + let mut policy = BenefitPolicy::new(PolicyConfig::default()); + for key in 1..=9u64 { + policy.consider_admission(key, 0, vec![(1u64, 4)], cost(400.0, 100.0), &[]); + } + assert!( + policy.probation_bytes() >= 4, + "charged {}", + policy.probation_bytes() + ); + // And with a 1-byte cap, admission must select shared-only victims. + let mut policy = BenefitPolicy::new(PolicyConfig { + probation_byte_budget: 0, + ..PolicyConfig::default() + }); + let decision = policy.consider_admission(1, 0, vec![(1u64, 4)], cost(400.0, 100.0), &[]); + assert!(!decision.probation_cap_repair.victims().is_empty()); +} + +#[test] +fn ghost_promoted_recurrence_returns_admit_persist() { + let mut policy = BenefitPolicy::new(PolicyConfig::default()); + policy.consider_admission(1, 1 << 20, vec![], cost(400.0, 100.0), &[]); + assert!(policy.record_hit(1, cost(400.0, 100.0)).is_none()); // hits = 1 + policy.remove(1, &[]); + let decision = policy.consider_admission(1, 1 << 20, vec![], cost(400.0, 100.0), &[]); + assert_eq!(decision.kind, AdmissionDecisionKind::AdmitPersist); + assert_eq!(policy.entry(1).unwrap().state, PolicyEntryState::Admitted); + // Non-promoted recurrence stays AdmitProbation/Probation. + policy.consider_admission(2, 1 << 20, vec![], cost(400.0, 100.0), &[]); + policy.remove(2, &[]); + let decision = policy.consider_admission(2, 1 << 20, vec![], cost(400.0, 100.0), &[]); + assert_eq!(decision.kind, AdmissionDecisionKind::AdmitProbation); + assert_eq!(policy.entry(2).unwrap().state, PolicyEntryState::Probation); +} + +#[test] +fn zero_ghost_capacity_is_a_real_zero_bound() { + let mut policy = BenefitPolicy::new(PolicyConfig { + ghost_capacity: 0, + ..PolicyConfig::default() + }); + assert!( + PolicyConfig { + ghost_capacity: 0, + ..PolicyConfig::default() + } + .is_valid() + ); + policy.consider_admission(1, 1 << 20, vec![], cost(400.0, 100.0), &[]); + policy.remove(1, &[]); + assert_eq!(policy.ghost_count(), 0, "zero capacity must retain nothing"); +} + +#[test] +fn cap_victim_selection_covers_recomputed_shared_shares() { + // Removing shared references raises survivors' shares: the selected set + // must actually bring the class under cap once committed. + let mut policy = BenefitPolicy::new(PolicyConfig { + probation_byte_budget: 9 << 20, + grace_observations: 0, + ..PolicyConfig::default() + }); + // Three no-hit probationers, each 1 MiB exclusive + a shared 9 MiB + // segment (3 MiB share each → class charge 3*(1+3) = 12 MiB > 9 MiB). + for key in 1..=3u64 { + let decision = + policy.consider_admission(key, 1 << 20, vec![(7u64, 9 << 20)], cost(400.0, 100.0), &[]); + for v in decision.probation_cap_repair.victims().iter().copied() { + policy.remove(v, &[]); + } + } + assert!( + policy.probation_bytes() <= 9 << 20, + "committed class charge {} over cap", + policy.probation_bytes() + ); + // Removal must have happened through the decision path. + assert!(policy.len() < 3); +} + +#[test] +fn duplicate_segment_references_are_rejected() { + let mut policy = BenefitPolicy::new(PolicyConfig::default()); + let decision = policy.consider_admission( + 1, + 0, + vec![(42u64, 100), (42u64, 100)], + cost(400.0, 100.0), + &[], + ); + assert_eq!(decision.verdict, AdmissionVerdict::Reject); + assert!( + decision + .reasons + .contains(&"duplicate-segment-reference".to_string()) + ); + assert!(policy.is_empty()); + // Conflicting size for a known segment is also rejected. + let mut policy = BenefitPolicy::new(PolicyConfig::default()); + policy.consider_admission(1, 0, vec![(42u64, 100)], cost(400.0, 100.0), &[]); + let decision = policy.consider_admission(2, 0, vec![(42u64, 200)], cost(400.0, 100.0), &[]); + assert_eq!(decision.verdict, AdmissionVerdict::Reject); + assert!( + decision + .reasons + .contains(&"segment-size-conflict".to_string()) + ); + // Same size for a known segment is fine. + let decision = policy.consider_admission(3, 0, vec![(42u64, 100)], cost(400.0, 100.0), &[]); + assert_eq!(decision.verdict, AdmissionVerdict::Admit); + // One 100-byte segment: class charge is exactly 100. + assert_eq!(policy.probation_bytes(), 100); +} + +#[test] +fn cap_victim_selection_reproduces_reported_counterexample() { + // Exact shape of the reported probe: before=10,590,618 over + // cap=9,437,184 (9 MiB); stale-share subtraction selected [1,2] and + // left the committed class at 10,354,688. Reproduce it with concrete + // numbers: 10 MiB cap basis scaled to 9 MiB via three probationers — + // 1 MiB exclusive each (3 MiB) plus one shared 8 MiB segment + // (8/3 MiB per share -> class ~3+8=11 MiB before, over cap). + let mut policy = BenefitPolicy::new(PolicyConfig { + probation_byte_budget: 9 << 20, + grace_observations: 0, + ..PolicyConfig::default() + }); + // Admission applies cap victims incrementally, so build the over-cap + // state with cap victims disabled (huge probation budget), then swap + // in the real cap and select. + for key in 1..=3u64 { + policy.consider_admission(key, 1 << 20, vec![(7u64, 8 << 20)], cost(400.0, 100.0), &[]); + } + let before = policy.probation_bytes(); + assert!(before > 9 << 20, "before {}", before); + // Now select against the real cap by constructing the over-cap state + // through the public API: reset the budget by direct selection. + let victims = policy + .with_probation_budget(9 << 20, |p| p.select_probation_cap_victims(&[])) + .victims() + .to_vec(); + assert!(!victims.is_empty()); + for key in &victims { + policy.remove(*key, &[]); + } + let after = policy.probation_bytes(); + assert!( + after <= 9 << 20, + "committed class {} still over cap after victims {:?}", + after, + victims + ); +} + +#[test] +fn conflict_on_a_later_segment_leaves_state_untouched() { + // Admit segment 1 at size 100, then offer a new key whose later segment + // conflicts: [(2,100),(1,200)]. The reject must leave the ledger without + // segment 2, no entry for the rejected key, and any matching ghost intact. + let mut policy = BenefitPolicy::new(PolicyConfig::default()); + policy.consider_admission(1, 0, vec![(1u64, 100)], cost(400.0, 100.0), &[]); + + // Seed a ghost for the key that will be rejected so ghost survival is + // observable. + policy.consider_admission(9, 1 << 20, vec![], cost(400.0, 100.0), &[]); + policy.remove(9, &[]); + assert!(policy.ghost(9).is_some()); + + let decision = policy.consider_admission( + 9, + 1 << 20, + vec![(2u64, 100), (1u64, 200)], + cost(400.0, 100.0), + &[], + ); + assert_eq!(decision.verdict, AdmissionVerdict::Reject); + assert!( + decision + .reasons + .contains(&"segment-size-conflict".to_string()) + ); + // No entry for the rejected key; the ghost survived the reject. + assert!(policy.entry(9).is_none()); + assert!( + policy.ghost(9).is_some(), + "rejected re-offer must not lose the ghost" + ); + // Segment 2 was never registered. + assert!(policy.segments.segment_record(2).is_none()); + // Segment 1 still has exactly one reference (the original entry). + let record = policy.segments.segment_record(1).expect("segment 1 intact"); + assert_eq!(record.references, std::collections::BTreeSet::from([1u64])); + assert_eq!(record.size, 100); + // Class charge unchanged: the original entry's shared 100 bytes only. + assert_eq!(policy.probation_bytes(), 100); +} + +#[test] +fn structural_rejects_leave_clock_and_zero_horizon_ghosts_untouched() { + // ghost_max_age_observations = 0: any real observation would expire the + // ghost. Duplicate-ID and size-conflict rejects must not. + let mk = || { + BenefitPolicy::new(PolicyConfig { + ghost_max_age_observations: 0, + grace_observations: 4, + ..PolicyConfig::default() + }) + }; + + // Duplicate segment IDs. + let mut policy = mk(); + policy.consider_admission(9, 1 << 20, vec![], cost(400.0, 100.0), &[]); + policy.remove(9, &[]); + assert!(policy.ghost(9).is_some()); + let clock = policy.clock_debug(); + let decision = policy.consider_admission( + 1, + 0, + vec![(42u64, 100), (42u64, 100)], + cost(400.0, 100.0), + &[], + ); + assert_eq!(decision.verdict, AdmissionVerdict::Reject); + assert_eq!( + policy.clock_debug(), + clock, + "duplicate reject must not advance clock" + ); + assert!( + policy.ghost(9).is_some(), + "duplicate reject must not expire ghosts" + ); + + // Size conflict on a later segment. + let mut policy = mk(); + policy.consider_admission(1, 0, vec![(1u64, 100)], cost(400.0, 100.0), &[]); + policy.consider_admission(9, 1 << 20, vec![], cost(400.0, 100.0), &[]); + policy.remove(9, &[]); + assert!(policy.ghost(9).is_some()); + let clock = policy.clock_debug(); + let decision = policy.consider_admission( + 9, + 1 << 20, + vec![(2u64, 100), (1u64, 200)], + cost(400.0, 100.0), + &[], + ); + assert_eq!(decision.verdict, AdmissionVerdict::Reject); + assert_eq!( + policy.clock_debug(), + clock, + "conflict reject must not advance clock" + ); + assert!( + policy.ghost(9).is_some(), + "conflict reject must not expire ghosts" + ); +} + +#[test] +fn already_resident_key_rejection_leaves_no_stale_references() { + let mut policy = BenefitPolicy::new(PolicyConfig::default()); + policy.consider_admission(1, 0, vec![(1u64, 100)], cost(400.0, 100.0), &[]); + // Re-admitting a resident key is rejected and must not swap the entry's + // segment list or leak old ledger references. + let decision = policy.consider_admission(1, 0, vec![(2u64, 50)], cost(400.0, 100.0), &[]); + assert_eq!(decision.verdict, AdmissionVerdict::Reject); + assert!( + decision + .reasons + .contains(&"already-resident-key".to_string()) + ); + let record = policy.segments.segment_record(1).expect("segment 1 intact"); + assert_eq!(record.references, std::collections::BTreeSet::from([1u64])); + assert!( + policy.segments.segment_record(2).is_none(), + "segment 2 must not be registered" + ); + assert_eq!(policy.probation_bytes(), 100); +} + +#[test] +fn admitted_coreference_removal_repairs_the_probation_cap() { + // Ghost-promoted admitted A on shared segment S; probation P shares S + // (half-share fits the cap); one hit on P (still probation at + // threshold 2); remove A -> P's charge rises to all of S and exceeds + // the cap. The removal response must carry P as a cap victim. + let mut policy = BenefitPolicy::new(PolicyConfig { + probation_byte_budget: 1 << 20, // 1 MiB cap; S is 1.5 MiB + grace_observations: 0, + ..PolicyConfig::default() + }); + // Build A as ghost-promoted: admit, one hit, evict (ghost), recur. + policy.consider_admission(1, 0, vec![(7u64, (3 << 20) / 2)], cost(400.0, 100.0), &[]); + policy.record_hit(1, cost(400.0, 100.0)); + policy.remove_without_cap_repair(1); + let decision = + policy.consider_admission(1, 0, vec![(7u64, (3 << 20) / 2)], cost(400.0, 100.0), &[]); + assert_eq!(decision.kind, AdmissionDecisionKind::AdmitPersist); + for v in decision.probation_cap_repair.victims().iter().copied() { + policy.remove_without_cap_repair(v); + } + assert_eq!(policy.entry(1).unwrap().state, PolicyEntryState::Admitted); + + // P shares S: half-share = 0.75 MiB fits the 1 MiB cap. + let decision = + policy.consider_admission(2, 0, vec![(7u64, (3 << 20) / 2)], cost(400.0, 100.0), &[]); + assert_eq!(decision.verdict, AdmissionVerdict::Admit); + for v in decision.probation_cap_repair.victims().iter().copied() { + policy.remove_without_cap_repair(v); + } + // One hit on P: still probation at threshold 2. + policy.record_hit(2, cost(400.0, 100.0)); + assert_eq!(policy.entry(2).unwrap().state, PolicyEntryState::Probation); + assert!( + policy.probation_bytes() <= 1 << 20, + "half-share state over cap: {}", + policy.probation_bytes() + ); + + // Remove the admitted co-reference: P's charge rises to all of S. + let outcome = policy.remove(1, &[]).expect("A removed"); + assert!( + !outcome.probation_cap_repair.victims().is_empty(), + "removal must carry cap victims for the risen share" + ); + for v in outcome.probation_cap_repair.victims() { + policy.remove_without_cap_repair(*v); + } + assert!( + policy.probation_bytes() <= 1 << 20, + "class {} still over cap after repair", + policy.probation_bytes() + ); + assert!(policy.entry(2).is_none(), "P must be the repair victim"); +} + +#[test] +fn cap_repair_prefers_unpinned_over_older_pinned_probationer() { + // Older pinned probationer (key 1) must never be selected while the + // younger unpinned probationer (key 2) can repair the cap (#1650). + let mut policy = BenefitPolicy::new(PolicyConfig { + probation_byte_budget: 1 << 20, + ..PolicyConfig::default() + }); + for key in 1..=2u64 { + let decision = policy.consider_admission(key, 1 << 20, vec![], cost(400.0, 100.0), &[1]); + // Do not commit yet: build the over-cap state first. + assert!( + !decision.probation_cap_repair.victims().contains(&1), + "admission never selects a pin" + ); + } + let repair = policy.select_probation_cap_victims(&[1]); + assert!( + !repair.victims().contains(&1), + "pinned key selected as victim: {:?}", + repair.victims() + ); + assert!(repair.victims().contains(&2)); + assert!(!repair.is_deferred()); + for v in repair.victims().iter().copied() { + policy.remove(v, &[1]); + } + assert!(policy.probation_bytes() <= 1 << 20); + assert!(policy.entries.contains_key(&1), "pin must survive repair"); +} + +#[test] +fn all_pinned_cap_is_deferred_with_shortfall_never_a_pin() { + let mut policy = BenefitPolicy::new(PolicyConfig { + probation_byte_budget: 1 << 20, + ..PolicyConfig::default() + }); + let pinned = [1u64, 2]; + for (n, key) in pinned.iter().enumerate() { + let decision = + policy.consider_admission(*key, 1 << 20, vec![], cost(400.0, 100.0), &pinned); + let repair = &decision.probation_cap_repair; + assert!( + repair.victims().is_empty(), + "all candidates pinned: no victim may be selected" + ); + // First admission is under cap; the second pushes the all-pinned + // class over cap with no selectable candidate. + if n == 1 { + assert!( + repair.is_deferred(), + "unsatisfiable cap must be Deferred, got {:?}", + repair + ); + assert_eq!(repair.shortfall_bytes(), 1 << 20); + } + assert!( + !repair.victims().contains(&1) && !repair.victims().contains(&2), + "pins never selected even when cap is unsatisfiable" + ); + } + // The removal path honors pins identically: a caller-forced removal of + // the pinned key 1 leaves key 2 at exactly the cap — satisfied, no pin + // selected. + let outcome = policy.remove(1, &pinned).expect("entry 1 removed"); + assert!(!outcome.probation_cap_repair.is_deferred()); + assert!(outcome.probation_cap_repair.victims().is_empty()); + assert!(policy.entries.contains_key(&2)); + // Once the pin releases, repair becomes satisfiable again (key 2 is + // selectable again). + let repair = policy.select_probation_cap_victims(&[]); + assert!(!repair.is_deferred()); +} + +/// Grace must expire under repeated misses: a miss is a real observation +/// that advances the clock but is a negative value signal (no recency +/// refresh), so a miss-only stream ages an entry out of grace and makes +/// it evictable instead of holding it indefinitely. +#[test] +fn grace_expires_under_repeated_misses() { + let grace = 8u64; + let mut policy = BenefitPolicy::new(PolicyConfig { + grace_observations: grace, + ..PolicyConfig::default() + }); + policy.consider_admission(1, 1 << 20, vec![], cost(400.0, 100.0), &[]); + policy.consider_admission(2, 1 << 20, vec![], cost(400.0, 100.0), &[]); + // Miss key 1 seven times (each miss advances the clock): the clock at + // admission was 2, so key 1's age becomes 8 (out of grace) while key + // 2's age becomes 7 (still in grace) — exactly one observation on + // either side of the grace boundary, not an ordering artifact. + let clock_at_admission = policy.clock_debug(); + for _ in 0..grace - 1 { + policy.record_miss(1); + } + assert!( + policy.clock_debug() == clock_at_admission + grace - 1, + "misses must advance the clock" + ); + let e1 = policy.entry(1).unwrap(); + let e2 = policy.entry(2).unwrap(); + assert_eq!( + policy.clock_debug().saturating_sub(e1.last_observation), + grace, + "key 1 exactly out of grace" + ); + assert_eq!( + policy.clock_debug().saturating_sub(e2.last_observation), + grace - 1, + "key 2 exactly one observation inside grace" + ); + // Out-of-grace entry 1 is now evictable in pass 1 (grace honored); + // requesting only its bytes must never touch in-grace entry 2. + let victims = policy.choose_victims(1 << 20, &[]); + let evictable: Vec = victims.iter().map(|(k, _)| *k).collect(); + assert!( + evictable.contains(&1), + "miss-only entry must age out of grace, victims {:?}", + evictable + ); + assert!( + !evictable.contains(&2), + "in-grace entry must not be evicted first, victims {:?}", + evictable + ); + // A miss-only stream cannot hold grace forever — already asserted above: + // entry 1's age is exactly `grace` (evictable) after only misses. +} diff --git a/crates/skippy-cache/src/policy/traces.rs b/crates/skippy-cache/src/policy/traces.rs new file mode 100644 index 0000000000..e4d8ca0c0b --- /dev/null +++ b/crates/skippy-cache/src/policy/traces.rs @@ -0,0 +1,137 @@ +//! Deterministic synthetic traces for policy comparison (#1650 first slice). +//! Seeded xorshift so every run replays identically. + +pub struct Rng(u64); + +impl Rng { + pub fn new(seed: u64) -> Self { + Self(seed.max(1)) + } + pub fn next_u64(&mut self) -> u64 { + let mut x = self.0; + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + self.0 = x; + x + } + pub fn next_f64(&mut self) -> f64 { + (self.next_u64() >> 11) as f64 / (1u64 << 53) as f64 + } +} + +/// One cache access in a trace: which entry, and how valuable its reuse is. +#[derive(Debug, Clone, Copy)] +pub struct TraceAccess { + pub entry: u64, + pub cold_prefill_cost: f64, + pub restore_cost: f64, + pub exclusive_bytes: u64, +} + +pub const HOT_ZIPF_ENTRIES: u64 = 64; + +/// Zipf-like hotset: entry i is accessed with probability proportional to +/// 1/(i+1), so a small hotset dominates while a long tail streams once. +pub fn zipf_hotset_trace(seed: u64, len: usize) -> Vec { + let mut rng = Rng::new(seed); + let weights: Vec = (0..HOT_ZIPF_ENTRIES) + .map(|i| 1.0 / (i as f64 + 1.0)) + .collect(); + let total: f64 = weights.iter().sum(); + let mut out = Vec::with_capacity(len); + for _ in 0..len { + let pick = rng.next_f64() * total; + let mut entry = 0; + let mut acc = 0.0; + for (i, w) in weights.iter().enumerate() { + acc += w; + if pick <= acc { + entry = i as u64; + break; + } + } + out.push(access_for(entry, &mut rng)); + } + out +} + +/// Turn growth: each "session" revisits its whole history, extending it by a +/// fresh entry — the realistic agentic chat shape. +pub fn turn_growth_trace(seed: u64, sessions: u64, turns: u64) -> Vec { + let mut rng = Rng::new(seed); + let mut out = Vec::new(); + for session in 0..sessions { + for turn in 0..turns { + for entry in 0..=turn { + out.push(access_for(session * turns + entry, &mut rng)); + } + } + } + out +} + +/// One-shot stream: every entry seen exactly once, occasionally interleaved +/// with a hot entry so the policy must not wreck the hotset. +pub fn one_shot_trace(seed: u64, len: usize) -> Vec { + let mut rng = Rng::new(seed); + let mut out = Vec::with_capacity(len); + for i in 0..len as u64 { + if i % 8 == 7 { + out.push(access_for(0, &mut rng)); // hot anchor + } + out.push(access_for(1_000_000 + i, &mut rng)); + } + out +} + +/// Mixed sizes: footprints span three orders of magnitude. +pub fn mixed_size_trace(seed: u64, len: usize) -> Vec { + let mut rng = Rng::new(seed); + let mut out = Vec::with_capacity(len); + for i in 0..len as u64 { + let class = (i / 32) % 3; + // Distinct key per size class: reusing one key would make the + // larger sizes hits on a 64 KiB resident entry, so the trace + // would never exercise mixed resident sizes or large-entry + // eviction. + let mut access = access_for(class * 32 + i % 32, &mut rng); + access.exclusive_bytes = match class { + 0 => 64 << 10, + 1 => 4 << 20, + _ => 256 << 20, + }; + out.push(access); + } + out +} + +fn access_for(entry: u64, rng: &mut Rng) -> TraceAccess { + let cold = 50.0 + rng.next_f64() * 400.0; + TraceAccess { + entry, + cold_prefill_cost: cold, + restore_cost: cold * 0.3, + exclusive_bytes: (1 + rng.next_u64() % 8) << 20, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn turn_growth_keys_do_not_collide_across_long_sessions() { + let turns = 1_001; + let trace = turn_growth_trace(7, 2, turns); + let first_session_len = (turns * (turns + 1) / 2) as usize; + let first_second_session_key = trace[first_session_len].entry; + + assert_eq!(first_second_session_key, turns); + assert!( + trace[..first_session_len] + .iter() + .all(|access| access.entry < first_second_session_key) + ); + } +} diff --git a/crates/skippy-cache/src/radix.rs b/crates/skippy-cache/src/radix.rs index 9386db599c..1a4255b8a7 100644 --- a/crates/skippy-cache/src/radix.rs +++ b/crates/skippy-cache/src/radix.rs @@ -1762,22 +1762,28 @@ mod tests { ); assert_eq!( stats.resident_logical_bytes, - resident.values().map(|entry| entry.bytes).sum(), + resident.values().map(|entry| entry.bytes).sum::(), "seed={seed:#x} step={step}" ); assert_eq!( stats.recurrent_logical_bytes, - recurrent.values().map(|entry| entry.bytes).sum(), + recurrent.values().map(|entry| entry.bytes).sum::(), "seed={seed:#x} step={step}" ); assert_eq!( stats.resident_active_refs, - resident.values().map(|entry| u64::from(entry.refs)).sum(), + resident + .values() + .map(|entry| u64::from(entry.refs)) + .sum::(), "seed={seed:#x} step={step}" ); assert_eq!( stats.recurrent_active_refs, - recurrent.values().map(|entry| u64::from(entry.refs)).sum(), + recurrent + .values() + .map(|entry| u64::from(entry.refs)) + .sum::(), "seed={seed:#x} step={step}" ); assert_eq!( diff --git a/crates/skippy-cache/src/source.rs b/crates/skippy-cache/src/source.rs new file mode 100644 index 0000000000..dd2c0fdfb5 --- /dev/null +++ b/crates/skippy-cache/src/source.rs @@ -0,0 +1,82 @@ +//! Read-only seam for future cache sources. +//! +//! Disk remains the writable L3 tier and owns admission. A network transport +//! can later implement these traits, verify the same manifest/segment format, +//! and commit fetched objects through the local manager before runtime fill. + +use anyhow::Result; + +use crate::l3::{HandoffManifest, HandoffSegmentStore}; + +pub trait ManifestSource: Send + Sync { + fn recorded_prefix_lengths(&self, namespace_key: &str) -> Result>; + + fn manifest_for_prefix( + &self, + namespace_key: &str, + token_len: u64, + prefix_key: &str, + ) -> Result>; + + fn load_manifest(&self, payload_digest: &str) -> Result>; +} + +pub trait SegmentSource: Send + Sync { + fn read_segment(&self, digest: &str) -> Result>>; +} + +impl ManifestSource for HandoffSegmentStore { + fn recorded_prefix_lengths(&self, namespace_key: &str) -> Result> { + HandoffSegmentStore::recorded_prefix_lengths(self, namespace_key) + } + + fn manifest_for_prefix( + &self, + namespace_key: &str, + token_len: u64, + prefix_key: &str, + ) -> Result> { + HandoffSegmentStore::manifest_for_prefix(self, namespace_key, token_len, prefix_key) + } + + fn load_manifest(&self, payload_digest: &str) -> Result> { + match HandoffSegmentStore::load_manifest(self, payload_digest) { + Ok(manifest) => Ok(Some(manifest)), + Err(error) + if error + .downcast_ref::() + .is_some_and(|error| error.kind() == std::io::ErrorKind::NotFound) => + { + Ok(None) + } + Err(error) => Err(error), + } + } +} + +impl SegmentSource for HandoffSegmentStore { + fn read_segment(&self, digest: &str) -> Result>> { + match HandoffSegmentStore::read_segment(self, digest) { + Ok(bytes) => Ok(Some(bytes)), + Err(error) + if error + .downcast_ref::() + .is_some_and(|error| error.kind() == std::io::ErrorKind::NotFound) => + { + Ok(None) + } + Err(error) => Err(error), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn local_store_satisfies_future_read_source_contracts() { + fn assert_sources() {} + assert_sources::(); + } +} diff --git a/crates/skippy-cache/src/tier.rs b/crates/skippy-cache/src/tier.rs new file mode 100644 index 0000000000..589585cb55 --- /dev/null +++ b/crates/skippy-cache/src/tier.rs @@ -0,0 +1,1407 @@ +//! The L3 tier under the radix cache. +//! +//! `UnifiedRadixCache` holds `ExactStatePayload` entries in RAM (L1/L2); +//! this tier gives them a durable floor: spill a payload under its prefix +//! identity when the radix cache evicts it, and fill it back from local disk +//! on a radix miss. A later peer source may supply the same manifest/segment +//! format, but it must commit through the local manager before this tier can +//! fill it. State never crosses a numerical-identity boundary: spills stamp +//! both model and exact-state identities, and fills require both. + +use std::sync::atomic::Ordering; + +use anyhow::{Context, Result, bail}; +use serde::Serialize; + +use crate::l3::{ + HandoffManifest, HandoffSegmentRef, HandoffSegmentStore, MANIFEST_VERSION, PayloadGeometry, + SegmentCodecIdentity, StoreLimits, StoreUsage, segment_digest, +}; +use crate::manager::{L3ActivitySnapshot, L3CacheManager, L3EffectiveStatus}; +use crate::payload::{ExactStatePayload, ExactStatePayloadKind}; + +/// Everything the status contract needs from the tier, in one read. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct L3Status { + pub model_identity: String, + pub state_identity: String, + pub effective: L3EffectiveStatus, + pub format_version: u32, + pub usage: StoreUsage, + /// Manifests stamped with this tier's identity: what a restart can reuse. + pub restorable_manifests: u64, + pub restorable_tokens: u64, + pub namespaces: u64, + pub activity: L3ActivitySnapshot, +} + +/// Key an L3 entry by the radix coordinates that identify it in RAM: the +/// namespace (which already binds the numerical stage identity) and the +/// exact token path. +pub fn l3_prefix_key(namespace: &str, token_ids: &[i32]) -> String { + let mut hasher = blake3::Hasher::new(); + hasher.update(b"l3-prefix-key-v1"); + hasher.update(namespace.as_bytes()); + for token_id in token_ids { + hasher.update(&token_id.to_le_bytes()); + } + format!("blake3:{}", hasher.finalize().to_hex()) +} + +/// The namespace's own index key. +pub fn l3_namespace_key(namespace: &str) -> String { + let mut hasher = blake3::Hasher::new(); + hasher.update(b"l3-namespace-key-v1"); + hasher.update(namespace.as_bytes()); + format!("blake3:{}", hasher.finalize().to_hex()) +} + +/// A located entry: the cheap index-probe result, addressing exactly one +/// recorded manifest. Splitting locate from load lets callers single-flight +/// the expensive load on the entry itself (namespace + recorded length + +/// manifest key) rather than on the query's shape. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct L3Location { + pub namespace_key: String, + pub prefix_key: String, + pub token_count: u64, + /// Manifest key (payload digest) — the identity of the physical entry; + /// the correct single-flight claim key. + pub manifest_key: String, + /// Metadata copied from the located manifest so the runtime can reject an + /// incompatible native page descriptor before reading segment bytes. + pub kv_desc_json: Option, + pub kv_bytes: u64, + pub native_kv_passthrough: bool, +} + +/// A successful fill from the tier. +pub struct L3Fill { + pub payload: ExactStatePayload, + /// How many of the query's leading tokens the filled state covers — the + /// length the entry was recorded at, which may be shorter than the + /// query (longest-recorded-prefix semantics, mirroring the radix). + pub token_count: u64, + pub kv_desc_json: Option, + pub payload_bytes: u64, +} + +pub struct L3Tier { + manager: L3CacheManager, + model_identity: String, + state_identity: String, + segment_bytes: usize, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum SegmentRepresentation { + Raw, + NativeKvPage, +} + +impl SegmentRepresentation { + fn identity(self, encoded_len: u64) -> SegmentCodecIdentity { + match self { + Self::Raw => SegmentCodecIdentity::raw(encoded_len), + Self::NativeKvPage => SegmentCodecIdentity::native_kv_page(encoded_len), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct SegmentCut { + offset: u64, + len: u64, + label: String, + representation: SegmentRepresentation, +} + +fn append_fixed_cuts( + cuts: &mut Vec, + offset: &mut u64, + bytes: u64, + segment_bytes: u64, + representation: SegmentRepresentation, +) { + let mut remaining = bytes; + while remaining > 0 { + let len = segment_bytes.max(1).min(remaining); + cuts.push(SegmentCut { + offset: *offset, + len, + label: String::new(), + representation, + }); + *offset = offset.saturating_add(len); + remaining -= len; + } +} + +impl L3Tier { + /// Open a tier capped by `budget_bytes` with no free-space reserve. + /// Prefer [`Self::open_with_limits`]. + pub fn open( + root: impl Into, + budget_bytes: u64, + state_identity: String, + segment_bytes: usize, + ) -> Result { + Self::open_with_limits( + root, + StoreLimits::new(budget_bytes, 0), + state_identity, + segment_bytes, + ) + } + + pub fn open_with_limits( + root: impl Into, + limits: StoreLimits, + state_identity: String, + segment_bytes: usize, + ) -> Result { + let manager = L3CacheManager::acquire(root.into(), limits)?; + Ok(manager.tier(state_identity, segment_bytes)) + } + + pub fn open_with_identities( + root: impl Into, + limits: StoreLimits, + model_identity: String, + state_identity: String, + segment_bytes: usize, + ) -> Result { + let manager = L3CacheManager::acquire(root.into(), limits)?; + Ok(manager.tier_for_model(model_identity, state_identity, segment_bytes)) + } + + pub(crate) fn from_manager( + manager: L3CacheManager, + model_identity: String, + state_identity: String, + segment_bytes: usize, + ) -> Self { + Self { + manager, + model_identity, + state_identity, + segment_bytes: segment_bytes.max(1), + } + } + + pub fn store(&self) -> &HandoffSegmentStore { + self.manager.store() + } + + pub fn manager(&self) -> &L3CacheManager { + &self.manager + } + + pub fn state_identity(&self) -> &str { + &self.state_identity + } + + pub fn model_identity(&self) -> &str { + &self.model_identity + } + + /// Point-in-time activity counters. + pub fn activity(&self) -> L3ActivitySnapshot { + self.manager.activity_snapshot() + } + + /// The full status snapshot. One pass over the manifests; safe to call + /// while serving, since it takes no lock a request path holds. + pub fn status(&self) -> Result { + let (restorable_manifests, restorable_tokens, _) = self.restorable_summary()?; + Ok(L3Status { + model_identity: self.model_identity.clone(), + state_identity: self.state_identity.clone(), + effective: self.manager.effective_status(), + format_version: MANIFEST_VERSION, + usage: self.store().usage()?, + restorable_manifests: restorable_manifests as u64, + restorable_tokens, + namespaces: self.store().namespace_count()?, + activity: self.activity(), + }) + } + + /// Spill a radix payload under its (namespace, token-path) coordinates. + /// Returns the manifest key. Entries at many lengths coexist — each is a + /// complete state for its own length, which is what longest-prefix fill + /// leans on. + /// + /// Zero-byte payloads are refused: a dense family whose native KV export + /// was unavailable would otherwise spill nothing and restore as a bare + /// position advance over missing state. + pub fn spill( + &self, + namespace: &str, + token_ids: &[i32], + payload: &ExactStatePayload, + kv_desc_json: Option, + geometry: Option<&PayloadGeometry>, + ) -> Result { + let _operation = self.manager.operation_guard(); + let result = self.spill_inner(namespace, token_ids, payload, kv_desc_json, geometry); + if let Err(error) = &result { + self.manager.activity_counters().record_error(error); + } + result + } + + fn spill_inner( + &self, + namespace: &str, + token_ids: &[i32], + payload: &ExactStatePayload, + kv_desc_json: Option, + geometry: Option<&PayloadGeometry>, + ) -> Result { + if payload.byte_len() == 0 { + bail!( + "refusing to spill an empty exact-state payload: no state component was exported" + ); + } + let token_count = token_ids.len() as u64; + let (kv, recurrent): (Vec, Vec) = match payload.kind() { + ExactStatePayloadKind::FullState => ( + payload + .full_state_bytes_timed() + .context("failed to reconstruct full state for spill")? + .0 + .into_owned(), + Vec::new(), + ), + ExactStatePayloadKind::RecurrentOnly => ( + Vec::new(), + payload + .recurrent_state_bytes() + .context("failed to reconstruct recurrent state for spill")? + .into_owned(), + ), + ExactStatePayloadKind::KvRecurrent => ( + { + let kv = payload + .kv_bytes() + .context("failed to reconstruct KV bytes for spill")? + .map(|bytes| bytes.into_owned()) + .unwrap_or_default(); + if kv.is_empty() { + bail!("refusing to spill kv-recurrent state without a KV component"); + } + kv + }, + payload + .recurrent_state_bytes() + .context("failed to reconstruct recurrent state for spill")? + .into_owned(), + ), + }; + + // For full-state and recurrent-only exactly one component is populated, + // and KV states reach gigabytes: concatenating would peak at twice the + // payload for no benefit. Only a genuine composite needs the copy. + let (kv_bytes, recurrent_bytes) = (kv.len() as u64, recurrent.len() as u64); + let wire = if recurrent.is_empty() { + kv + } else if kv.is_empty() { + recurrent + } else { + let mut wire = Vec::with_capacity(kv.len() + recurrent.len()); + wire.extend_from_slice(&kv); + wire.extend_from_slice(&recurrent); + wire + }; + let payload_digest = segment_digest(&wire); + + let mut manifest = HandoffManifest::new_for_model( + self.model_identity.clone(), + self.state_identity.clone(), + payload.kind().as_str().to_string(), + ); + manifest.total_bytes = wire.len() as u64; + manifest.payload_digest = payload_digest.clone(); + manifest.kv_bytes = kv_bytes; + manifest.recurrent_bytes = recurrent_bytes; + let native_kv_passthrough = payload.kind() == ExactStatePayloadKind::KvRecurrent + && kv_bytes > 0 + && kv_desc_json + .as_deref() + .is_some_and(|descriptor| !descriptor.trim().is_empty()); + manifest.kv_desc_json = kv_desc_json; + manifest.token_count = token_count; + // Cut on the payload's own geometry when the caller knows it, so a + // longer prefix reuses the segments of the shorter one it extends. A + // geometry that does not describe these exact bytes is ignored rather + // than trusted: mis-cutting would still reassemble, but silently write + // the whole payload again every turn. + let geometry = geometry.filter(|geometry| { + let geometry_kv_bytes = geometry.total_bytes().saturating_sub(geometry.tail_bytes); + let matches = geometry.matches(wire.len() as u64) + && (!native_kv_passthrough || geometry_kv_bytes == kv_bytes); + if !matches { + self.manager + .activity_counters() + .geometry_rejected + .fetch_add(1, Ordering::Relaxed); + } + matches + }); + let cuts = match geometry { + Some(geometry) => geometry + .plan(self.segment_bytes as u64) + .into_iter() + .map(|(offset, len, label)| SegmentCut { + offset, + len, + representation: if native_kv_passthrough && offset < kv_bytes { + SegmentRepresentation::NativeKvPage + } else { + SegmentRepresentation::Raw + }, + label, + }) + .collect(), + None => { + let mut cuts = Vec::new(); + let mut offset = 0u64; + if native_kv_passthrough { + append_fixed_cuts( + &mut cuts, + &mut offset, + kv_bytes, + self.segment_bytes as u64, + SegmentRepresentation::NativeKvPage, + ); + append_fixed_cuts( + &mut cuts, + &mut offset, + recurrent_bytes, + self.segment_bytes as u64, + SegmentRepresentation::Raw, + ); + } else { + append_fixed_cuts( + &mut cuts, + &mut offset, + wire.len() as u64, + self.segment_bytes as u64, + SegmentRepresentation::Raw, + ); + } + cuts + } + }; + let segment_slices = cuts + .iter() + .map(|cut| { + let start = usize::try_from(cut.offset).context("segment offset exceeds usize")?; + let end = start + .checked_add(usize::try_from(cut.len).context("segment length exceeds usize")?) + .context("segment range overflows")?; + Ok(&wire[start..end]) + }) + .collect::>>()?; + let stored_segments = match self.store().try_put_segments(&segment_slices) { + Ok(Ok(stored)) => stored, + Ok(Err(refusal)) => { + self.manager.record_write_refusal(refusal); + bail!("cannot store packed segments: {}", refusal.reason()); + } + Err(error) => { + self.manager.record_storage_error(); + return Err(error); + } + }; + let mut new_bytes = 0u64; + // Held until after the commit below: until the manifest names them + // these segments are unreferenced, and an eviction triggered by + // another writer would collect them mid-build. + let mut held = Vec::with_capacity(stored_segments.len()); + for ((index, cut), stored) in cuts.into_iter().enumerate().zip(stored_segments) { + if stored.put.new { + new_bytes = new_bytes.saturating_add(stored.put.bytes); + } + manifest.segments.push(HandoffSegmentRef { + index: index as u32, + offset: cut.offset, + bytes: cut.len, + digest: stored.digest.clone(), + codec_identity: Some(cut.representation.identity(cut.len)), + meta_json: (!cut.label.is_empty()).then_some(cut.label), + }); + held.push(stored); + } + // Pin before publishing the manifest so another stage cannot evict + // it in the gap between commit and prefix-link publication. + let _manifest_pin = self.store().pin(&payload_digest); + match self.store().try_commit(&manifest) { + Ok(Ok(())) => {} + Ok(Err(refusal)) => { + self.manager.record_write_refusal(refusal); + bail!( + "cannot commit manifest {}: {}", + manifest.payload_digest, + refusal.reason() + ); + } + Err(error) => { + self.manager.record_storage_error(); + return Err(error); + } + } + if let Err(error) = self.store().link_prefix( + &l3_namespace_key(namespace), + token_count, + &l3_prefix_key(namespace, token_ids), + &payload_digest, + ) { + self.manager.record_storage_error(); + return Err(error); + } + drop(held); + self.manager + .activity_counters() + .writes + .fetch_add(1, Ordering::Relaxed); + self.manager + .activity_counters() + .bytes_written + .fetch_add(new_bytes, Ordering::Relaxed); + self.manager.record_successful_write(); + Ok(payload_digest) + } + + /// Locate the longest recorded prefix of the query, mirroring the radix + /// cache's longest-component-prefix semantics: probe recorded lengths + /// for this namespace from longest to shortest (capped at `max_probes`), + /// hashing the query's own leading tokens at each length — so a + /// recorded entry only matches when the query genuinely starts with the + /// tokens it was recorded for. Index probes only; the expensive load is + /// `load`, so callers can single-flight on the returned entry. + pub fn locate_longest( + &self, + namespace: &str, + token_ids: &[i32], + max_probes: usize, + ) -> Result> { + let result = self.locate_longest_inner(namespace, token_ids, max_probes); + match &result { + Ok(Some(_)) => { + self.manager + .activity_counters() + .hits + .fetch_add(1, Ordering::Relaxed); + } + Ok(None) => { + self.manager + .activity_counters() + .misses + .fetch_add(1, Ordering::Relaxed); + } + Err(error) => self.manager.activity_counters().record_error(error), + } + result + } + + fn locate_longest_inner( + &self, + namespace: &str, + token_ids: &[i32], + max_probes: usize, + ) -> Result> { + let namespace_key = l3_namespace_key(namespace); + let query_len = token_ids.len() as u64; + let lengths = self.store().recorded_prefix_lengths(&namespace_key)?; + for length in lengths + .into_iter() + .filter(|length| *length > 0 && *length <= query_len) + .take(max_probes.max(1)) + { + let prefix_key = l3_prefix_key(namespace, &token_ids[..length as usize]); + let Some(manifest) = + self.store() + .manifest_for_prefix(&namespace_key, length, &prefix_key)? + else { + continue; + }; + if manifest.model_identity != self.model_identity + || manifest.state_identity != self.state_identity + { + bail!( + "L3 entry for this prefix was spilled under state identity {} but the tier serves {}", + manifest.state_identity, + self.state_identity + ); + } + if manifest.token_count != length { + bail!( + "L3 index length {length} disagrees with manifest token count {}", + manifest.token_count + ); + } + let native_kv_passthrough = manifest.uses_native_kv_passthrough(); + return Ok(Some(L3Location { + namespace_key, + prefix_key, + token_count: length, + manifest_key: manifest.payload_digest, + kv_desc_json: manifest.kv_desc_json.clone(), + kv_bytes: manifest.kv_bytes, + native_kv_passthrough, + })); + } + Ok(None) + } + + /// Load a located entry. Memory bound: `assemble` materializes the + /// payload once (`total_bytes`); the kv/recurrent split below reuses + /// that allocation via `split_off`, so peak extra memory is the payload + /// itself. Full-state fills are whole-blob by nature; kv-recurrent + /// fills could stream per segment later. + pub fn load(&self, location: &L3Location) -> Result { + let Some(_operation) = self.manager.try_operation_guard() else { + bail!("disk cache lifecycle operation in progress"); + }; + // Pinned for the whole load: eviction under a concurrent spill must + // not remove the segments this fill is assembling. + let _pin = self.store().pin(&location.manifest_key); + let result = self.load_inner(location); + match &result { + Ok(fill) => { + self.manager + .activity_counters() + .fills + .fetch_add(1, Ordering::Relaxed); + self.manager + .activity_counters() + .bytes_read + .fetch_add(fill.payload_bytes, Ordering::Relaxed); + } + Err(error) => self.manager.activity_counters().record_error(error), + } + result + } + + fn load_inner(&self, location: &L3Location) -> Result { + let manifest = self.store().load_manifest(&location.manifest_key)?; + if manifest.model_identity != self.model_identity + || manifest.state_identity != self.state_identity + { + bail!( + "L3 manifest {} was spilled under state identity {} but the tier serves {}", + location.manifest_key, + manifest.state_identity, + self.state_identity + ); + } + let mut wire = self.store().assemble(&manifest)?; + let payload_bytes = wire.len() as u64; + let kv_bytes = usize::try_from(manifest.kv_bytes).context("kv bytes exceed usize")?; + let payload = match manifest.payload_kind.as_str() { + "full-state" => ExactStatePayload::full_state(wire), + "recurrent-only" => ExactStatePayload::recurrent_only(wire), + "kv-recurrent" => { + // kv_bytes is an independent manifest field: a corrupt or + // truncated one must be a recorded miss, not a panic on a + // serving thread. + if kv_bytes > wire.len() { + bail!( + "L3 manifest claims {kv_bytes} kv bytes but the payload holds {}", + wire.len() + ); + } + let recurrent = wire.split_off(kv_bytes); + ExactStatePayload::kv_recurrent(wire, recurrent) + } + other => bail!("L3 manifest holds unknown payload kind {other}"), + }; + Ok(L3Fill { + payload, + token_count: manifest.token_count, + kv_desc_json: manifest.kv_desc_json, + payload_bytes, + }) + } + + /// Locate and load in one step. + pub fn fill_longest( + &self, + namespace: &str, + token_ids: &[i32], + max_probes: usize, + ) -> Result> { + match self.locate_longest(namespace, token_ids, max_probes)? { + Some(location) => Ok(Some(self.load(&location)?)), + None => Ok(None), + } + } + + /// What the tier can restore right now: (manifest count, restorable + /// token total, segment footprint bytes). Startup visibility so warm + /// state is never invisible. + pub fn restorable_summary(&self) -> Result<(usize, u64, u64)> { + let keys = self.store().list_manifests()?; + let mut tokens = 0u64; + let mut count = 0usize; + for key in &keys { + if let Ok(manifest) = self.store().load_manifest(key) + && manifest.model_identity == self.model_identity + && manifest.state_identity == self.state_identity + { + tokens = tokens.saturating_add(manifest.token_count); + count += 1; + } + } + Ok((count, tokens, self.store().segment_footprint_bytes()?)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::l3::{ + CODEC_NATIVE_KV_PAGE, CODEC_NATIVE_KV_PAGE_VERSION, GeometryBlock, GeometryKind, + }; + + fn temp_root(name: &str) -> std::path::PathBuf { + let root = std::env::temp_dir() + .join("skippy-l3-tier-tests") + .join(format!("{name}-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&root); + root + } + + fn tier(name: &str, identity: &str) -> L3Tier { + L3Tier::open(temp_root(name), 0, identity.to_string(), 4096).expect("open tier") + } + + fn tokens(len: usize) -> Vec { + (0..len as i32).collect() + } + + /// Mirrors the runtime's export layout: every layer's K rows, then every + /// layer's V rows, each run holding one row per token in token order. + fn kv_geometry( + layers: u32, + rows: u64, + k_stride: u64, + v_stride: u64, + window: u64, + ) -> PayloadGeometry { + let mut blocks = Vec::new(); + for layer in 0..layers { + blocks.push(GeometryBlock { + stride: k_stride, + kind: GeometryKind::Key, + layer, + column: 0, + }); + } + for layer in 0..layers { + blocks.push(GeometryBlock { + stride: v_stride, + kind: GeometryKind::Value, + layer, + column: 0, + }); + } + PayloadGeometry { + blocks, + rows, + window_rows: window, + tail_bytes: 0, + } + } + + /// A growing prefix keeps earlier tokens' rows byte-identical, so the + /// export for turn N+1 is turn N's runs each extended in place. + fn kv_payload(layers: u32, rows: u64, k_stride: u64, v_stride: u64) -> ExactStatePayload { + let mut wire = Vec::new(); + for (kind, stride) in [(0u8, k_stride), (1u8, v_stride)] { + for layer in 0..layers { + for row in 0..rows { + // Stamp kind, layer and row into every row so no two runs + // can hash alike: content addressing would legitimately + // collapse them, and the test would be measuring that + // instead of the windowing. + for byte in 0..stride { + wire.push(match byte { + 0 => kind, + 1 => layer as u8, + 2 => row as u8, + 3 => (row >> 8) as u8, + _ => (byte as u8).wrapping_mul(7) ^ (row as u8).wrapping_mul(31), + }); + } + } + } + } + ExactStatePayload::kv_recurrent(wire, Vec::new()) + } + + fn runtime_kv_desc(kv_type: u32, token_count: u64, payload_bytes: u64) -> String { + serde_json::json!({ + "version": 1, + "layer_start": 0, + "layer_end": 1, + "token_start": 0, + "token_count": token_count, + "layer_count": 1, + "k_type": kv_type, + "v_type": kv_type, + "k_row_bytes": 16, + "v_row_bytes": 16, + "v_element_bytes": 2, + "k_idx_row_bytes": 0, + "payload_bytes": payload_bytes, + "flags": 0, + "codec": 0, + "component_count": 0 + }) + .to_string() + } + + fn assert_native_kv_identity(segment: &HandoffSegmentRef) { + let identity = segment.codec_identity.as_ref().expect("segment identity"); + assert_eq!(identity.name, CODEC_NATIVE_KV_PAGE); + assert_eq!(identity.version, CODEC_NATIVE_KV_PAGE_VERSION); + assert_eq!(identity.class, crate::l3::CodecClass::Exact); + assert_eq!(identity.decoded_len, segment.bytes); + assert_eq!(identity.calibration_digest, None); + } + + #[test] + fn native_passthrough_preserves_supported_runtime_kv_representations() { + // ggml type ids used by the runtime: F32, F16, Q8_0 and Q4_0. The + // store treats all four as opaque native bytes and never transcodes. + for (name, kv_type) in [("f32", 0u32), ("f16", 1), ("q8_0", 8), ("q4_0", 2)] { + let tier = tier(&format!("native-{name}"), "blake3:native"); + let kv = vec![kv_type as u8; 5_000]; + let desc = runtime_kv_desc(kv_type, 4, kv.len() as u64); + let digest = tier + .spill( + "ns", + &tokens(4), + &ExactStatePayload::kv_recurrent(kv.clone(), Vec::new()), + Some(desc.clone()), + None, + ) + .expect("spill native KV"); + let manifest = tier.store().load_manifest(&digest).expect("load manifest"); + assert!(manifest.uses_native_kv_passthrough()); + assert_eq!(manifest.kv_desc_json.as_deref(), Some(desc.as_str())); + assert!( + manifest + .segments + .iter() + .all(|segment| segment.offset + segment.bytes <= manifest.kv_bytes) + ); + for segment in &manifest.segments { + assert_native_kv_identity(segment); + } + + let fill = tier + .fill_longest("ns", &tokens(4), 8) + .expect("fill") + .expect("native entry"); + assert_eq!( + fill.payload.kv_bytes().unwrap().unwrap().as_ref(), + kv.as_slice(), + "{name} bytes changed during the storage round trip" + ); + } + } + + #[test] + fn fixed_native_and_recurrent_segments_do_not_cross_representation_boundary() { + let tier = tier("native-fixed-boundary", "blake3:native"); + let kv = vec![1u8; 5_000]; + let recurrent = vec![2u8; 5_000]; + let digest = tier + .spill( + "ns", + &tokens(4), + &ExactStatePayload::kv_recurrent(kv.clone(), recurrent.clone()), + Some(runtime_kv_desc(1, 4, kv.len() as u64)), + None, + ) + .expect("spill mixed exact state"); + let manifest = tier.store().load_manifest(&digest).expect("load manifest"); + assert_eq!(manifest.kv_bytes, 5_000); + assert_eq!(manifest.recurrent_bytes, 5_000); + assert_eq!( + manifest.segments[1].offset + manifest.segments[1].bytes, + 5_000 + ); + for segment in &manifest.segments { + let end = segment.offset + segment.bytes; + assert!(end <= manifest.kv_bytes || segment.offset >= manifest.kv_bytes); + if segment.offset < manifest.kv_bytes { + assert_native_kv_identity(segment); + } else { + assert_eq!( + segment.codec_identity, + Some(SegmentCodecIdentity::raw(segment.bytes)) + ); + } + } + let fill = tier + .fill_longest("ns", &tokens(4), 8) + .expect("fill") + .expect("mixed entry"); + assert_eq!(fill.payload.kv_bytes().unwrap().unwrap().as_ref(), kv); + assert_eq!( + fill.payload.recurrent_state_bytes().unwrap().as_ref(), + recurrent + ); + } + + #[test] + fn geometry_marks_kv_windows_native_and_recurrent_tail_raw() { + let tier = tier("native-geometry-boundary", "blake3:native"); + let geometry = PayloadGeometry { + tail_bytes: 50, + ..kv_geometry(1, 8, 16, 16, 4) + }; + let kv = vec![3u8; 256]; + let recurrent = vec![4u8; 50]; + let digest = tier + .spill( + "ns", + &tokens(8), + &ExactStatePayload::kv_recurrent(kv, recurrent), + Some(runtime_kv_desc(1, 8, 256)), + Some(&geometry), + ) + .expect("spill geometry-native state"); + let manifest = tier.store().load_manifest(&digest).expect("load manifest"); + for segment in &manifest.segments { + if segment.meta_json.as_deref() == Some("tail") { + assert_eq!( + segment.codec_identity, + Some(SegmentCodecIdentity::raw(segment.bytes)) + ); + assert!(segment.offset >= manifest.kv_bytes); + } else { + assert_native_kv_identity(segment); + assert!(segment.offset + segment.bytes <= manifest.kv_bytes); + } + } + } + + #[test] + fn geometry_cut_segments_survive_prefix_growth() { + // The measured failure: the export is layer-major, so a fixed byte cut + // lands in a different place every turn and nothing is reused. + let (layers, k_stride, v_stride, window) = (4u32, 64u64, 64u64, 8u64); + let tier = tier("geometry-growth", "blake3:geometry"); + + let first = kv_payload(layers, 32, k_stride, v_stride); + tier.spill( + "ns", + &tokens(32), + &first, + None, + Some(&kv_geometry(layers, 32, k_stride, v_stride, window)), + ) + .expect("spill 32"); + let after_first = tier.status().expect("status").activity.bytes_written; + assert_eq!( + after_first, + first.byte_len(), + "first spill writes everything" + ); + + // Eight more tokens: only the new rows are new bytes. + let second = kv_payload(layers, 40, k_stride, v_stride); + tier.spill( + "ns", + &tokens(40), + &second, + None, + Some(&kv_geometry(layers, 40, k_stride, v_stride, window)), + ) + .expect("spill 40"); + let status = tier.status().expect("status"); + let physical = status.activity.bytes_written - after_first; + let ideal = second.byte_len() - first.byte_len(); + assert_eq!(status.activity.geometry_rejected, 0); + assert_eq!( + physical, ideal, + "growth wrote {physical} bytes for {ideal} bytes of new state" + ); + + // And the entry still reassembles to exactly what was spilled. + let fill = tier + .fill_longest("ns", &tokens(40), 64) + .expect("fill") + .expect("entry must be loadable"); + assert_eq!( + fill.payload + .kv_bytes() + .expect("kv") + .expect("some") + .into_owned(), + second.kv_bytes().expect("kv").expect("some").into_owned() + ); + } + + #[test] + fn fixed_offset_cutting_is_what_fails_the_gate() { + // The same growth without geometry: this is the 8x the probe measured, + // kept as a test so the regression is visible rather than remembered. + let (layers, k_stride, v_stride) = (4u32, 64u64, 64u64); + let tier = tier("geometry-fixed", "blake3:geometry"); + let first = kv_payload(layers, 32, k_stride, v_stride); + tier.spill("ns", &tokens(32), &first, None, None) + .expect("spill 32"); + let after_first = tier.status().expect("status").activity.bytes_written; + let second = kv_payload(layers, 40, k_stride, v_stride); + tier.spill("ns", &tokens(40), &second, None, None) + .expect("spill 40"); + let physical = tier.status().expect("status").activity.bytes_written - after_first; + let ideal = second.byte_len() - first.byte_len(); + assert!( + physical > ideal * 2, + "fixed cutting unexpectedly deduped: {physical} vs {ideal}" + ); + } + + #[test] + fn a_geometry_that_does_not_describe_the_payload_is_ignored() { + let tier = tier("geometry-mismatch", "blake3:geometry"); + let payload = kv_payload(2, 16, 32, 32); + // Wrong row count: cutting to it would silently mis-window every turn. + let wrong = kv_geometry(2, 15, 32, 32, 4); + tier.spill("ns", &tokens(16), &payload, None, Some(&wrong)) + .expect("spill still succeeds"); + let status = tier.status().expect("status"); + assert_eq!(status.activity.geometry_rejected, 1); + let fill = tier + .fill_longest("ns", &tokens(16), 8) + .expect("fill") + .expect("entry must be loadable"); + assert_eq!(fill.payload.byte_len(), payload.byte_len()); + } + + #[test] + fn status_counters_reconcile_with_what_the_tier_did() { + let root = temp_root("status"); + let tier = L3Tier::open(&root, 0, "blake3:status".to_string(), 4096).unwrap(); + // Every 4096-byte segment distinct, so within-spill dedup does not + // hide what the duplicate spill below is meant to show. + let bytes: Vec = (0..10_000u32) + .map(|value| (value / 4096) as u8 ^ value as u8) + .collect(); + let payload = ExactStatePayload::full_state(bytes); + + tier.spill("ns", &[1, 2, 3], &payload, None, None).unwrap(); + // Identical bytes again: a write, but no new segment bytes. + tier.spill("ns", &[1, 2, 3], &payload, None, None).unwrap(); + assert!( + tier.locate_longest("ns", &[1, 2, 3, 4], 8) + .unwrap() + .is_some() + ); + assert!(tier.locate_longest("other", &[9], 8).unwrap().is_none()); + let location = tier.locate_longest("ns", &[1, 2, 3], 8).unwrap().unwrap(); + let fill = tier.load(&location).unwrap(); + + let status = tier.status().unwrap(); + assert_eq!(status.format_version, MANIFEST_VERSION); + assert_eq!(status.restorable_manifests, 1); + assert_eq!(status.restorable_tokens, 3); + assert_eq!(status.namespaces, 1); + assert_eq!(status.activity.writes, 2); + assert_eq!( + status.activity.bytes_written, 10_000, + "duplicate spill re-wrote bytes" + ); + assert_eq!(status.activity.hits, 2); + assert_eq!(status.activity.misses, 1); + assert_eq!(status.activity.fills, 1); + assert_eq!(status.activity.bytes_read, fill.payload_bytes); + assert_eq!(status.activity.evictions, 0); + assert_eq!(status.activity.corrupt_entries, 0); + assert_eq!(status.activity.last_error, None); + } + + #[test] + fn spill_and_fill_roundtrip_all_payload_kinds() { + let tier = tier("roundtrip", "blake3:identity-a"); + let cases = vec![ + ( + "namespace-full", + ExactStatePayload::full_state((0..50_000u32).map(|v| v as u8).collect()), + ), + ( + "namespace-recurrent", + ExactStatePayload::recurrent_only(vec![9u8; 10_000]), + ), + ( + "namespace-kv", + ExactStatePayload::kv_recurrent(vec![1u8; 20_000], vec![2u8; 5_000]), + ), + ]; + for (namespace, payload) in cases { + let digest = tier + .spill( + namespace, + &tokens(512), + &payload, + Some("{\"desc\":1}".to_string()), + None, + ) + .expect("spill"); + let fill = tier + .fill_longest(namespace, &tokens(512), 64) + .expect("fill") + .expect("tier must hold the prefix"); + assert_eq!(fill.token_count, 512); + assert_eq!(fill.kv_desc_json.as_deref(), Some("{\"desc\":1}")); + assert_eq!(fill.payload.kind(), payload.kind()); + let manifest = tier + .store() + .load_manifest(&digest) + .expect("load roundtrip manifest"); + if payload.kind() != ExactStatePayloadKind::KvRecurrent { + assert!( + manifest + .segments + .iter() + .all(|segment| segment.codec_identity + == Some(SegmentCodecIdentity::raw(segment.bytes))) + ); + } + match payload.kind() { + ExactStatePayloadKind::KvRecurrent => { + assert_eq!( + fill.payload.kv_bytes().unwrap().unwrap().into_owned(), + payload.kv_bytes().unwrap().unwrap().into_owned() + ); + assert_eq!( + fill.payload.recurrent_state_bytes().unwrap().into_owned(), + payload.recurrent_state_bytes().unwrap().into_owned() + ); + } + ExactStatePayloadKind::RecurrentOnly => assert_eq!( + fill.payload.recurrent_state_bytes().unwrap().into_owned(), + payload.recurrent_state_bytes().unwrap().into_owned() + ), + ExactStatePayloadKind::FullState => assert_eq!( + fill.payload + .full_state_bytes_timed() + .unwrap() + .0 + .into_owned(), + payload.full_state_bytes_timed().unwrap().0.into_owned() + ), + } + } + } + + /// The sacrament case: a later, longer prompt (multi-turn growth) must + /// find the longest recorded shorter prefix, not just an exact match. + #[test] + fn longer_query_fills_from_longest_recorded_prefix() { + let tier = tier("longest", "blake3:identity-a"); + tier.spill( + "ns", + &tokens(800), + &ExactStatePayload::full_state(vec![1u8; 2048]), + None, + None, + ) + .expect("spill 800"); + tier.spill( + "ns", + &tokens(1200), + &ExactStatePayload::full_state(vec![2u8; 2048]), + None, + None, + ) + .expect("spill 1200"); + + // Query extends the 1200-token path: the longest entry wins. + let fill = tier + .fill_longest("ns", &tokens(1900), 64) + .expect("fill") + .expect("hit"); + assert_eq!(fill.token_count, 1200); + assert_eq!( + fill.payload + .full_state_bytes_timed() + .unwrap() + .0 + .into_owned(), + vec![2u8; 2048] + ); + + // Query between the two recorded lengths: the shorter entry wins. + let fill = tier + .fill_longest("ns", &tokens(1000), 64) + .expect("fill") + .expect("hit"); + assert_eq!(fill.token_count, 800); + } + + /// A recorded length only matches when the query genuinely starts with + /// the recorded tokens — a divergent prompt of the same length must + /// miss, not corrupt. + #[test] + fn divergent_tokens_at_a_recorded_length_miss() { + let tier = tier("divergent", "blake3:identity-a"); + tier.spill( + "ns", + &tokens(600), + &ExactStatePayload::full_state(vec![3u8; 1024]), + None, + None, + ) + .expect("spill"); + let mut divergent = tokens(600); + divergent[100] = 999_999; + assert!( + tier.fill_longest("ns", &divergent, 64) + .expect("fill") + .is_none() + ); + } + + #[test] + fn unknown_namespace_fills_none() { + let tier = tier("miss", "blake3:identity-a"); + assert!( + tier.fill_longest("never-spilled", &tokens(64), 64) + .expect("fill") + .is_none() + ); + } + + #[test] + fn identity_mismatch_is_refused_not_served() { + let root = temp_root("identity"); + let writer = L3Tier::open(&root, 0, "blake3:identity-a".to_string(), 4096).unwrap(); + writer + .spill( + "ns", + &tokens(128), + &ExactStatePayload::full_state(vec![5u8; 1024]), + None, + None, + ) + .expect("spill"); + let reader = L3Tier::open(&root, 0, "blake3:identity-b".to_string(), 4096).unwrap(); + assert!(reader.fill_longest("ns", &tokens(128), 64).is_err()); + } + + #[test] + fn model_identity_mismatch_is_refused_even_when_state_identity_matches() { + let root = temp_root("model-identity-mismatch"); + let manager = L3CacheManager::acquire(&root, StoreLimits::new(1_000_000, 0)).unwrap(); + let writer = manager.tier_for_model( + "blake3:model-a".to_string(), + "blake3:state".to_string(), + 4096, + ); + writer + .spill( + "ns", + &tokens(128), + &ExactStatePayload::full_state(vec![1u8; 2048]), + None, + None, + ) + .expect("spill"); + let reader = manager.tier_for_model( + "blake3:model-b".to_string(), + "blake3:state".to_string(), + 4096, + ); + + assert!(reader.fill_longest("ns", &tokens(128), 64).is_err()); + } + + #[test] + fn respilling_a_length_supersedes_the_older_entry() { + let tier = tier("supersede", "blake3:identity-a"); + tier.spill( + "ns", + &tokens(100), + &ExactStatePayload::full_state(vec![1u8; 2048]), + None, + None, + ) + .expect("first spill"); + tier.spill( + "ns", + &tokens(100), + &ExactStatePayload::full_state(vec![2u8; 2048]), + None, + None, + ) + .expect("second spill"); + let fill = tier + .fill_longest("ns", &tokens(100), 64) + .expect("fill") + .expect("present"); + assert_eq!(fill.token_count, 100); + assert_eq!( + fill.payload + .full_state_bytes_timed() + .unwrap() + .0 + .into_owned(), + vec![2u8; 2048] + ); + } + + /// A dense family whose KV export was unavailable must not persist an + /// empty entry that would later restore as position-without-state. + #[test] + fn empty_payloads_are_refused_at_spill() { + let tier = tier("empty", "blake3:identity-a"); + let empty = ExactStatePayload::kv_recurrent(Vec::new(), Vec::new()); + assert!(tier.spill("ns", &tokens(128), &empty, None, None).is_err()); + assert!( + tier.fill_longest("ns", &tokens(128), 64) + .expect("fill") + .is_none() + ); + } + + #[test] + fn kv_recurrent_payload_requires_kv_bytes() { + let tier = tier("empty-kv", "blake3:identity-a"); + let payload = ExactStatePayload::kv_recurrent(Vec::new(), vec![1u8; 128]); + assert!( + tier.spill("ns", &tokens(128), &payload, None, None) + .is_err() + ); + } + + /// The locate/load split: locate is a cheap index probe addressing one + /// physical entry (the single-flight claim key), and load returns the + /// same payload fill_longest would. + #[test] + fn locate_addresses_one_entry_and_load_fetches_it() { + let tier = tier("locate", "blake3:identity-a"); + tier.spill( + "ns", + &tokens(400), + &ExactStatePayload::full_state(vec![7u8; 4096]), + None, + None, + ) + .expect("spill"); + let location = tier + .locate_longest("ns", &tokens(500), 64) + .expect("locate") + .expect("hit"); + assert_eq!(location.token_count, 400); + // Two queries of different lengths that resolve to the same recorded + // prefix share one claim key. + let other = tier + .locate_longest("ns", &tokens(450), 64) + .expect("locate") + .expect("hit"); + assert_eq!(location.manifest_key, other.manifest_key); + let fill = tier.load(&location).expect("load"); + assert_eq!(fill.token_count, 400); + assert_eq!( + fill.payload + .full_state_bytes_timed() + .unwrap() + .0 + .into_owned(), + vec![7u8; 4096] + ); + } + + #[test] + fn restorable_summary_counts_matching_identity_only() { + let root = temp_root("summary"); + let tier_a = L3Tier::open(&root, 0, "blake3:identity-a".to_string(), 4096).unwrap(); + let tier_b = L3Tier::open(&root, 0, "blake3:identity-b".to_string(), 4096).unwrap(); + tier_a + .spill( + "ns", + &tokens(300), + &ExactStatePayload::full_state(vec![1u8; 512]), + None, + None, + ) + .unwrap(); + tier_b + .spill( + "ns", + &tokens(700), + &ExactStatePayload::full_state(vec![2u8; 512]), + None, + None, + ) + .unwrap(); + let (count, restorable_tokens, footprint) = tier_a.restorable_summary().unwrap(); + assert_eq!(count, 1); + assert_eq!(restorable_tokens, 300); + assert!(footprint >= 1024); + } + + /// Rewrite the on-disk codec identity of a committed manifest, simulating a + /// future/remote entry this build cannot decode. The manifest layout is + /// `/manifests/.json`. + fn corrupt_manifest_codec(store: &HandoffSegmentStore, digest: &str, name: &str, version: u32) { + let path = store + .root() + .join("manifests") + .join(format!("{digest}.json")); + let mut value: serde_json::Value = + serde_json::from_slice(&std::fs::read(&path).expect("read manifest")) + .expect("parse manifest"); + value["codec"] = serde_json::json!({ "name": name, "version": version }); + std::fs::write(&path, serde_json::to_vec(&value).expect("serialize")) + .expect("write manifest"); + } + + #[test] + fn locate_longest_skips_unsupported_codec_and_falls_back_to_shorter() { + let tier = tier("codec-locate-fallback", "blake3:codec"); + let short = ExactStatePayload::full_state(vec![1u8; 4096]); + let long = ExactStatePayload::full_state(vec![2u8; 8192]); + tier.spill("ns", &tokens(3), &short, None, None) + .expect("spill short"); + tier.spill("ns", &tokens(6), &long, None, None) + .expect("spill long"); + + // The longer prefix is the natural longest match before tampering. + let long_loc = tier + .locate_longest("ns", &tokens(6), 8) + .expect("locate") + .expect("long entry present"); + assert_eq!(long_loc.token_count, 6); + + // Make only the longer entry's codec unsupported on disk. + corrupt_manifest_codec(tier.store(), &long_loc.manifest_key, "zstd", 1); + + // locate_longest must skip the unsupported longer prefix (pruning its + // link) and fall back to the shorter supported one, not error or stop. + let fell_back = tier + .locate_longest("ns", &tokens(6), 8) + .expect("locate after tamper") + .expect("shorter entry still serves"); + assert_eq!( + fell_back.token_count, 3, + "fell back to the supported shorter prefix" + ); + + // The bad longer link was pruned, so the shorter prefix is the answer. + let again = tier + .locate_longest("ns", &tokens(6), 8) + .expect("locate again") + .expect("shorter entry still serves"); + assert_eq!(again.token_count, 3); + } +} diff --git a/crates/skippy-correctness/Cargo.toml b/crates/skippy-correctness/Cargo.toml index 4cad1a79ed..22f901f5bd 100644 --- a/crates/skippy-correctness/Cargo.toml +++ b/crates/skippy-correctness/Cargo.toml @@ -7,7 +7,9 @@ version.workspace = true [dependencies] hex = "0.4" anyhow.workspace = true +blake3.workspace = true clap.workspace = true +skippy-cache = { path = "../skippy-cache" } skippy-protocol = { path = "../skippy-protocol" } skippy-runtime = { path = "../skippy-runtime" } model-artifact = { path = "../model-artifact" } diff --git a/crates/skippy-correctness/src/cli.rs b/crates/skippy-correctness/src/cli.rs index 96b56508b4..a55b5cc15b 100644 --- a/crates/skippy-correctness/src/cli.rs +++ b/crates/skippy-correctness/src/cli.rs @@ -16,10 +16,33 @@ pub enum CommandKind { Chain(ChainArgs), SplitScan(SplitScanArgs), StateHandoff(StateHandoffArgs), + RemoteHandoff(RemoteHandoffArgs), SplitPrefixHit(SplitPrefixHitArgs), NativeMtpOpenAiAb(Box), GlmDsaStage0Trace(Box), StageFaParity(StageFaParityArgs), + KvPageGrowth(KvPageGrowthArgs), +} + +#[derive(Args)] +pub struct KvPageGrowthArgs { + #[command(flatten)] + pub runtime: RuntimeArgs, + /// Tokens prefilled before the first export, standing in for an agent's + /// system prefix. + #[arg(long, default_value_t = 2048)] + pub base_tokens: usize, + /// Tokens appended per turn. + #[arg(long, default_value_t = 512)] + pub turn_tokens: usize, + /// Turns appended after the base prefix. + #[arg(long, default_value_t = 4)] + pub turns: usize, + /// Segment size the L3 store cuts at. + #[arg(long, default_value_t = 8 * 1024 * 1024)] + pub segment_bytes: u64, + #[arg(long)] + pub json: Option, } #[derive(Args, Clone)] @@ -67,6 +90,29 @@ pub enum FlashAttentionArg { Enabled, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] +pub enum CacheTypeArg { + #[value(name = "f16")] + F16, + #[value(name = "f32")] + F32, + #[value(name = "q8_0")] + Q8Zero, + #[value(name = "q4_0")] + Q4Zero, +} + +impl CacheTypeArg { + pub(crate) const fn ggml_type(self) -> u32 { + match self { + Self::F16 => skippy_runtime::GGML_TYPE_F16, + Self::F32 => skippy_runtime::GGML_TYPE_F32, + Self::Q8Zero => skippy_runtime::GGML_TYPE_Q8_0, + Self::Q4Zero => skippy_runtime::GGML_TYPE_Q4_0, + } + } +} + #[derive(Args, Clone)] pub struct ServerArgs { #[arg(long, default_value = "target/debug/skippy-server")] @@ -188,6 +234,123 @@ pub struct StateHandoffArgs { pub synthetic_input_activation: bool, #[arg(long)] pub binary_control: bool, + /// Run the experimental CacheGen acceptance gate against the native + /// KV-page control. Requires a local full-model kv-recurrent handoff and + /// direct device decode support; unsupported backends fail without a + /// scalar restore fallback. + #[arg(long)] + pub cachegen_gate: bool, + /// Native K cache type used by the state-handoff and CacheGen control arms. + #[arg(long, value_enum, default_value = "f16")] + pub cache_type_k: CacheTypeArg, + /// Native V cache type used by the state-handoff and CacheGen control arms. + #[arg(long, value_enum, default_value = "f16")] + pub cache_type_v: CacheTypeArg, + /// Teacher-forced continuation steps used for CacheGen quality and + /// steady-state decode measurements. + #[arg(long, default_value_t = 64)] + pub cachegen_continuation_steps: usize, + /// Minimum fraction of greedy tokens that must agree with native. + #[arg(long, default_value_t = 0.95)] + pub cachegen_min_token_agreement: f64, + /// Maximum allowed CacheGen/native p99 decode latency regression. + #[arg(long, default_value_t = 0.05)] + pub cachegen_max_p99_decode_regression: f64, + /// Optional maximum estimated codec working bytes. If omitted, peak + /// memory is reported without adding a pass/fail criterion. + #[arg(long)] + pub cachegen_max_peak_working_bytes: Option, + #[arg(long)] + pub allow_mismatch: bool, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] +#[value(rename_all = "kebab-case")] +pub enum RemoteHandoffRole { + Send, + Recv, + Restore, + Serve, + Fetch, +} + +#[derive(Args)] +pub struct RemoteHandoffArgs { + #[command(flatten)] + pub runtime: RuntimeArgs, + #[command(flatten)] + pub output: OutputArgs, + #[arg(long, value_enum, help = "send = prefill node, recv = decode node")] + pub role: RemoteHandoffRole, + #[arg( + long, + default_value = "127.0.0.1:19081", + help = "Address the receiver listens on; binding beyond loopback exposes the unauthenticated lab transport and requires a trusted private network" + )] + pub listen: SocketAddr, + #[arg(long, help = "Receiver address the sender connects to")] + pub peer: Option, + #[arg(long, value_enum, default_value = "full-state")] + pub state_payload_kind: StatePayloadKind, + #[arg( + long, + help = "Expand or truncate the prompt to this many prefix tokens" + )] + pub prefix_token_count: Option, + #[arg( + long, + default_value_t = 32, + help = "Greedy continuation length compared token-for-token across nodes" + )] + pub decode_tokens: usize, + #[arg(long, default_value_t = 8 * 1024 * 1024)] + pub segment_bytes: usize, + #[arg( + long, + help = "Also measure prefill-in-place on the receiver for a TTFT baseline" + )] + pub baseline: bool, + #[arg(long)] + pub runtime_lane_count: Option, + #[arg( + long, + default_value_t = 600, + help = "Per-read socket timeout for the whole connection (not just the handshake): a peer that stalls mid-stream errors out after this long" + )] + pub handshake_timeout_secs: u64, + #[arg( + long, + default_value_t = 1, + help = "Receiver only: handoffs to serve before exiting (0 = until killed); reports get a -N suffix when not 1" + )] + pub accept_count: usize, + #[arg( + long, + help = "L3 segment store directory: sender spills exported state, receiver write-behinds incoming segments and imports from the store, restore reattaches from it" + )] + pub store_dir: Option, + #[arg( + long, + default_value_t = 0, + help = "Segment footprint cap for the store; oldest manifests evict first (0 = unlimited)" + )] + pub store_budget_bytes: u64, + #[arg( + long, + help = "Restore only: manifest key (payload digest) to reattach; defaults to the newest manifest" + )] + pub manifest: Option, + #[arg( + long, + help = "Stream KV pages per prefill chunk, overlapping transfer with the remaining prefill; the receiver stages pages but cannot generate until the commit record validates (pass on both sides)" + )] + pub streaming: bool, + #[arg( + long, + default_value_t = 512, + help = "Prefill chunk size in tokens for --streaming" + )] + pub stream_chunk_tokens: usize, #[arg(long)] pub allow_mismatch: bool, } @@ -331,6 +494,30 @@ pub enum StatePayloadKind { KvRecurrent, } +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn remote_handoff_listener_defaults_to_loopback() { + let cli = Cli::try_parse_from([ + "skippy-correctness", + "remote-handoff", + "--model", + "model.gguf", + "--role", + "serve", + ]) + .expect("parse remote handoff defaults"); + let CommandKind::RemoteHandoff(args) = cli.command else { + panic!("expected remote handoff command"); + }; + + assert!(args.listen.ip().is_loopback()); + assert_eq!(args.listen.port(), 19081); + } +} + #[derive(Args)] pub struct StageFaParityArgs { #[arg(long)] diff --git a/crates/skippy-correctness/src/main.rs b/crates/skippy-correctness/src/main.rs index 2127d4259a..b04aa099f5 100644 --- a/crates/skippy-correctness/src/main.rs +++ b/crates/skippy-correctness/src/main.rs @@ -12,7 +12,10 @@ use crate::{ cli::{Cli, CommandKind}, glm_dsa_trace::glm_dsa_stage0_trace, native_mtp_openai::native_mtp_openai_ab, - runner::{chain, single_step, split_prefix_hit, split_scan, stage_fa_parity, state_handoff}, + runner::{ + chain, kv_page_growth, remote_handoff, single_step, split_prefix_hit, split_scan, + stage_fa_parity, state_handoff, + }, }; fn prepare_model_download_directories() { @@ -40,9 +43,11 @@ fn main() -> Result<()> { CommandKind::Chain(args) => chain(args), CommandKind::SplitScan(args) => split_scan(args), CommandKind::StateHandoff(args) => state_handoff(args), + CommandKind::RemoteHandoff(args) => remote_handoff(args), CommandKind::SplitPrefixHit(args) => split_prefix_hit(args), CommandKind::NativeMtpOpenAiAb(args) => native_mtp_openai_ab(*args), CommandKind::GlmDsaStage0Trace(args) => glm_dsa_stage0_trace(*args), CommandKind::StageFaParity(args) => stage_fa_parity(args), + CommandKind::KvPageGrowth(args) => kv_page_growth(args), } } diff --git a/crates/skippy-correctness/src/report.rs b/crates/skippy-correctness/src/report.rs index b2ba771ec3..33f2f861b4 100644 --- a/crates/skippy-correctness/src/report.rs +++ b/crates/skippy-correctness/src/report.rs @@ -1,7 +1,64 @@ -use serde::Serialize; +use serde::{Deserialize, Serialize}; pub use model_artifact::ModelIdentity; +#[derive(Debug, Serialize)] +pub struct RemoteHandoffReport { + pub mode: &'static str, + pub status: &'static str, + pub role: &'static str, + pub model_identity: ModelIdentity, + pub matches: bool, + pub tokens_match: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub baseline_matches: Option, + pub state_payload_kind: &'static str, + pub prompt_token_count: usize, + pub decode_token_count: usize, + pub continuation_token: i32, + pub source_tokens: Vec, + pub restored_tokens: Vec, + pub state_bytes: usize, + pub state_bytes_per_prompt_token: f64, + pub kv_bytes: usize, + pub recurrent_bytes: usize, + pub segment_count: usize, + pub segment_bytes: usize, + pub payload_digest: String, + pub model_load_ms: f64, + pub tokenize_ms: f64, + pub source_prefill_ms: f64, + pub state_export_ms: f64, + pub transfer_ms: f64, + pub transfer_gbps: f64, + pub source_decode_ms: f64, + #[serde(skip_serializing_if = "Option::is_none")] + pub store_ms: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub overlap_wall_ms: Option, + pub receiver: RemoteHandoffReceiverTimings, + pub ttft_disaggregated_ms: f64, + #[serde(skip_serializing_if = "Option::is_none")] + pub ttft_local_ms: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub ttft_speedup: Option, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct RemoteHandoffReceiverTimings { + pub model_load_ms: f64, + pub transfer_receive_ms: f64, + pub kv_attach_ms: f64, + #[serde(default)] + pub store_ms: f64, + #[serde(default)] + pub attach_residual_ms: f64, + pub first_decode_ms: f64, + pub decode_ms: f64, + pub baseline_prefill_ms: f64, + pub baseline_first_decode_ms: f64, +} + #[derive(Debug, Serialize)] pub struct BaselineReport { pub token_id: i32, @@ -203,9 +260,55 @@ pub struct StateHandoffReport { pub cache_hit_import_ms: Vec, #[serde(skip_serializing_if = "Vec::is_empty")] pub cache_hit_decode_ms: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub cachegen_gate: Option, pub stage_models: Vec, } +#[derive(Debug, Serialize)] +pub struct CacheGenGateReport { + pub passed: bool, + pub failure_reasons: Vec, + pub restore_path: &'static str, + pub cache_type_k: &'static str, + pub cache_type_v: &'static str, + pub continuation_steps: usize, + pub native_storage_bytes: usize, + pub cachegen_storage_bytes: usize, + pub compression_ratio: f64, + pub tile_count: usize, + pub encode_ms: f64, + pub scalar_oracle_decode_ms: f64, + pub native_write_ms: f64, + pub cachegen_write_ms: f64, + pub native_persist_ms: f64, + pub cachegen_persist_ms: f64, + pub native_read_ms: f64, + pub cachegen_read_ms: f64, + pub native_import_ms: f64, + pub cachegen_import_ms: f64, + pub native_ttft_ms: f64, + pub cachegen_ttft_ms: f64, + pub native_decode_tokens_per_second: f64, + pub cachegen_decode_tokens_per_second: f64, + pub native_p99_decode_ms: f64, + pub cachegen_p99_decode_ms: f64, + pub p99_decode_regression: f64, + pub matching_tokens: usize, + pub token_agreement: f64, + #[serde(skip_serializing_if = "Option::is_none")] + pub first_token_mismatch_step: Option, + pub mean_entropy_abs_drift: f64, + pub max_entropy_abs_drift: f64, + pub mean_top_logprob_abs_drift: f64, + pub max_top_logprob_abs_drift: f64, + pub estimated_peak_codec_working_bytes: usize, + pub min_token_agreement: f64, + pub max_p99_decode_regression: f64, + #[serde(skip_serializing_if = "Option::is_none")] + pub max_peak_codec_working_bytes: Option, +} + #[derive(Debug, Serialize, Clone)] pub struct StatePayloadDigestReport { pub payload_kind: &'static str, diff --git a/crates/skippy-correctness/src/runner/cachegen_gate.rs b/crates/skippy-correctness/src/runner/cachegen_gate.rs new file mode 100644 index 0000000000..75e1b63fb6 --- /dev/null +++ b/crates/skippy-correctness/src/runner/cachegen_gate.rs @@ -0,0 +1,606 @@ +use std::{ + fs::{self, File}, + hint::black_box, + io::Write, + path::{Path, PathBuf}, + time::{Instant, SystemTime, UNIX_EPOCH}, +}; + +use anyhow::{Context, Result, bail}; +use skippy_cache::cachegen::archive::{ + CacheGenArchive, ComponentLayout, PageLayout, ValueType, decode_page, encode_page, +}; +use skippy_runtime::{ + GGML_TYPE_F16, GGML_TYPE_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, KV_PAGE_FLAG_V_TRANSPOSED, + RuntimeKvPageDesc, StageModel, StageSession, TokenSignal, +}; + +use crate::report::CacheGenGateReport; + +use super::{ + stage_execution::{BinaryStateHandoffConfig, elapsed_ms}, + state_handoff::LocalStatePayload, +}; + +struct PersistedPayloads { + native_kv: Vec, + native_recurrent: Vec, + cachegen_archive: Vec, + cachegen_recurrent: Vec, + native_write_ms: f64, + cachegen_write_ms: f64, + native_read_ms: f64, + cachegen_read_ms: f64, +} + +pub(in crate::runner) fn run_cachegen_gate( + model: &StageModel, + args: &BinaryStateHandoffConfig, + payload: &LocalStatePayload, + prefix: &[i32], + continuation: i32, +) -> Result { + let LocalStatePayload::KvRecurrent { + kv_desc: Some(kv_desc), + kv, + recurrent, + } = payload + else { + bail!("CacheGen gate requires an exported KV page descriptor and bytes"); + }; + kv_desc.validate_payload(kv.len())?; + + let encode_started = Instant::now(); + let archive = encode_kv_archive(kv_desc, kv)?; + let encode_ms = elapsed_ms(encode_started); + let native_storage_bytes = kv.len().saturating_add(recurrent.len()); + let cachegen_storage_bytes = archive.bytes.len().saturating_add(recurrent.len()); + + let persisted = persist_and_read_payloads(kv, recurrent, &archive.bytes)?; + // Keep the scalar decoder in the gate as an independently timed oracle, + // but never feed its multi-gigabyte output into the measured restore. The + // accelerated path must consume the persisted archive directly or fail. + let scalar_oracle_decode_started = Instant::now(); + let scalar_oracle = decode_kv_archive(kv_desc, &persisted.cachegen_archive)?; + let scalar_oracle_decode_ms = elapsed_ms(scalar_oracle_decode_started); + black_box(&scalar_oracle); + drop(scalar_oracle); + + let native_payload = LocalStatePayload::KvRecurrent { + kv_desc: Some(kv_desc.clone()), + kv: persisted.native_kv, + recurrent: persisted.native_recurrent, + }; + let native_import_started = Instant::now(); + let mut native = import_session(model, &native_payload, prefix)?; + let native_import_ms = elapsed_ms(native_import_started); + let native_continuation = + run_native_continuation(&mut native, continuation, args.cachegen_continuation_steps)?; + drop(native); + + let cachegen_import_started = Instant::now(); + let mut cachegen = import_cachegen_session( + model, + kv_desc, + &persisted.cachegen_archive, + &persisted.cachegen_recurrent, + prefix, + )?; + let cachegen_import_ms = elapsed_ms(cachegen_import_started); + let continuation = + compare_cachegen_continuation(&mut cachegen, continuation, native_continuation)?; + let native_p99_decode_ms = percentile_99(&continuation.native_decode_ms); + let cachegen_p99_decode_ms = percentile_99(&continuation.cachegen_decode_ms); + let p99_decode_regression = relative_regression(cachegen_p99_decode_ms, native_p99_decode_ms); + let native_ttft_ms = restore_to_first_token_ms( + persisted.native_read_ms, + native_import_ms, + &continuation.native_decode_ms, + ); + let cachegen_ttft_ms = restore_to_first_token_ms( + persisted.cachegen_read_ms, + cachegen_import_ms, + &continuation.cachegen_decode_ms, + ); + let token_agreement = + continuation.matching_tokens as f64 / args.cachegen_continuation_steps as f64; + let compression_ratio = cachegen_storage_bytes as f64 / native_storage_bytes.max(1) as f64; + + let mut failure_reasons = Vec::new(); + if cachegen_storage_bytes >= native_storage_bytes { + failure_reasons.push("encoded payload is not smaller than native".to_string()); + } + if cachegen_ttft_ms >= native_ttft_ms { + failure_reasons.push(format!( + "restore-to-first-token did not beat native ({cachegen_ttft_ms:.3} ms >= {native_ttft_ms:.3} ms)" + )); + } + if token_agreement < args.cachegen_min_token_agreement { + failure_reasons.push(format!( + "token agreement {token_agreement:.4} is below {:.4}", + args.cachegen_min_token_agreement + )); + } + if p99_decode_regression > args.cachegen_max_p99_decode_regression { + failure_reasons.push(format!( + "p99 decode regression {p99_decode_regression:.4} exceeds {:.4}", + args.cachegen_max_p99_decode_regression + )); + } + if let Some(limit) = args.cachegen_max_peak_working_bytes + && archive.estimated_peak_codec_working_bytes > limit + { + failure_reasons.push(format!( + "estimated codec working set {} bytes exceeds {limit} bytes", + archive.estimated_peak_codec_working_bytes + )); + } + + Ok(CacheGenGateReport { + passed: failure_reasons.is_empty(), + failure_reasons, + restore_path: "native-device", + cache_type_k: cache_type_name(kv_desc.k_type)?, + cache_type_v: cache_type_name(kv_desc.v_type)?, + continuation_steps: args.cachegen_continuation_steps, + native_storage_bytes, + cachegen_storage_bytes, + compression_ratio, + tile_count: archive.tile_count, + encode_ms, + scalar_oracle_decode_ms, + native_write_ms: persisted.native_write_ms, + cachegen_write_ms: persisted.cachegen_write_ms, + native_persist_ms: persisted.native_write_ms, + cachegen_persist_ms: encode_ms + persisted.cachegen_write_ms, + native_read_ms: persisted.native_read_ms, + cachegen_read_ms: persisted.cachegen_read_ms, + native_import_ms, + cachegen_import_ms, + native_ttft_ms, + cachegen_ttft_ms, + native_decode_tokens_per_second: tokens_per_second(&continuation.native_decode_ms), + cachegen_decode_tokens_per_second: tokens_per_second(&continuation.cachegen_decode_ms), + native_p99_decode_ms, + cachegen_p99_decode_ms, + p99_decode_regression, + matching_tokens: continuation.matching_tokens, + token_agreement, + first_token_mismatch_step: continuation.first_token_mismatch_step, + mean_entropy_abs_drift: mean(&continuation.entropy_abs_drift), + max_entropy_abs_drift: max_or_zero(&continuation.entropy_abs_drift), + mean_top_logprob_abs_drift: mean(&continuation.top_logprob_abs_drift), + max_top_logprob_abs_drift: max_or_zero(&continuation.top_logprob_abs_drift), + estimated_peak_codec_working_bytes: archive.estimated_peak_codec_working_bytes, + min_token_agreement: args.cachegen_min_token_agreement, + max_p99_decode_regression: args.cachegen_max_p99_decode_regression, + max_peak_codec_working_bytes: args.cachegen_max_peak_working_bytes, + }) +} + +fn import_cachegen_session( + model: &StageModel, + kv_desc: &RuntimeKvPageDesc, + archive: &[u8], + recurrent: &[u8], + prefix: &[i32], +) -> Result { + let mut session = model + .create_session() + .context("create CacheGen gate session")?; + session + .import_cachegen_kv_page(kv_desc, archive) + .context("import CacheGen archive directly into resident KV")?; + session + .import_recurrent_state_for_token_count(recurrent, prefix.len() as u64) + .context("import CacheGen gate recurrent state")?; + Ok(session) +} + +fn import_session( + model: &StageModel, + payload: &LocalStatePayload, + prefix: &[i32], +) -> Result { + let LocalStatePayload::KvRecurrent { + kv_desc: Some(kv_desc), + kv, + recurrent, + } = payload + else { + bail!("CacheGen gate internal payload is not KV-recurrent"); + }; + let mut session = model.create_session().context("create gate session")?; + session + .import_kv_page(kv_desc, kv) + .context("import gate KV page")?; + session + .import_recurrent_state_for_token_count(recurrent, prefix.len() as u64) + .context("import gate recurrent state")?; + Ok(session) +} + +struct ContinuationComparison { + native_decode_ms: Vec, + cachegen_decode_ms: Vec, + matching_tokens: usize, + first_token_mismatch_step: Option, + entropy_abs_drift: Vec, + top_logprob_abs_drift: Vec, +} + +struct NativeContinuation { + predicted_tokens: Vec, + signals: Vec, + decode_ms: Vec, +} + +fn run_native_continuation( + native: &mut StageSession, + mut token: i32, + steps: usize, +) -> Result { + let mut predicted_tokens = Vec::with_capacity(steps); + let mut signals = Vec::with_capacity(steps); + let mut decode_ms = Vec::with_capacity(steps); + for _ in 0..steps { + let started = Instant::now(); + let prediction = native.decode_step(token).context("native gate decode")?; + decode_ms.push(elapsed_ms(started)); + signals.push( + native + .last_token_signal() + .context("native gate token signal")?, + ); + predicted_tokens.push(prediction); + token = prediction; + } + Ok(NativeContinuation { + predicted_tokens, + signals, + decode_ms, + }) +} + +fn compare_cachegen_continuation( + cachegen: &mut StageSession, + first_token: i32, + native: NativeContinuation, +) -> Result { + let steps = native.predicted_tokens.len(); + let mut cachegen_decode_ms = Vec::with_capacity(steps); + let mut matching_tokens = 0usize; + let mut first_token_mismatch_step = None; + let mut entropy_abs_drift = Vec::with_capacity(steps); + let mut top_logprob_abs_drift = Vec::with_capacity(steps); + for step in 0..steps { + let token = if step == 0 { + first_token + } else { + native.predicted_tokens[step - 1] + }; + let started = Instant::now(); + let cachegen_prediction = cachegen + .decode_step(token) + .context("CacheGen gate decode")?; + cachegen_decode_ms.push(elapsed_ms(started)); + let cachegen_signal = cachegen + .last_token_signal() + .context("CacheGen gate token signal")?; + let native_prediction = native.predicted_tokens[step]; + let native_signal = native.signals[step]; + if native_prediction == cachegen_prediction { + matching_tokens += 1; + } else if first_token_mismatch_step.is_none() { + first_token_mismatch_step = Some(step); + } + entropy_abs_drift.push(f64::from( + (native_signal.entropy - cachegen_signal.entropy).abs(), + )); + top_logprob_abs_drift.push(f64::from( + (native_signal.top_logprob - cachegen_signal.top_logprob).abs(), + )); + } + Ok(ContinuationComparison { + native_decode_ms: native.decode_ms, + cachegen_decode_ms, + matching_tokens, + first_token_mismatch_step, + entropy_abs_drift, + top_logprob_abs_drift, + }) +} + +fn encode_kv_archive(desc: &RuntimeKvPageDesc, raw: &[u8]) -> Result { + desc.validate_payload(raw.len())?; + encode_page(&page_layout(desc)?, raw) +} + +fn decode_kv_archive(desc: &RuntimeKvPageDesc, archive: &[u8]) -> Result> { + let raw_len = usize::try_from(desc.payload_bytes).context("descriptor length exceeds usize")?; + desc.validate_payload(raw_len)?; + decode_page(archive, raw_len) +} + +fn page_layout(desc: &RuntimeKvPageDesc) -> Result { + let components = if desc.component_count == 0 { + vec![component_layout( + desc.token_count, + desc.layer_count, + desc.k_type, + desc.v_type, + desc.k_row_bytes, + desc.v_row_bytes, + desc.v_element_bytes, + desc.k_idx_row_bytes, + 0, + desc.payload_bytes, + desc.flags, + )?] + } else { + desc.components + .iter() + .take(desc.component_count as usize) + .map(|component| { + component_layout( + component.token_count, + component.layer_count, + component.k_type, + component.v_type, + component.k_row_bytes, + component.v_row_bytes, + component.v_element_bytes, + component.k_idx_row_bytes, + component.payload_offset, + component.payload_bytes, + component.flags, + ) + }) + .collect::>>()? + }; + Ok(PageLayout { + payload_bytes: desc.payload_bytes, + components, + }) +} + +#[allow(clippy::too_many_arguments)] +fn component_layout( + token_count: u64, + layer_count: u32, + k_type: u32, + v_type: u32, + k_row_bytes: u32, + v_row_bytes: u32, + v_element_bytes: u32, + k_idx_row_bytes: u32, + payload_offset: u64, + payload_bytes: u64, + flags: u64, +) -> Result { + Ok(ComponentLayout { + token_count, + layer_count, + k_type: cachegen_value_type(k_type)?, + v_type: cachegen_value_type(v_type)?, + k_row_bytes, + v_row_bytes, + v_element_bytes, + k_idx_row_bytes, + payload_offset, + payload_bytes, + v_transposed: flags & KV_PAGE_FLAG_V_TRANSPOSED != 0, + }) +} + +fn cachegen_value_type(value: u32) -> Result { + match value { + GGML_TYPE_F32 => Ok(ValueType::F32), + GGML_TYPE_F16 => Ok(ValueType::F16), + GGML_TYPE_Q8_0 => Ok(ValueType::Q8_0), + GGML_TYPE_Q4_0 => Ok(ValueType::Q4_0), + _ => bail!("CacheGen gate does not support runtime K/V type {value}"), + } +} + +fn cache_type_name(value: u32) -> Result<&'static str> { + match value { + GGML_TYPE_F32 => Ok("f32"), + GGML_TYPE_F16 => Ok("f16"), + GGML_TYPE_Q8_0 => Ok("q8_0"), + GGML_TYPE_Q4_0 => Ok("q4_0"), + _ => bail!("CacheGen gate does not support runtime K/V type {value}"), + } +} + +fn persist_and_read_payloads( + native_kv: &[u8], + recurrent: &[u8], + cachegen_archive: &[u8], +) -> Result { + let nonce = SystemTime::now() + .duration_since(UNIX_EPOCH) + .context("system clock precedes Unix epoch")? + .as_nanos(); + let root = std::env::temp_dir().join(format!( + "skippy-cachegen-gate-{}-{nonce}", + std::process::id() + )); + fs::create_dir(&root).with_context(|| format!("create {}", root.display()))?; + let native_path = root.join("native.bin"); + let cachegen_path = root.join("cachegen.bin"); + let result = (|| { + let native_write_ms = write_payload(&native_path, native_kv, recurrent)?; + let cachegen_write_ms = write_payload(&cachegen_path, cachegen_archive, recurrent)?; + let (native_bytes, native_read_ms) = read_payload(&native_path)?; + let (cachegen_bytes, cachegen_read_ms) = read_payload(&cachegen_path)?; + let native_split = native_kv.len(); + let cachegen_split = cachegen_archive.len(); + if native_bytes.len() != native_split + recurrent.len() + || cachegen_bytes.len() != cachegen_split + recurrent.len() + { + bail!("persisted gate payload length mismatch"); + } + Ok(PersistedPayloads { + native_kv: native_bytes[..native_split].to_vec(), + native_recurrent: native_bytes[native_split..].to_vec(), + cachegen_archive: cachegen_bytes[..cachegen_split].to_vec(), + cachegen_recurrent: cachegen_bytes[cachegen_split..].to_vec(), + native_write_ms, + cachegen_write_ms, + native_read_ms, + cachegen_read_ms, + }) + })(); + let _ = fs::remove_dir_all(root); + result +} + +fn write_payload(path: &Path, first: &[u8], second: &[u8]) -> Result { + let started = Instant::now(); + let mut file = File::create(path).with_context(|| format!("create {}", path.display()))?; + file.write_all(first)?; + file.write_all(second)?; + file.sync_all()?; + Ok(elapsed_ms(started)) +} + +fn read_payload(path: &PathBuf) -> Result<(Vec, f64)> { + let started = Instant::now(); + let bytes = fs::read(path).with_context(|| format!("read {}", path.display()))?; + black_box(&bytes); + Ok((bytes, elapsed_ms(started))) +} + +fn percentile_99(samples: &[f64]) -> f64 { + let mut sorted = samples.to_vec(); + sorted.sort_by(f64::total_cmp); + let index = ((sorted.len() as f64 * 0.99).ceil() as usize) + .saturating_sub(1) + .min(sorted.len().saturating_sub(1)); + sorted.get(index).copied().unwrap_or(0.0) +} + +fn relative_regression(candidate: f64, baseline: f64) -> f64 { + if baseline <= f64::EPSILON { + if candidate <= baseline { + 0.0 + } else { + f64::INFINITY + } + } else { + (candidate - baseline) / baseline + } +} + +fn tokens_per_second(samples: &[f64]) -> f64 { + let total_ms: f64 = samples.iter().sum(); + if total_ms <= f64::EPSILON { + 0.0 + } else { + samples.len() as f64 * 1000.0 / total_ms + } +} + +fn restore_to_first_token_ms(read_ms: f64, import_ms: f64, decode_ms: &[f64]) -> f64 { + read_ms + import_ms + decode_ms.first().copied().unwrap_or(0.0) +} + +fn mean(values: &[f64]) -> f64 { + if values.is_empty() { + 0.0 + } else { + values.iter().sum::() / values.len() as f64 + } +} + +fn max_or_zero(values: &[f64]) -> f64 { + values.iter().copied().fold(0.0, f64::max) +} + +#[cfg(test)] +mod tests { + use super::*; + use skippy_cache::cachegen::lmcache::MAX_TOKENS_PER_CHUNK; + + fn f16_bytes(values: usize) -> Vec { + (0..values) + .flat_map(|index| { + skippy_protocol::binary::f32_to_f16_bits(index as f32 / 17.0).to_le_bytes() + }) + .collect() + } + + fn descriptor(flags: u64) -> RuntimeKvPageDesc { + descriptor_with_tokens(flags, 4) + } + + fn descriptor_with_tokens(flags: u64, token_count: u64) -> RuntimeKvPageDesc { + let payload_bytes = token_count * 2 * 2 * 6; + RuntimeKvPageDesc { + version: 1, + layer_start: 0, + layer_end: 2, + token_start: 0, + token_count, + layer_count: 2, + k_type: GGML_TYPE_F16, + v_type: GGML_TYPE_F16, + k_row_bytes: 6, + v_row_bytes: 6, + v_element_bytes: 2, + k_idx_row_bytes: 0, + payload_bytes, + flags, + codec: 1, + component_count: 0, + components: Default::default(), + } + } + + #[test] + fn archive_roundtrip_preserves_geometry_and_length() { + let desc = descriptor(0); + let raw = f16_bytes(48); + let archive = encode_kv_archive(&desc, &raw).expect("encode"); + let decoded = decode_kv_archive(&desc, &archive.bytes).expect("decode"); + assert_eq!(decoded.len(), raw.len()); + assert_eq!(archive.tile_count, 4); + assert_ne!(decoded, raw, "fixture must exercise lossy quantization"); + } + + #[test] + fn transposed_v_layout_is_restored_before_native_import() { + let desc = descriptor(KV_PAGE_FLAG_V_TRANSPOSED); + let raw = f16_bytes(48); + let archive = encode_kv_archive(&desc, &raw).expect("encode"); + let decoded = decode_kv_archive(&desc, &archive.bytes).expect("decode"); + assert_eq!(decoded.len(), raw.len()); + assert_eq!(archive.tile_count, 4); + } + + #[test] + fn archive_rejects_uncovered_descriptor_bytes() { + let mut desc = descriptor(0); + desc.payload_bytes += 2; + let raw = f16_bytes(49); + assert!(encode_kv_archive(&desc, &raw).is_err()); + } + + #[test] + fn archive_chunks_long_transposed_pages_at_the_reference_tile_size() { + let token_count = MAX_TOKENS_PER_CHUNK as u64 + 4; + let desc = descriptor_with_tokens(KV_PAGE_FLAG_V_TRANSPOSED, token_count); + let raw = f16_bytes(desc.payload_bytes as usize / 2); + let archive = encode_kv_archive(&desc, &raw).expect("encode"); + let decoded = decode_kv_archive(&desc, &archive.bytes).expect("decode"); + assert_eq!(decoded.len(), raw.len()); + assert_eq!(archive.tile_count, 8); + } + + #[test] + fn restore_timing_contains_only_persisted_read_native_import_and_first_decode() { + let ttft_ms = restore_to_first_token_ms(12.0, 34.0, &[5.0, 999.0]); + assert_eq!(ttft_ms, 51.0); + } +} diff --git a/crates/skippy-correctness/src/runner/kv_page_growth.rs b/crates/skippy-correctness/src/runner/kv_page_growth.rs new file mode 100644 index 0000000000..048289febf --- /dev/null +++ b/crates/skippy-correctness/src/runner/kv_page_growth.rs @@ -0,0 +1,433 @@ +//! Does a growing prefix append to the exported KV page, or re-lay it out? +//! +//! Issue #1576 §13.4 gates the disk tier at 1.2x write amplification: physical +//! bytes written must stay close to the newly committed unique payload bytes. +//! The L3 store cuts content-addressed segments at fixed byte offsets, so that +//! gate holds only if turn N+1's payload keeps turn N's bytes at the same +//! offsets. If the runtime lays the page out layer-major, every layer block +//! after the first shifts as tokens are added, every downstream segment +//! re-digests, and each turn rewrites the whole prefix. +//! +//! This runner answers that empirically before the on-disk format is fixed: it +//! grows one session a turn at a time, exports the full page after each turn, +//! and reports how much of the previous export survives — byte-identical +//! prefix, and reusable segments at the store's segment size. It also probes +//! the alternative design, exporting only the new token window +//! (`token_start > 0`), which would make growth append-only at the manifest +//! level regardless of internal layout. + +use anyhow::{Context, Result, bail}; +use serde::Serialize; +use skippy_cache::{ + ExactStatePayload, GeometryBlock, GeometryKind, L3Tier, PayloadGeometry, StoreLimits, +}; +use skippy_runtime::{ + GGML_TYPE_F16, MtpSource, RuntimeConfig, RuntimeKvPage, RuntimeKvPageDesc, StageModel, + StageSession, +}; + +use super::{ + stage_execution::{runtime_flash_attn, runtime_load_mode}, + state_handoff::state_handoff_tokens, +}; +use crate::cli::KvPageGrowthArgs; + +#[derive(Debug, Serialize)] +struct TurnReport { + turn: usize, + token_count: u64, + payload_bytes: u64, + /// Bytes of this payload identical to the previous turn's from offset 0. + common_prefix_bytes: u64, + segments: usize, + /// Segments whose digest was already stored by an earlier turn. + reused_segments: usize, + /// What the store would physically write for this turn, cutting at fixed + /// byte offsets. + physical_bytes: u64, + /// What the real L3 tier writes for this turn, cutting on the page's own + /// geometry. This is the number that ships. + tier_physical_bytes: u64, + /// `tier_physical_bytes` over `ideal_bytes`. + tier_amplification: f64, + /// Whether the tier accepted the geometry rather than falling back. + geometry_accepted: bool, + /// What an ideal append-only layout would write: the payload growth. + ideal_bytes: u64, + /// physical / ideal. The §13.4 gate is 1.2. + amplification: f64, + /// Whether exporting just this turn's token window succeeded. + windowed_export: WindowedProbe, +} + +#[derive(Debug, Serialize)] +#[serde(tag = "status", rename_all = "kebab-case")] +enum WindowedProbe { + Ok { payload_bytes: u64 }, + Unsupported { error: String }, + Skipped, +} + +#[derive(Debug, Serialize)] +struct GrowthReport { + model: String, + segment_bytes: u64, + base_tokens: usize, + turn_tokens: usize, + turns: Vec, + /// Worst per-turn amplification: the number the gate has to clear. + max_amplification: f64, + /// The same for the real tier, which is what §13.4 actually gates. + max_tier_amplification: f64, + /// True when every turn kept the previous payload byte-identical at offset 0. + append_only: bool, + /// Context the probe actually ran with, raised from --ctx-size when that + /// was too small for the configured workload. + ctx_size: u32, +} + +pub fn kv_page_growth(args: KvPageGrowthArgs) -> Result<()> { + if args.turns == 0 { + bail!("--turns must be at least 1"); + } + if args.turn_tokens == 0 { + bail!("--turn-tokens must be at least 1"); + } + if args.segment_bytes == 0 { + bail!("--segment-bytes must be greater than zero"); + } + let total_tokens = args + .base_tokens + .checked_add(args.turns.saturating_mul(args.turn_tokens)) + .context("token budget overflow")?; + // The shared --ctx-size default (128) is far below this probe's own + // workload default (2048 + 4 x 512), so the documented invocation would + // always fail its own budget check. The context is an implementation + // detail of the measurement, not a workload parameter: size it to fit. + // The effective value is reported, so a raise is never silent. + let required_ctx = u32::try_from(total_tokens).context("token budget exceeds u32")?; + let ctx_size = args.runtime.ctx_size.max(required_ctx); + + let config = RuntimeConfig { + stage_index: 0, + layer_start: 0, + layer_end: args.runtime.layer_end, + ctx_size, + lane_count: 1, + n_batch: args.runtime.n_batch, + n_ubatch: args.runtime.n_ubatch, + n_threads: None, + n_threads_batch: None, + n_gpu_layers: args.runtime.n_gpu_layers, + mmap: None, + mlock: false, + repack: false, + op_offload: None, + no_host_buffer: false, + check_tensors: false, + direct_io: false, + main_gpu: None, + split_mode: skippy_runtime::SplitMode::Auto, + selected_backend_device: None, + load_mode: runtime_load_mode(args.runtime.stage_load_mode), + projector_path: None, + projector_use_gpu: None, + checkpoint_quantization: skippy_runtime::CheckpointQuantization::Preserve, + checkpoint_imatrix: None, + checkpoint_imatrix_sha256: None, + media_marker: None, + image_min_tokens: None, + image_max_tokens: None, + batch_max_tokens: None, + glm_dsa_policy: skippy_runtime::GlmDsaPolicy::Auto, + include_embeddings: true, + include_output: false, + mtp_source: MtpSource::Disabled, + filter_tensors_on_load: true, + resident_tensor_names: Vec::new(), + cache_type_k: GGML_TYPE_F16, + cache_type_v: GGML_TYPE_F16, + flash_attn_type: runtime_flash_attn(args.runtime.flash_attn), + kv_offload: None, + kv_unified: None, + swa_full: None, + }; + let model = StageModel::open(&args.runtime.model, &config) + .context("failed to open stage model for KV page growth probe")?; + let tokens = state_handoff_tokens(&model, &args.runtime.prompt, Some(total_tokens)) + .context("failed to build growth prefix")?; + if tokens.len() < total_tokens { + bail!( + "prefix expansion produced {} tokens, needed {total_tokens}", + tokens.len() + ); + } + let layer_end = i32::try_from(args.runtime.layer_end).context("layer_end exceeds i32")?; + let mut session = model + .create_session() + .context("failed to create growth probe session")?; + + // The real store, so the reported number is the one that ships rather than + // a simulation of it. A fresh root per run: leftover segments from an + // earlier run would dedupe against this one and flatter the result. + let tier_root = std::env::temp_dir().join(format!( + "skippy-kv-page-growth-{}-{}", + std::process::id(), + args.segment_bytes + )); + let _ = std::fs::remove_dir_all(&tier_root); + let tier = L3Tier::open_with_limits( + &tier_root, + StoreLimits::new(0, 0), + "blake3:kv-page-growth".to_string(), + usize::try_from(args.segment_bytes).context("segment size exceeds usize")?, + ) + .context("open L3 tier for the probe")?; + let mut tier_written = 0u64; + + let mut seen_digests: std::collections::HashSet = std::collections::HashSet::new(); + let mut previous_payload: Vec = Vec::new(); + let mut turns = Vec::with_capacity(args.turns); + let mut cursor = 0usize; + + for turn in 0..=args.turns { + let next = if turn == 0 { + args.base_tokens + } else { + cursor + args.turn_tokens + }; + session + .prefill_chunked(&tokens[cursor..next]) + .with_context(|| format!("prefill failed on turn {turn}"))?; + let previous_cursor = cursor; + cursor = next; + + let page = export_page(&mut session, layer_end, 0, cursor as u64) + .with_context(|| format!("full-prefix KV export failed on turn {turn}"))?; + let payload = page.payload.clone(); + + let common_prefix_bytes = common_prefix_len(&previous_payload, &payload) as u64; + let mut segments = 0usize; + let mut reused_segments = 0usize; + let mut physical_bytes = 0u64; + for chunk in payload.chunks(usize::try_from(args.segment_bytes)?) { + segments += 1; + let digest = blake3::hash(chunk).to_hex().to_string(); + if seen_digests.insert(digest) { + physical_bytes += chunk.len() as u64; + } else { + reused_segments += 1; + } + } + let ideal_bytes = (payload.len() as u64).saturating_sub(previous_payload.len() as u64); + // Turn 0 commits the whole base prefix, which is genuinely new data. + let ideal_bytes = if turn == 0 { + payload.len() as u64 + } else { + ideal_bytes + }; + let amplification = if ideal_bytes == 0 { + f64::INFINITY + } else { + physical_bytes as f64 / ideal_bytes as f64 + }; + + // Spill through the real tier with the geometry the server derives + // from the page descriptor, and read what it actually wrote. + let geometry = page_geometry(&page.desc, payload.len() as u64, args.segment_bytes); + let exact_state = ExactStatePayload::kv_recurrent(payload.clone(), Vec::new()); + let geometry_rejected_before = tier + .status() + .context("tier status before spill")? + .activity + .geometry_rejected; + tier.spill( + "probe", + &tokens[..cursor], + &exact_state, + None, + geometry.as_ref(), + ) + .with_context(|| format!("tier spill failed on turn {turn}"))?; + let activity = tier.status().context("tier status")?.activity; + let geometry_rejected_delta = activity + .geometry_rejected + .saturating_sub(geometry_rejected_before); + let tier_physical_bytes = activity.bytes_written - tier_written; + tier_written = activity.bytes_written; + let tier_amplification = if ideal_bytes == 0 { + f64::INFINITY + } else { + tier_physical_bytes as f64 / ideal_bytes as f64 + }; + let geometry_accepted = geometry.is_some() && geometry_rejected_delta == 0; + + // The alternative design: store only the new window rather than + // re-cutting the whole prefix. Turn 0 has no preceding window. + let windowed_export = if turn == 0 { + WindowedProbe::Skipped + } else { + match export_page( + &mut session, + layer_end, + previous_cursor as u64, + (cursor - previous_cursor) as u64, + ) { + Ok(page) => WindowedProbe::Ok { + payload_bytes: page.payload.len() as u64, + }, + Err(error) => WindowedProbe::Unsupported { + error: format!("{error:#}"), + }, + } + }; + + println!( + "turn={turn} tokens={cursor} payload={} common_prefix={common_prefix_bytes} \ + segments={segments} reused={reused_segments} physical={physical_bytes} \ + ideal={ideal_bytes} amplification={amplification:.2} \ + tier_physical={tier_physical_bytes} tier_amplification={tier_amplification:.2} \ + geometry={geometry_accepted}", + payload.len() + ); + + turns.push(TurnReport { + turn, + token_count: cursor as u64, + payload_bytes: payload.len() as u64, + common_prefix_bytes, + segments, + reused_segments, + physical_bytes, + ideal_bytes, + amplification, + tier_physical_bytes, + tier_amplification, + geometry_accepted, + windowed_export, + }); + previous_payload = payload; + } + + let append_only = turns.windows(2).all(|pair| { + let [previous, current] = pair else { + return true; + }; + current.common_prefix_bytes >= previous.payload_bytes + }); + let max_amplification = turns + .iter() + .skip(1) + .map(|turn| turn.amplification) + .fold(0.0_f64, f64::max); + let max_tier_amplification = turns + .iter() + .skip(1) + .map(|turn| turn.tier_amplification) + .fold(0.0_f64, f64::max); + + let report = GrowthReport { + model: args.runtime.model.display().to_string(), + segment_bytes: args.segment_bytes, + base_tokens: args.base_tokens, + turn_tokens: args.turn_tokens, + turns, + max_amplification, + max_tier_amplification, + append_only, + ctx_size, + }; + println!( + "kv_page_growth append_only={append_only} ctx_size={ctx_size} \ + fixed_amplification={max_amplification:.2} \ + tier_amplification={max_tier_amplification:.2} gate=1.20 verdict={}", + if max_tier_amplification <= 1.2 { + "geometry cutting holds the gate" + } else { + "FAILS the gate" + } + ); + let _ = std::fs::remove_dir_all(&tier_root); + if let Some(path) = args.json.as_deref() { + let rendered = serde_json::to_string_pretty(&report).context("render growth report")?; + std::fs::write(path, rendered) + .with_context(|| format!("write growth report {}", path.display()))?; + } + // The report is written first so a failing run is still diagnosable, but a + // failed gate must fail the command: certification automation reads the + // exit status, not the printed verdict. + if max_tier_amplification > 1.2 { + bail!( + "kv page growth gate failed: tier amplification {max_tier_amplification:.2} exceeds 1.20" + ); + } + Ok(()) +} + +fn export_page( + session: &mut StageSession, + layer_end: i32, + token_start: u64, + token_count: u64, +) -> Result { + session.export_kv_page(0, layer_end, token_start, token_count) +} + +/// Mirrors `kv_page_geometry` in `skippy-server`'s kv_integration: every +/// layer's K rows, then every layer's V rows, in fixed token windows. +fn page_geometry( + desc: &RuntimeKvPageDesc, + payload_bytes: u64, + segment_bytes: u64, +) -> Option { + if desc.component_count != 0 || desc.token_count == 0 || desc.layer_count == 0 { + return None; + } + let k_stride = u64::from(desc.k_row_bytes); + let v_stride = u64::from(desc.v_row_bytes); + if k_stride == 0 || desc.flags & skippy_runtime::KV_PAGE_FLAG_V_TRANSPOSED != 0 { + return None; + } + let mut blocks = Vec::new(); + for layer in 0..desc.layer_count { + blocks.push(GeometryBlock { + stride: k_stride, + kind: GeometryKind::Key, + layer, + column: 0, + }); + } + if v_stride > 0 { + for layer in 0..desc.layer_count { + blocks.push(GeometryBlock { + stride: v_stride, + kind: GeometryKind::Value, + layer, + column: 0, + }); + } + } + let widest = blocks.iter().map(|block| block.stride).max()?; + let window_rows = (segment_bytes / widest.max(1)) + .clamp(1, 512) + .next_power_of_two() + .min(512); + let geometry = PayloadGeometry { + blocks, + rows: desc.token_count, + window_rows, + tail_bytes: 0, + }; + let tail = payload_bytes.checked_sub(geometry.total_bytes())?; + Some(PayloadGeometry { + tail_bytes: tail, + ..geometry + }) +} + +fn common_prefix_len(left: &[u8], right: &[u8]) -> usize { + left.iter() + .zip(right) + .take_while(|(left, right)| left == right) + .count() +} diff --git a/crates/skippy-correctness/src/runner/mod.rs b/crates/skippy-correctness/src/runner/mod.rs index a851a36131..5d7f440604 100644 --- a/crates/skippy-correctness/src/runner/mod.rs +++ b/crates/skippy-correctness/src/runner/mod.rs @@ -1,5 +1,8 @@ +mod cachegen_gate; +mod kv_page_growth; pub(crate) mod native_mtp; mod prediction_return; +mod remote_handoff; mod single_step; mod split_chain; mod split_prefix_hit; @@ -7,6 +10,8 @@ mod stage_execution; mod stage_fa_parity; mod state_handoff; +pub use kv_page_growth::kv_page_growth; +pub use remote_handoff::remote_handoff; pub use single_step::single_step; pub use split_chain::{chain, split_scan}; pub use split_prefix_hit::split_prefix_hit; diff --git a/crates/skippy-correctness/src/runner/remote_handoff.rs b/crates/skippy-correctness/src/runner/remote_handoff.rs new file mode 100644 index 0000000000..3c638a7010 --- /dev/null +++ b/crates/skippy-correctness/src/runner/remote_handoff.rs @@ -0,0 +1,1833 @@ +use std::{ + io::{BufReader, BufWriter, Read, Write}, + net::{TcpListener, TcpStream}, + time::{Duration, Instant}, +}; + +use anyhow::{Context, Result, bail}; +use serde::{Deserialize, Serialize, de::DeserializeOwned}; +use skippy_cache::{ + HandoffManifest, HandoffSegmentRef, HandoffSegmentStore, KvFetchClient, SegmentCodecIdentity, + serve_store_with_timeout, +}; +use skippy_runtime::{ + GGML_TYPE_F16, MtpSource, RuntimeConfig, RuntimeKvPageDesc, StageModel, StageSession, +}; + +use crate::{ + cli::{RemoteHandoffArgs, RemoteHandoffRole, StatePayloadKind}, + report::{RemoteHandoffReceiverTimings, RemoteHandoffReport}, +}; + +use super::native_mtp::emit_report; +use super::stage_execution::{ + PackageStageSpec, elapsed_ms, ensure_matches, protocol_load_mode, runtime_flash_attn, + runtime_load_mode, runtime_model_identity, stage_id_for_index, stage_model_resolution, status, +}; +use super::state_handoff::state_handoff_tokens; + +mod identity; +use identity::{effective_payload_kind, state_identity_for}; + +const PROTOCOL_VERSION: u32 = 1; +const MAX_HEADER_BYTES: u64 = 16 * 1024 * 1024; +const MAX_SEGMENT_BYTES: u64 = 256 * 1024 * 1024; +const STREAM_BUFFER_BYTES: usize = 1024 * 1024; + +mod frame_kind { + pub const HELLO: u8 = 1; + pub const HELLO_ACK: u8 = 2; + pub const SEGMENT: u8 = 3; + pub const COMMIT: u8 = 4; + pub const VERIFY: u8 = 5; + pub const RESULT: u8 = 6; + pub const PAGE: u8 = 7; + pub const RECURRENT: u8 = 8; +} + +#[derive(Serialize, Deserialize)] +struct HelloHeader { + protocol_version: u32, + model_id: String, + layer_end: u32, + ctx_size: u32, + state_payload_kind: String, + flash_attn: String, + decode_token_count: usize, + lane_count: u32, + state_identity: String, +} + +#[derive(Serialize, Deserialize)] +struct HelloAckHeader { + ok: bool, + reason: Option, +} + +#[derive(Serialize, Deserialize)] +struct SegmentHeader { + index: usize, + offset: u64, + payload_bytes: u64, + blake3: String, +} + +#[derive(Serialize, Deserialize)] +struct PageHeader { + index: usize, + token_start: u64, + token_count: u64, + payload_bytes: u64, + blake3: String, + kv_desc: RuntimeKvPageDesc, +} + +#[derive(Serialize, Deserialize)] +struct RecurrentHeader { + payload_bytes: u64, + blake3: String, +} + +/// Per-segment metadata carried in `HandoffSegmentRef::meta_json` for +/// page-stream manifests. +#[derive(Serialize, Deserialize)] +struct PageSegmentMeta { + kind: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + kv_desc: Option, + #[serde(default)] + token_start: u64, + #[serde(default)] + token_count: u64, +} + +#[derive(Serialize, Deserialize)] +struct CommitHeader { + segment_count: usize, + total_bytes: u64, + payload_blake3: String, + prompt_token_count: u64, + continuation_token: i32, + kv_bytes: u64, + recurrent_bytes: u64, + kv_desc: Option, + prefix_tokens: Vec, + run_baseline: bool, + #[serde(default)] + streaming: bool, + #[serde(default)] + page_count: usize, +} + +#[derive(Serialize, Deserialize)] +struct VerifyHeader { + source_tokens: Vec, +} + +#[derive(Serialize, Deserialize)] +struct ResultHeader { + ok: bool, + reason: Option, + restored_tokens: Vec, + baseline_tokens: Vec, + tokens_match: bool, + baseline_matches: Option, + timings: RemoteHandoffReceiverTimings, +} + +struct StatePayload { + kv_desc: Option, + kv: Vec, + recurrent: Vec, +} + +impl StatePayload { + fn full_state(bytes: Vec) -> Self { + Self { + kv_desc: None, + kv: bytes, + recurrent: Vec::new(), + } + } + + fn wire_bytes(&self) -> Vec { + let mut bytes = Vec::with_capacity(self.kv.len() + self.recurrent.len()); + bytes.extend_from_slice(&self.kv); + bytes.extend_from_slice(&self.recurrent); + bytes + } +} + +fn effective_lane_count(args: &RemoteHandoffArgs) -> u32 { + // Full-state blobs cover the whole context, whose KV layout depends on + // the lane count, so sender and receiver must open identically shaped + // contexts. + args.runtime_lane_count.unwrap_or(2).max(2) +} + +pub fn remote_handoff(args: RemoteHandoffArgs) -> Result<()> { + match args.role { + RemoteHandoffRole::Send if args.streaming => run_sender_streaming(args), + RemoteHandoffRole::Send => run_sender(args), + RemoteHandoffRole::Recv => run_receiver(args), + RemoteHandoffRole::Restore => run_restore(args), + RemoteHandoffRole::Serve => run_serve(args), + RemoteHandoffRole::Fetch => run_fetch(args), + } +} + +/// Serve the local L3 store to peers over the `skippy-kv/1` fetch protocol. +fn run_serve(args: RemoteHandoffArgs) -> Result<()> { + let store_dir = args + .store_dir + .clone() + .context("--role serve requires --store-dir")?; + let store = HandoffSegmentStore::open(&store_dir, args.store_budget_bytes)?; + let listener = std::net::TcpListener::bind(args.listen) + .with_context(|| format!("failed to bind skippy-kv listener on {}", args.listen))?; + eprintln!( + "skippy-kv store server ready on {} serving {} ({} manifests)", + args.listen, + store_dir.display(), + store.list_manifests()?.len() + ); + serve_store_with_timeout( + &store, + &listener, + args.accept_count, + Duration::from_secs(args.handshake_timeout_secs.max(1)), + ) +} + +/// Pull a manifest and its segments from a peer's store into the local one, +/// then restore and decode from it to prove the fetched state is usable — +/// cross-node prefix reuse without a push handoff. +fn run_fetch(mut args: RemoteHandoffArgs) -> Result<()> { + let peer = args + .peer + .context("--role fetch requires --peer ")?; + let store_dir = args + .store_dir + .clone() + .context("--role fetch requires --store-dir")?; + let store = HandoffSegmentStore::open(&store_dir, args.store_budget_bytes)?; + let fetch_started = Instant::now(); + let mut client = KvFetchClient::connect_with_timeout( + &peer.to_string(), + Duration::from_secs(args.handshake_timeout_secs.max(1)), + )?; + let (manifest, stats) = client + .fetch_into_store(args.manifest.as_deref(), &store) + .context("failed to fetch manifest from peer")?; + drop(client); + let fetch_ms = elapsed_ms(fetch_started); + eprintln!( + "skippy-kv fetched manifest {} from {peer}: {} segments pulled, {} already local, {:.1} MiB in {fetch_ms:.0} ms ({:.2} Gbps)", + manifest.payload_digest, + stats.segments_fetched, + stats.segments_skipped, + stats.bytes_fetched as f64 / (1024.0 * 1024.0), + transfer_gbps(stats.bytes_fetched as usize, fetch_ms), + ); + args.manifest = Some(manifest.payload_digest.clone()); + args.streaming = manifest.payload_kind == "kv-page-stream"; + run_restore(args) +} + +fn open_store(args: &RemoteHandoffArgs) -> Result> { + args.store_dir + .as_ref() + .map(|dir| HandoffSegmentStore::open(dir, args.store_budget_bytes)) + .transpose() +} + +fn manifest_from_commit( + commit: &CommitHeader, + segments: Vec, + state_identity: String, + payload_kind: &str, + expected_tokens: Vec, +) -> Result { + let mut manifest = HandoffManifest::new(state_identity, payload_kind.to_string()); + manifest.total_bytes = commit.total_bytes; + manifest.payload_digest = commit.payload_blake3.clone(); + manifest.segments = segments; + manifest.kv_bytes = commit.kv_bytes; + manifest.recurrent_bytes = commit.recurrent_bytes; + manifest.kv_desc_json = commit + .kv_desc + .as_ref() + .map(serde_json::to_string) + .transpose() + .context("failed to serialize kv desc for manifest")?; + manifest.token_count = commit.prompt_token_count; + manifest.continuation_token = commit.continuation_token; + manifest.expected_tokens = expected_tokens; + Ok(manifest) +} + +fn open_full_model(args: &RemoteHandoffArgs, lane_count: u32) -> Result { + let model_identity = runtime_model_identity(&args.runtime)?; + let spec = PackageStageSpec { + topology_id: "correctness-remote-handoff", + stage_id: stage_id_for_index(0), + stage_index: 0, + layer_start: 0, + layer_end: args.runtime.layer_end, + include_embeddings: true, + include_output: true, + }; + let resolution = stage_model_resolution( + &args.runtime.model, + args.runtime.stage_model.as_ref(), + args.runtime.stage_load_mode, + &model_identity, + spec, + )?; + let runtime_config = RuntimeConfig { + stage_index: 0, + layer_start: 0, + layer_end: args.runtime.layer_end, + ctx_size: args.runtime.ctx_size, + lane_count, + n_batch: args.runtime.n_batch, + n_ubatch: args.runtime.n_ubatch, + n_threads: None, + n_threads_batch: None, + n_gpu_layers: args.runtime.n_gpu_layers, + mmap: None, + mlock: false, + repack: false, + op_offload: None, + no_host_buffer: false, + check_tensors: false, + direct_io: false, + main_gpu: None, + split_mode: skippy_runtime::SplitMode::Auto, + selected_backend_device: None, + load_mode: runtime_load_mode(args.runtime.stage_load_mode), + projector_path: None, + projector_use_gpu: None, + media_marker: None, + image_min_tokens: None, + image_max_tokens: None, + batch_max_tokens: None, + glm_dsa_policy: skippy_runtime::GlmDsaPolicy::Auto, + include_embeddings: true, + include_output: true, + mtp_source: MtpSource::Disabled, + filter_tensors_on_load: false, + resident_tensor_names: Vec::new(), + checkpoint_quantization: skippy_runtime::CheckpointQuantization::Preserve, + checkpoint_imatrix: None, + checkpoint_imatrix_sha256: None, + cache_type_k: GGML_TYPE_F16, + cache_type_v: GGML_TYPE_F16, + flash_attn_type: runtime_flash_attn(args.runtime.flash_attn), + kv_offload: None, + kv_unified: None, + swa_full: None, + }; + StageModel::open(&resolution.path, &runtime_config) + .context("failed to open remote handoff model") +} + +fn greedy_decode(session: &mut StageSession, first_token: i32, count: usize) -> Result> { + let mut tokens = Vec::with_capacity(count); + let mut next = first_token; + for _ in 0..count { + let predicted = session + .decode_step(next) + .context("greedy decode step failed")?; + tokens.push(predicted); + next = predicted; + } + Ok(tokens) +} + +fn export_state_payload( + session: &mut StageSession, + args: &RemoteHandoffArgs, + token_count: u64, +) -> Result { + match args.state_payload_kind { + StatePayloadKind::FullState => Ok(StatePayload::full_state( + session.export_full_state(0, args.runtime.layer_end as i32)?, + )), + StatePayloadKind::KvRecurrent => { + let page = session.export_kv_page(0, args.runtime.layer_end as i32, 0, token_count)?; + let recurrent = session.export_recurrent_state()?; + Ok(StatePayload { + kv_desc: Some(page.desc), + kv: page.payload, + recurrent, + }) + } + other => bail!("remote handoff does not support state payload kind {other:?}"), + } +} + +fn import_state_payload( + session: &mut StageSession, + args: &RemoteHandoffArgs, + commit: &CommitHeader, + bytes: &[u8], +) -> Result<()> { + let kv_bytes = usize::try_from(commit.kv_bytes).context("kv byte count exceeds usize")?; + let recurrent_bytes = + usize::try_from(commit.recurrent_bytes).context("recurrent byte count exceeds usize")?; + if kv_bytes + recurrent_bytes != bytes.len() { + bail!( + "commit component sizes {} + {} do not cover payload of {} bytes", + kv_bytes, + recurrent_bytes, + bytes.len() + ); + } + match &commit.kv_desc { + None => session.import_full_state_for_token_count( + 0, + args.runtime.layer_end as i32, + &bytes[..kv_bytes], + commit.prompt_token_count, + ), + Some(kv_desc) => { + session.import_kv_page(kv_desc, &bytes[..kv_bytes])?; + session.import_recurrent_state_for_token_count( + &bytes[kv_bytes..], + commit.prompt_token_count, + ) + } + } +} + +fn run_sender(args: RemoteHandoffArgs) -> Result<()> { + let peer = args + .peer + .context("--role send requires --peer ")?; + let model_identity = runtime_model_identity(&args.runtime)?; + let report_out = args.output.report_out.clone(); + + let model_load_started = Instant::now(); + let model = open_full_model(&args, effective_lane_count(&args))?; + let model_load_ms = elapsed_ms(model_load_started); + + let tokenize_started = Instant::now(); + let tokens = state_handoff_tokens(&model, &args.runtime.prompt, args.prefix_token_count) + .context("failed to tokenize remote handoff prompt")?; + let split = args.prefix_token_count.unwrap_or(tokens.len() - 1); + let prefix = tokens[..split].to_vec(); + let continuation = tokens[split]; + let tokenize_ms = elapsed_ms(tokenize_started); + + let stream = TcpStream::connect(peer) + .with_context(|| format!("failed to connect to receiver at {peer}"))?; + stream.set_nodelay(true).ok(); + let mut reader = BufReader::with_capacity(STREAM_BUFFER_BYTES, stream.try_clone()?); + let mut writer = BufWriter::with_capacity(STREAM_BUFFER_BYTES, stream); + + write_frame( + &mut writer, + frame_kind::HELLO, + &HelloHeader { + protocol_version: PROTOCOL_VERSION, + model_id: model_identity.model_id.clone(), + layer_end: args.runtime.layer_end, + ctx_size: args.runtime.ctx_size, + state_payload_kind: effective_payload_kind(&args).to_string(), + flash_attn: format!("{:?}", args.runtime.flash_attn), + decode_token_count: args.decode_tokens, + lane_count: effective_lane_count(&args), + state_identity: state_identity_for(&args, &model_identity), + }, + &[], + )?; + writer.flush().context("failed to flush hello frame")?; + let (ack, _) = read_frame_expect::(&mut reader, frame_kind::HELLO_ACK)?; + if !ack.ok { + bail!( + "receiver rejected handoff: {}", + ack.reason.unwrap_or_else(|| "no reason given".to_string()) + ); + } + + let mut session = model + .create_session() + .context("failed to create sender session")?; + let prefill_started = Instant::now(); + session + .prefill_chunked(&prefix) + .context("sender prefill failed")?; + let source_prefill_ms = elapsed_ms(prefill_started); + + let export_started = Instant::now(); + let payload = export_state_payload(&mut session, &args, prefix.len() as u64)?; + let state_export_ms = elapsed_ms(export_started); + + let wire_bytes = payload.wire_bytes(); + let payload_digest = digest_hex(&wire_bytes); + let segment_bytes = args.segment_bytes.max(1); + let transfer_started = Instant::now(); + let mut segment_count = 0usize; + for (index, chunk) in wire_bytes.chunks(segment_bytes).enumerate() { + write_frame( + &mut writer, + frame_kind::SEGMENT, + &SegmentHeader { + index, + offset: (index * segment_bytes) as u64, + payload_bytes: chunk.len() as u64, + blake3: digest_hex(chunk), + }, + chunk, + )?; + segment_count += 1; + } + write_frame( + &mut writer, + frame_kind::COMMIT, + &CommitHeader { + segment_count, + total_bytes: wire_bytes.len() as u64, + payload_blake3: payload_digest.clone(), + prompt_token_count: prefix.len() as u64, + continuation_token: continuation, + kv_bytes: payload.kv.len() as u64, + recurrent_bytes: payload.recurrent.len() as u64, + kv_desc: payload.kv_desc.clone(), + prefix_tokens: if args.baseline { + prefix.clone() + } else { + Vec::new() + }, + run_baseline: args.baseline, + streaming: false, + page_count: 0, + }, + &[], + )?; + writer.flush().context("failed to flush handoff stream")?; + let transfer_ms = elapsed_ms(transfer_started); + + // Spill to the local L3 store off the transfer critical path: the + // receiver is already importing while these writes land on disk. + let store = open_store(&args)?; + let store_started = Instant::now(); + let mut segment_refs = Vec::new(); + if let Some(store) = &store { + for (index, chunk) in wire_bytes.chunks(segment_bytes).enumerate() { + let digest = store + .put_segment(chunk) + .context("sender segment spill failed")? + .digest; + segment_refs.push(HandoffSegmentRef { + index: index as u32, + offset: (index * segment_bytes) as u64, + bytes: chunk.len() as u64, + digest, + codec_identity: Some(SegmentCodecIdentity::raw(chunk.len() as u64)), + meta_json: None, + }); + } + } + let mut store_ms = elapsed_ms(store_started); + + // The receiver imports and decodes while the sender produces the + // reference continuation, mirroring how both phases overlap in serving. + let source_decode_started = Instant::now(); + let source_tokens = greedy_decode(&mut session, continuation, args.decode_tokens)?; + let source_decode_ms = elapsed_ms(source_decode_started); + + if let Some(store) = &store { + let commit_started = Instant::now(); + let manifest = manifest_from_commit( + &CommitHeader { + segment_count, + total_bytes: wire_bytes.len() as u64, + payload_blake3: payload_digest.clone(), + prompt_token_count: prefix.len() as u64, + continuation_token: continuation, + kv_bytes: payload.kv.len() as u64, + recurrent_bytes: payload.recurrent.len() as u64, + kv_desc: payload.kv_desc.clone(), + prefix_tokens: Vec::new(), + run_baseline: false, + streaming: false, + page_count: 0, + }, + segment_refs, + state_identity_for(&args, &model_identity), + payload_kind_name(args.state_payload_kind), + source_tokens.clone(), + )?; + store + .commit(&manifest) + .context("sender manifest commit failed")?; + store_ms += elapsed_ms(commit_started); + } + + write_frame( + &mut writer, + frame_kind::VERIFY, + &VerifyHeader { + source_tokens: source_tokens.clone(), + }, + &[], + )?; + writer.flush().context("failed to flush verify frame")?; + + let (result, _) = read_frame_expect::(&mut reader, frame_kind::RESULT)?; + if !result.ok { + bail!( + "receiver failed to complete handoff: {}", + result + .reason + .unwrap_or_else(|| "no reason given".to_string()) + ); + } + + let timings = result.timings; + let ttft_disaggregated_ms = source_prefill_ms + + state_export_ms + + transfer_ms + + timings.kv_attach_ms + + timings.first_decode_ms; + let ttft_local_ms = (timings.baseline_prefill_ms > 0.0) + .then_some(timings.baseline_prefill_ms + timings.baseline_first_decode_ms); + let matches = result.tokens_match && result.baseline_matches.unwrap_or(true); + let report = RemoteHandoffReport { + mode: "remote-handoff", + status: status(matches), + role: "send", + model_identity, + matches, + tokens_match: result.tokens_match, + baseline_matches: result.baseline_matches, + state_payload_kind: payload_kind_name(args.state_payload_kind), + prompt_token_count: prefix.len(), + decode_token_count: args.decode_tokens, + continuation_token: continuation, + source_tokens, + restored_tokens: result.restored_tokens, + state_bytes: wire_bytes.len(), + state_bytes_per_prompt_token: wire_bytes.len() as f64 / prefix.len().max(1) as f64, + kv_bytes: payload.kv.len(), + recurrent_bytes: payload.recurrent.len(), + segment_count, + segment_bytes, + payload_digest, + model_load_ms, + tokenize_ms, + source_prefill_ms, + state_export_ms, + transfer_ms, + transfer_gbps: transfer_gbps(wire_bytes.len(), transfer_ms), + source_decode_ms, + store_ms: store.is_some().then_some(store_ms), + overlap_wall_ms: None, + receiver: timings, + ttft_disaggregated_ms, + ttft_local_ms, + ttft_speedup: ttft_local_ms.map(|local| local / ttft_disaggregated_ms.max(f64::EPSILON)), + }; + emit_report(&report, report_out.as_deref())?; + ensure_matches(matches, args.allow_mismatch)?; + Ok(()) +} + +/// Streaming sender: after each prefill chunk, the KV page for that token +/// range is exported and streamed while the next chunk computes, so transfer +/// hides behind prefill. The recurrent snapshot (when the family has one) +/// and the commit record are the uncovered tail; the receiver stages pages +/// as they arrive but cannot generate until the commit validates. +fn run_sender_streaming(args: RemoteHandoffArgs) -> Result<()> { + let peer = args + .peer + .context("--role send requires --peer ")?; + let model_identity = runtime_model_identity(&args.runtime)?; + let report_out = args.output.report_out.clone(); + + let model_load_started = Instant::now(); + let model = open_full_model(&args, effective_lane_count(&args))?; + let model_load_ms = elapsed_ms(model_load_started); + + let tokenize_started = Instant::now(); + let tokens = state_handoff_tokens(&model, &args.runtime.prompt, args.prefix_token_count) + .context("failed to tokenize remote handoff prompt")?; + let split = args.prefix_token_count.unwrap_or(tokens.len() - 1); + let prefix = tokens[..split].to_vec(); + let continuation = tokens[split]; + let tokenize_ms = elapsed_ms(tokenize_started); + + let stream = TcpStream::connect(peer) + .with_context(|| format!("failed to connect to receiver at {peer}"))?; + stream.set_nodelay(true).ok(); + let mut reader = BufReader::with_capacity(STREAM_BUFFER_BYTES, stream.try_clone()?); + let mut writer = BufWriter::with_capacity(STREAM_BUFFER_BYTES, stream); + + write_frame( + &mut writer, + frame_kind::HELLO, + &HelloHeader { + protocol_version: PROTOCOL_VERSION, + model_id: model_identity.model_id.clone(), + layer_end: args.runtime.layer_end, + ctx_size: args.runtime.ctx_size, + state_payload_kind: effective_payload_kind(&args).to_string(), + flash_attn: format!("{:?}", args.runtime.flash_attn), + decode_token_count: args.decode_tokens, + lane_count: effective_lane_count(&args), + state_identity: state_identity_for(&args, &model_identity), + }, + &[], + )?; + writer.flush().context("failed to flush hello frame")?; + let (ack, _) = read_frame_expect::(&mut reader, frame_kind::HELLO_ACK)?; + if !ack.ok { + bail!( + "receiver rejected handoff: {}", + ack.reason.unwrap_or_else(|| "no reason given".to_string()) + ); + } + + let mut session = model + .create_session() + .context("failed to create sender session")?; + let chunk_tokens = args.stream_chunk_tokens.max(1); + let store = open_store(&args)?; + let overlap_started = Instant::now(); + let mut source_prefill_ms = 0.0f64; + let mut state_export_ms = 0.0f64; + let mut transfer_ms = 0.0f64; + let mut store_ms = 0.0f64; + let mut payload_hasher = blake3::Hasher::new(); + let mut segment_refs: Vec = Vec::new(); + let mut total_bytes = 0u64; + let mut page_count = 0usize; + for (index, chunk) in prefix.chunks(chunk_tokens).enumerate() { + let token_start = (index * chunk_tokens) as u64; + let prefill_started = Instant::now(); + session + .prefill_chunked(chunk) + .context("sender streaming prefill chunk failed")?; + source_prefill_ms += elapsed_ms(prefill_started); + + let export_started = Instant::now(); + let page = session + .export_kv_page( + 0, + args.runtime.layer_end as i32, + token_start, + chunk.len() as u64, + ) + .with_context(|| format!("failed to export KV page for chunk {index}"))?; + state_export_ms += elapsed_ms(export_started); + + let digest = digest_hex(&page.payload); + payload_hasher.update(&page.payload); + let stream_started = Instant::now(); + write_frame( + &mut writer, + frame_kind::PAGE, + &PageHeader { + index, + token_start, + token_count: chunk.len() as u64, + payload_bytes: page.payload.len() as u64, + blake3: digest.clone(), + kv_desc: page.desc.clone(), + }, + &page.payload, + )?; + // Flush per page so the receiver imports while later chunks prefill. + writer.flush().context("failed to flush page frame")?; + transfer_ms += elapsed_ms(stream_started); + + if let Some(store) = &store { + let store_started = Instant::now(); + store + .put_segment(&page.payload) + .context("sender page spill failed")?; + store_ms += elapsed_ms(store_started); + segment_refs.push(HandoffSegmentRef { + index: index as u32, + offset: total_bytes, + bytes: page.payload.len() as u64, + digest, + codec_identity: Some(SegmentCodecIdentity::native_kv_page( + page.payload.len() as u64 + )), + meta_json: Some(serde_json::to_string(&PageSegmentMeta { + kind: "kv-page".to_string(), + kv_desc: Some(page.desc), + token_start, + token_count: chunk.len() as u64, + })?), + }); + } + total_bytes += page.payload.len() as u64; + page_count += 1; + } + + // Recurrent/SSM state is only final once the whole prompt has prefilled — + // the uncovered tail of the overlap. + let export_started = Instant::now(); + let recurrent: Vec = session.export_recurrent_state().unwrap_or_default(); + state_export_ms += elapsed_ms(export_started); + if !recurrent.is_empty() { + let digest = digest_hex(&recurrent); + payload_hasher.update(&recurrent); + let stream_started = Instant::now(); + write_frame( + &mut writer, + frame_kind::RECURRENT, + &RecurrentHeader { + payload_bytes: recurrent.len() as u64, + blake3: digest.clone(), + }, + &recurrent, + )?; + transfer_ms += elapsed_ms(stream_started); + if let Some(store) = &store { + store + .put_segment(&recurrent) + .context("sender recurrent spill failed")?; + segment_refs.push(HandoffSegmentRef { + index: page_count as u32, + offset: total_bytes, + bytes: recurrent.len() as u64, + digest, + codec_identity: Some(SegmentCodecIdentity::raw(recurrent.len() as u64)), + meta_json: Some(serde_json::to_string(&PageSegmentMeta { + kind: "recurrent".to_string(), + kv_desc: None, + token_start: 0, + token_count: prefix.len() as u64, + })?), + }); + } + total_bytes += recurrent.len() as u64; + } + + let payload_digest = format!("{}", payload_hasher.finalize().to_hex()); + let commit = CommitHeader { + segment_count: page_count + usize::from(!recurrent.is_empty()), + total_bytes, + payload_blake3: payload_digest.clone(), + prompt_token_count: prefix.len() as u64, + continuation_token: continuation, + kv_bytes: total_bytes - recurrent.len() as u64, + recurrent_bytes: recurrent.len() as u64, + kv_desc: None, + prefix_tokens: if args.baseline { + prefix.clone() + } else { + Vec::new() + }, + run_baseline: args.baseline, + streaming: true, + page_count, + }; + write_frame(&mut writer, frame_kind::COMMIT, &commit, &[])?; + writer.flush().context("failed to flush commit frame")?; + let overlap_wall_ms = elapsed_ms(overlap_started); + + let source_decode_started = Instant::now(); + let source_tokens = greedy_decode(&mut session, continuation, args.decode_tokens)?; + let source_decode_ms = elapsed_ms(source_decode_started); + + if let Some(store) = &store { + let commit_started = Instant::now(); + let manifest = manifest_from_commit( + &commit, + segment_refs, + state_identity_for(&args, &model_identity), + effective_payload_kind(&args), + source_tokens.clone(), + )?; + store + .commit(&manifest) + .context("sender manifest commit failed")?; + store_ms += elapsed_ms(commit_started); + } + + write_frame( + &mut writer, + frame_kind::VERIFY, + &VerifyHeader { + source_tokens: source_tokens.clone(), + }, + &[], + )?; + writer.flush().context("failed to flush verify frame")?; + + let (result, _) = read_frame_expect::(&mut reader, frame_kind::RESULT)?; + if !result.ok { + bail!( + "receiver failed to complete handoff: {}", + result + .reason + .unwrap_or_else(|| "no reason given".to_string()) + ); + } + + let timings = result.timings; + // Streaming TTFT: page transfer and import hide inside the prefill wall; + // only the commit tail and the first decode remain serial. + let ttft_disaggregated_ms = + overlap_wall_ms + timings.attach_residual_ms + timings.first_decode_ms; + let ttft_local_ms = (timings.baseline_prefill_ms > 0.0) + .then_some(timings.baseline_prefill_ms + timings.baseline_first_decode_ms); + let matches = result.tokens_match && result.baseline_matches.unwrap_or(true); + let report = RemoteHandoffReport { + mode: "remote-handoff", + status: status(matches), + role: "send", + model_identity, + matches, + tokens_match: result.tokens_match, + baseline_matches: result.baseline_matches, + state_payload_kind: effective_payload_kind(&args), + prompt_token_count: prefix.len(), + decode_token_count: args.decode_tokens, + continuation_token: continuation, + source_tokens, + restored_tokens: result.restored_tokens, + state_bytes: total_bytes as usize, + state_bytes_per_prompt_token: total_bytes as f64 / prefix.len().max(1) as f64, + kv_bytes: (total_bytes - recurrent.len() as u64) as usize, + recurrent_bytes: recurrent.len(), + segment_count: page_count, + segment_bytes: chunk_tokens, + payload_digest, + model_load_ms, + tokenize_ms, + source_prefill_ms, + state_export_ms, + transfer_ms, + transfer_gbps: transfer_gbps(total_bytes as usize, transfer_ms), + source_decode_ms, + store_ms: store.is_some().then_some(store_ms), + overlap_wall_ms: Some(overlap_wall_ms), + receiver: timings, + ttft_disaggregated_ms, + ttft_local_ms, + ttft_speedup: ttft_local_ms.map(|local| local / ttft_disaggregated_ms.max(f64::EPSILON)), + }; + emit_report(&report, report_out.as_deref())?; + ensure_matches(matches, args.allow_mismatch)?; + Ok(()) +} + +fn run_receiver(args: RemoteHandoffArgs) -> Result<()> { + let model_identity = runtime_model_identity(&args.runtime)?; + + let model_load_started = Instant::now(); + let model = open_full_model(&args, effective_lane_count(&args))?; + let model_load_ms = elapsed_ms(model_load_started); + + let store = open_store(&args)?; + let local_state_identity = state_identity_for(&args, &model_identity); + let listener = TcpListener::bind(args.listen) + .with_context(|| format!("failed to bind receiver listener on {}", args.listen))?; + eprintln!( + "remote-handoff receiver ready on {} (model loaded in {model_load_ms:.0} ms)", + args.listen + ); + let mut served = 0usize; + let mut failures = 0usize; + loop { + let (stream, sender_addr) = listener.accept().context("failed to accept sender")?; + served += 1; + eprintln!("remote-handoff sender connected from {sender_addr} (handoff {served})"); + let report_out = per_connection_report_path(&args, served); + match handle_receiver_connection( + &model, + &args, + &model_identity, + model_load_ms, + stream, + report_out.as_deref(), + store.as_ref(), + &local_state_identity, + ) { + Ok(true) => eprintln!("remote-handoff handoff {served}: MATCH"), + Ok(false) => { + failures += 1; + eprintln!("remote-handoff handoff {served}: MISMATCH"); + } + Err(error) => { + failures += 1; + eprintln!("remote-handoff handoff {served} failed: {error:#}"); + } + } + if args.accept_count != 0 && served >= args.accept_count { + break; + } + } + if failures > 0 { + ensure_matches(false, args.allow_mismatch) + .with_context(|| format!("{failures} of {served} handoffs failed"))?; + } + Ok(()) +} + +fn per_connection_report_path( + args: &RemoteHandoffArgs, + connection_index: usize, +) -> Option { + let base = args.output.report_out.as_ref()?; + if args.accept_count == 1 { + return Some(base.clone()); + } + let stem = base.file_stem().unwrap_or_default().to_string_lossy(); + let extension = base + .extension() + .map(|extension| format!(".{}", extension.to_string_lossy())) + .unwrap_or_default(); + Some(base.with_file_name(format!("{stem}-{connection_index}{extension}"))) +} + +#[allow(clippy::too_many_arguments)] +fn handle_receiver_connection( + model: &StageModel, + args: &RemoteHandoffArgs, + model_identity: &model_artifact::ModelIdentity, + model_load_ms: f64, + stream: TcpStream, + report_out: Option<&std::path::Path>, + store: Option<&HandoffSegmentStore>, + local_state_identity: &str, +) -> Result { + stream.set_nodelay(true).ok(); + // A per-read socket property: it stays in force for every subsequent + // read on this connection, so a sender that stalls mid-stream errors + // out after this long rather than wedging the accept loop. + stream + .set_read_timeout(Some(Duration::from_secs(args.handshake_timeout_secs))) + .ok(); + let mut reader = BufReader::with_capacity(STREAM_BUFFER_BYTES, stream.try_clone()?); + let mut writer = BufWriter::with_capacity(STREAM_BUFFER_BYTES, stream); + + let (hello, _) = read_frame_expect::(&mut reader, frame_kind::HELLO)?; + if let Err(error) = validate_hello(&hello, args, &model_identity.model_id) + .and_then(|()| validate_state_identity(&hello, local_state_identity)) + { + write_frame( + &mut writer, + frame_kind::HELLO_ACK, + &HelloAckHeader { + ok: false, + reason: Some(error.to_string()), + }, + &[], + )?; + writer.flush().ok(); + return Err(error); + } + write_frame( + &mut writer, + frame_kind::HELLO_ACK, + &HelloAckHeader { + ok: true, + reason: None, + }, + &[], + )?; + writer.flush().context("failed to flush hello ack")?; + + // Phase one: accumulate segments — into the L3 store when configured + // (write-behind), in memory otherwise. Nothing touches a session until + // the commit record validates completeness. + let receive_started = Instant::now(); + let mut payload: Vec = Vec::new(); + let mut received_bytes = 0u64; + let mut segment_refs: Vec = Vec::new(); + let mut store_ms = 0.0f64; + let mut segments_seen = 0usize; + let mut staged: Option = None; + let commit = loop { + let (kind, header, body) = read_frame(&mut reader)?; + match kind { + frame_kind::PAGE => { + let page: PageHeader = + serde_json::from_value(header).context("malformed page header")?; + if digest_hex(&body) != page.blake3 { + bail!("page {} failed digest verification", page.index); + } + let stage = match staged.as_mut() { + Some(stage) => stage, + None => staged.insert(StagedStream::new(model)?), + }; + if page.token_start != stage.imported_tokens { + bail!( + "page {} starts at token {} but {} tokens are staged", + page.index, + page.token_start, + stage.imported_tokens + ); + } + stage.hasher.update(&body); + // Import into the staging session while the sender is still + // prefilling later chunks. Staged state cannot generate: + // decode is gated on the commit record validating below. + let import_started = Instant::now(); + stage + .session + .import_kv_page(&page.kv_desc, &body) + .with_context(|| format!("failed to stage KV page {}", page.index))?; + stage.kv_attach_ms += elapsed_ms(import_started); + stage.imported_tokens += page.token_count; + if let Some(store) = store { + let put_started = Instant::now(); + let digest = store + .put_segment(&body) + .context("receiver page write-behind failed")? + .digest; + store_ms += elapsed_ms(put_started); + segment_refs.push(HandoffSegmentRef { + index: page.index as u32, + offset: received_bytes, + bytes: body.len() as u64, + digest, + codec_identity: Some(SegmentCodecIdentity::native_kv_page( + body.len() as u64 + )), + meta_json: Some(serde_json::to_string(&PageSegmentMeta { + kind: "kv-page".to_string(), + kv_desc: Some(page.kv_desc), + token_start: page.token_start, + token_count: page.token_count, + })?), + }); + } + received_bytes += body.len() as u64; + segments_seen += 1; + } + frame_kind::RECURRENT => { + let recurrent: RecurrentHeader = + serde_json::from_value(header).context("malformed recurrent header")?; + if digest_hex(&body) != recurrent.blake3 { + bail!("recurrent snapshot failed digest verification"); + } + let stage = match staged.as_mut() { + Some(stage) => stage, + None => staged.insert(StagedStream::new(model)?), + }; + stage.hasher.update(&body); + if let Some(store) = store { + let put_started = Instant::now(); + let digest = store + .put_segment(&body) + .context("receiver recurrent write-behind failed")? + .digest; + store_ms += elapsed_ms(put_started); + segment_refs.push(HandoffSegmentRef { + index: segments_seen as u32, + offset: received_bytes, + bytes: body.len() as u64, + digest, + codec_identity: Some(SegmentCodecIdentity::raw(body.len() as u64)), + meta_json: Some(serde_json::to_string(&PageSegmentMeta { + kind: "recurrent".to_string(), + kv_desc: None, + token_start: 0, + token_count: stage.imported_tokens, + })?), + }); + } + received_bytes += body.len() as u64; + segments_seen += 1; + stage.recurrent = body; + } + frame_kind::SEGMENT => { + let segment: SegmentHeader = + serde_json::from_value(header).context("malformed segment header")?; + if segment.index != segments_seen { + bail!( + "segment {} arrived out of order (expected {})", + segment.index, + segments_seen + ); + } + if segment.offset != received_bytes { + bail!( + "segment {} offset {} does not match received byte count {received_bytes}", + segment.index, + segment.offset, + ); + } + match store { + Some(store) => { + let put_started = Instant::now(); + let digest = store + .put_segment(&body) + .context("receiver segment write-behind failed")? + .digest; + store_ms += elapsed_ms(put_started); + if digest != segment.blake3 { + bail!("segment {} failed digest verification", segment.index); + } + segment_refs.push(HandoffSegmentRef { + index: segment.index as u32, + offset: segment.offset, + bytes: body.len() as u64, + digest, + codec_identity: Some(SegmentCodecIdentity::raw(body.len() as u64)), + meta_json: None, + }); + } + None => { + if digest_hex(&body) != segment.blake3 { + bail!("segment {} failed digest verification", segment.index); + } + payload.extend_from_slice(&body); + } + } + received_bytes += body.len() as u64; + segments_seen += 1; + } + frame_kind::COMMIT => { + break serde_json::from_value::(header) + .context("malformed commit header")?; + } + other => bail!("unexpected frame kind {other} while receiving segments"), + } + }; + let transfer_receive_ms = elapsed_ms(receive_started); + + let mut committed_manifest = None; + let outcome = if commit.streaming { + (|| -> Result { + let mut stage = staged + .take() + .context("streaming commit arrived before any staged pages")?; + if stage.imported_tokens != commit.prompt_token_count { + bail!( + "staged pages cover {} tokens but commit records {}", + stage.imported_tokens, + commit.prompt_token_count + ); + } + if segments_seen != commit.segment_count { + bail!( + "commit expected {} stream segments but {segments_seen} arrived", + commit.segment_count + ); + } + if stage.recurrent.len() as u64 != commit.recurrent_bytes { + bail!( + "commit records {} recurrent bytes but {} arrived", + commit.recurrent_bytes, + stage.recurrent.len() + ); + } + let running_digest = format!("{}", stage.hasher.finalize().to_hex()); + if running_digest != commit.payload_blake3 { + bail!("streamed payload failed commit digest verification"); + } + let residual_started = Instant::now(); + if stage.recurrent.is_empty() { + stage + .session + .set_position(commit.prompt_token_count) + .context("failed to finalize staged position")?; + } else { + let recurrent = std::mem::take(&mut stage.recurrent); + stage + .session + .import_recurrent_state_for_token_count(&recurrent, commit.prompt_token_count) + .context("failed to import staged recurrent state")?; + } + if let Some(store) = store { + let commit_started = Instant::now(); + let manifest = manifest_from_commit( + &commit, + std::mem::take(&mut segment_refs), + local_state_identity.to_string(), + effective_payload_kind(args), + Vec::new(), + )?; + store + .commit(&manifest) + .context("receiver streaming manifest commit failed")?; + store_ms += elapsed_ms(commit_started); + committed_manifest = Some(manifest); + } + let mut timings = RemoteHandoffReceiverTimings { + kv_attach_ms: stage.kv_attach_ms, + ..RemoteHandoffReceiverTimings::default() + }; + timings.attach_residual_ms = elapsed_ms(residual_started); + decode_and_baseline(model, args, &commit, stage.session, timings) + })() + } else { + (|| -> Result { + let payload_bytes = match store { + Some(store) => { + let commit_started = Instant::now(); + let manifest = manifest_from_commit( + &commit, + std::mem::take(&mut segment_refs), + local_state_identity.to_string(), + payload_kind_name(args.state_payload_kind), + Vec::new(), + )?; + store + .commit(&manifest) + .context("receiver manifest commit failed")?; + // Import from the store, not the socket buffer: the disk + // backend is the path under test, and `assemble` re-verifies + // every segment plus the whole-payload digest. + let assembled = store + .assemble(&manifest) + .context("receiver manifest assembly failed")?; + store_ms += elapsed_ms(commit_started); + committed_manifest = Some(manifest); + assembled + } + None => { + validate_commit(&commit, segments_seen, &payload)?; + std::mem::take(&mut payload) + } + }; + run_receiver_attach_and_decode(model, args, &commit, &payload_bytes) + })() + }; + let (attach, reason) = match outcome { + Ok(attach) => (Some(attach), None), + Err(error) => (None, Some(format!("{error:#}"))), + }; + + let (verify, _) = read_frame_expect::(&mut reader, frame_kind::VERIFY)?; + if let (Some(store), Some(mut manifest)) = (store, committed_manifest) { + // Second-phase manifest update: record the reference continuation so + // an offline restore can self-verify. Not part of the handoff's + // success criteria. + manifest.expected_tokens = verify.source_tokens.clone(); + if let Err(error) = store.commit(&manifest) { + eprintln!("remote-handoff: failed to record expected tokens in manifest: {error:#}"); + } + } + let (restored_tokens, baseline_tokens, mut timings) = match attach { + Some(attach) => ( + attach.restored_tokens, + attach.baseline_tokens, + attach.timings, + ), + None => ( + Vec::new(), + Vec::new(), + RemoteHandoffReceiverTimings::default(), + ), + }; + timings.model_load_ms = model_load_ms; + timings.transfer_receive_ms = transfer_receive_ms; + timings.store_ms = store_ms; + let tokens_match = !restored_tokens.is_empty() && restored_tokens == verify.source_tokens; + let baseline_matches = commit + .run_baseline + .then(|| !baseline_tokens.is_empty() && baseline_tokens == verify.source_tokens); + let ok = reason.is_none(); + write_frame( + &mut writer, + frame_kind::RESULT, + &ResultHeader { + ok, + reason: reason.clone(), + restored_tokens: restored_tokens.clone(), + baseline_tokens, + tokens_match, + baseline_matches, + timings: timings.clone(), + }, + &[], + )?; + writer.flush().context("failed to flush result frame")?; + + let matches = ok && tokens_match && baseline_matches.unwrap_or(true); + let report = RemoteHandoffReport { + mode: "remote-handoff", + status: status(matches), + role: "recv", + model_identity: model_identity.clone(), + matches, + tokens_match, + baseline_matches, + state_payload_kind: payload_kind_name(args.state_payload_kind), + prompt_token_count: commit.prompt_token_count as usize, + decode_token_count: hello.decode_token_count, + continuation_token: commit.continuation_token, + source_tokens: verify.source_tokens, + restored_tokens, + state_bytes: received_bytes as usize, + state_bytes_per_prompt_token: received_bytes as f64 + / (commit.prompt_token_count as f64).max(1.0), + kv_bytes: commit.kv_bytes as usize, + recurrent_bytes: commit.recurrent_bytes as usize, + segment_count: segments_seen, + segment_bytes: args.segment_bytes, + payload_digest: commit.payload_blake3.clone(), + model_load_ms, + tokenize_ms: 0.0, + source_prefill_ms: 0.0, + state_export_ms: 0.0, + transfer_ms: transfer_receive_ms, + transfer_gbps: transfer_gbps(received_bytes as usize, transfer_receive_ms), + source_decode_ms: 0.0, + store_ms: None, + overlap_wall_ms: None, + receiver: timings, + ttft_disaggregated_ms: 0.0, + ttft_local_ms: None, + ttft_speedup: None, + }; + emit_report(&report, report_out)?; + if let Some(reason) = reason { + bail!("remote handoff receive failed: {reason}"); + } + Ok(matches) +} + +/// Reattach continuation state purely from the local L3 store — no network, +/// no exporter. Restart survival: any manifest the store holds can be +/// imported into a fresh process and decoded, and when the manifest records +/// the exporter's continuation the run self-verifies determinism. +fn run_restore(mut args: RemoteHandoffArgs) -> Result<()> { + let store_dir = args + .store_dir + .clone() + .context("--role restore requires --store-dir")?; + let store = HandoffSegmentStore::open(&store_dir, args.store_budget_bytes)?; + let key = match args.manifest.clone() { + Some(key) => key, + None => store + .list_manifests()? + .into_iter() + .next() + .with_context(|| format!("store at {} holds no manifests", store_dir.display()))?, + }; + let manifest = store.load_manifest(&key)?; + // The persisted manifest is authoritative for the encoded payload shape. + // A restart should not require the operator to repeat --streaming merely + // to reconstruct the identity of state already on disk. + args.streaming = manifest.payload_kind == "kv-page-stream"; + let model_identity = runtime_model_identity(&args.runtime)?; + let local_state_identity = state_identity_for(&args, &model_identity); + if manifest.state_identity != local_state_identity { + bail!( + "manifest {key} was produced under state identity {} but this configuration is {local_state_identity}", + manifest.state_identity + ); + } + + let model_load_started = Instant::now(); + let model = open_full_model(&args, effective_lane_count(&args))?; + let model_load_ms = elapsed_ms(model_load_started); + + let commit = CommitHeader { + segment_count: manifest.segments.len(), + total_bytes: manifest.total_bytes, + payload_blake3: manifest.payload_digest.clone(), + prompt_token_count: manifest.token_count, + continuation_token: manifest.continuation_token, + kv_bytes: manifest.kv_bytes, + recurrent_bytes: manifest.recurrent_bytes, + kv_desc: manifest + .kv_desc_json + .as_deref() + .map(serde_json::from_str) + .transpose() + .context("malformed kv desc in manifest")?, + prefix_tokens: Vec::new(), + run_baseline: false, + streaming: manifest.payload_kind == "kv-page-stream", + page_count: 0, + }; + let (attach, state_bytes, store_ms) = if commit.streaming { + // Page-stream manifests restore page by page, exactly as the + // network path staged them. + let restore_started = Instant::now(); + let mut stage = StagedStream::new(&model)?; + let mut recurrent: Vec = Vec::new(); + let mut restored_bytes = 0usize; + for segment in &manifest.segments { + let meta: PageSegmentMeta = serde_json::from_str( + segment + .meta_json + .as_deref() + .context("page-stream segment is missing metadata")?, + ) + .context("malformed page-stream segment metadata")?; + let bytes = store.read_segment(&segment.digest)?; + restored_bytes += bytes.len(); + match meta.kind.as_str() { + "kv-page" => { + let kv_desc = meta + .kv_desc + .context("kv-page segment is missing its descriptor")?; + if meta.token_start != stage.imported_tokens { + bail!( + "page-stream segment starts at token {} but {} tokens are staged", + meta.token_start, + stage.imported_tokens + ); + } + stage + .session + .import_kv_page(&kv_desc, &bytes) + .context("failed to restore staged KV page")?; + stage.imported_tokens += meta.token_count; + } + "recurrent" => recurrent = bytes, + other => bail!("unknown page-stream segment kind {other}"), + } + } + if stage.imported_tokens != manifest.token_count { + bail!( + "page-stream manifest covers {} tokens but records {}", + stage.imported_tokens, + manifest.token_count + ); + } + if recurrent.is_empty() { + stage + .session + .set_position(manifest.token_count) + .context("failed to finalize restored position")?; + } else { + stage + .session + .import_recurrent_state_for_token_count(&recurrent, manifest.token_count) + .context("failed to restore recurrent state")?; + } + let store_ms = elapsed_ms(restore_started); + let timings = RemoteHandoffReceiverTimings::default(); + ( + decode_and_baseline(&model, &args, &commit, stage.session, timings)?, + restored_bytes, + store_ms, + ) + } else { + let assemble_started = Instant::now(); + let payload = store + .assemble(&manifest) + .context("manifest assembly failed")?; + let store_ms = elapsed_ms(assemble_started); + let state_bytes = payload.len(); + ( + run_receiver_attach_and_decode(&model, &args, &commit, &payload)?, + state_bytes, + store_ms, + ) + }; + + let expected = &manifest.expected_tokens; + let compared = expected.len().min(attach.restored_tokens.len()); + let tokens_match = compared > 0 && attach.restored_tokens[..compared] == expected[..compared]; + let matches = expected.is_empty() || tokens_match; + let mut timings = attach.timings; + timings.model_load_ms = model_load_ms; + timings.store_ms = store_ms; + let report = RemoteHandoffReport { + mode: "remote-handoff", + status: status(matches), + role: "restore", + model_identity, + matches, + tokens_match, + baseline_matches: None, + state_payload_kind: effective_payload_kind(&args), + prompt_token_count: manifest.token_count as usize, + decode_token_count: args.decode_tokens, + continuation_token: manifest.continuation_token, + source_tokens: expected.clone(), + restored_tokens: attach.restored_tokens, + state_bytes, + state_bytes_per_prompt_token: state_bytes as f64 / (manifest.token_count as f64).max(1.0), + kv_bytes: manifest.kv_bytes as usize, + recurrent_bytes: manifest.recurrent_bytes as usize, + segment_count: manifest.segments.len(), + segment_bytes: args.segment_bytes, + payload_digest: manifest.payload_digest.clone(), + model_load_ms, + tokenize_ms: 0.0, + source_prefill_ms: 0.0, + state_export_ms: 0.0, + transfer_ms: 0.0, + transfer_gbps: 0.0, + source_decode_ms: 0.0, + store_ms: Some(store_ms), + overlap_wall_ms: None, + receiver: timings, + ttft_disaggregated_ms: 0.0, + ttft_local_ms: None, + ttft_speedup: None, + }; + emit_report(&report, args.output.report_out.as_deref())?; + ensure_matches(matches, args.allow_mismatch)?; + Ok(()) +} + +struct ReceiverAttachOutcome { + restored_tokens: Vec, + baseline_tokens: Vec, + timings: RemoteHandoffReceiverTimings, +} + +/// A session staged from streamed KV pages. Holds no generation authority: +/// callers only decode after the commit record validates, and dropping the +/// struct discards uncommitted state. +struct StagedStream { + session: StageSession, + imported_tokens: u64, + hasher: blake3::Hasher, + recurrent: Vec, + kv_attach_ms: f64, +} + +impl StagedStream { + fn new(model: &StageModel) -> Result { + Ok(Self { + session: model + .create_session() + .context("failed to create staging session")?, + imported_tokens: 0, + hasher: blake3::Hasher::new(), + recurrent: Vec::new(), + kv_attach_ms: 0.0, + }) + } +} + +fn run_receiver_attach_and_decode( + model: &StageModel, + args: &RemoteHandoffArgs, + commit: &CommitHeader, + payload: &[u8], +) -> Result { + let attach_started = Instant::now(); + let mut session = model + .create_session() + .context("failed to create receiver session")?; + import_state_payload(&mut session, args, commit, payload) + .context("failed to import handoff state")?; + let kv_attach_ms = elapsed_ms(attach_started); + let timings = RemoteHandoffReceiverTimings { + kv_attach_ms, + ..RemoteHandoffReceiverTimings::default() + }; + decode_and_baseline(model, args, commit, session, timings) +} + +fn decode_and_baseline( + model: &StageModel, + args: &RemoteHandoffArgs, + commit: &CommitHeader, + mut session: StageSession, + mut timings: RemoteHandoffReceiverTimings, +) -> Result { + let first_decode_started = Instant::now(); + let first = session + .decode_step(commit.continuation_token) + .context("receiver first decode failed")?; + let first_decode_ms = elapsed_ms(first_decode_started); + let mut restored_tokens = vec![first]; + let decode_started = Instant::now(); + if args.decode_tokens > 1 { + restored_tokens.extend(greedy_decode(&mut session, first, args.decode_tokens - 1)?); + } + timings.first_decode_ms = first_decode_ms; + timings.decode_ms = first_decode_ms + elapsed_ms(decode_started); + drop(session); + + let mut baseline_tokens = Vec::new(); + if commit.run_baseline { + if commit.prefix_tokens.is_empty() { + bail!("baseline requested but commit carried no prefix tokens"); + } + let mut baseline = model + .create_session() + .context("failed to create baseline session")?; + let baseline_prefill_started = Instant::now(); + baseline + .prefill_chunked(&commit.prefix_tokens) + .context("baseline prefill failed")?; + timings.baseline_prefill_ms = elapsed_ms(baseline_prefill_started); + let baseline_first_started = Instant::now(); + let first = baseline + .decode_step(commit.continuation_token) + .context("baseline first decode failed")?; + timings.baseline_first_decode_ms = elapsed_ms(baseline_first_started); + baseline_tokens.push(first); + if args.decode_tokens > 1 { + baseline_tokens.extend(greedy_decode(&mut baseline, first, args.decode_tokens - 1)?); + } + } + + Ok(ReceiverAttachOutcome { + restored_tokens, + baseline_tokens, + timings, + }) +} + +fn validate_hello(hello: &HelloHeader, args: &RemoteHandoffArgs, model_id: &str) -> Result<()> { + if hello.protocol_version != PROTOCOL_VERSION { + bail!( + "protocol version mismatch: sender {} vs receiver {PROTOCOL_VERSION}", + hello.protocol_version + ); + } + if hello.model_id != model_id { + bail!( + "model mismatch: sender serves {} but receiver serves {model_id}", + hello.model_id + ); + } + if hello.layer_end != args.runtime.layer_end { + bail!( + "layer_end mismatch: sender {} vs receiver {}", + hello.layer_end, + args.runtime.layer_end + ); + } + if hello.ctx_size != args.runtime.ctx_size { + bail!( + "ctx_size mismatch: sender {} vs receiver {}", + hello.ctx_size, + args.runtime.ctx_size + ); + } + if hello.state_payload_kind != effective_payload_kind(args) { + bail!( + "state payload kind mismatch: sender {} vs receiver {}", + hello.state_payload_kind, + effective_payload_kind(args) + ); + } + if hello.decode_token_count != args.decode_tokens { + bail!( + "decode token count mismatch: sender {} vs receiver {}", + hello.decode_token_count, + args.decode_tokens + ); + } + if hello.lane_count != effective_lane_count(args) { + bail!( + "lane count mismatch: sender {} vs receiver {}", + hello.lane_count, + effective_lane_count(args) + ); + } + Ok(()) +} + +fn validate_state_identity(hello: &HelloHeader, local_state_identity: &str) -> Result<()> { + if hello.state_identity != local_state_identity { + bail!( + "state identity mismatch: sender {} vs receiver {local_state_identity} — the \ + numerical configurations differ even though the per-field checks passed", + hello.state_identity + ); + } + Ok(()) +} + +fn validate_commit(commit: &CommitHeader, segments_seen: usize, payload: &[u8]) -> Result<()> { + if commit.segment_count != segments_seen { + bail!( + "commit expected {} segments but {} arrived", + commit.segment_count, + segments_seen + ); + } + if commit.total_bytes != payload.len() as u64 { + bail!( + "commit expected {} bytes but {} arrived", + commit.total_bytes, + payload.len() + ); + } + if digest_hex(payload) != commit.payload_blake3 { + bail!("assembled payload failed commit digest verification"); + } + Ok(()) +} + +fn payload_kind_name(kind: StatePayloadKind) -> &'static str { + match kind { + StatePayloadKind::ResidentKv => "resident-kv", + StatePayloadKind::FullState => "full-state", + StatePayloadKind::RecurrentOnly => "recurrent-only", + StatePayloadKind::KvRecurrent => "kv-recurrent", + } +} + +fn transfer_gbps(bytes: usize, elapsed_ms: f64) -> f64 { + if elapsed_ms <= 0.0 { + return 0.0; + } + (bytes as f64 * 8.0) / (elapsed_ms / 1000.0) / 1e9 +} + +fn write_frame( + writer: &mut impl Write, + kind: u8, + header: &impl Serialize, + payload: &[u8], +) -> Result<()> { + let header_bytes = serde_json::to_vec(header).context("failed to encode frame header")?; + if header_bytes.len() as u64 > MAX_HEADER_BYTES { + bail!("frame header of {} bytes exceeds limit", header_bytes.len()); + } + writer.write_all(&[kind])?; + writer.write_all(&(header_bytes.len() as u32).to_le_bytes())?; + writer.write_all(&header_bytes)?; + writer.write_all(&(payload.len() as u64).to_le_bytes())?; + writer.write_all(payload)?; + Ok(()) +} + +fn read_frame(reader: &mut impl Read) -> Result<(u8, serde_json::Value, Vec)> { + let mut kind = [0u8; 1]; + reader + .read_exact(&mut kind) + .context("handoff stream closed while reading frame kind")?; + let mut header_len = [0u8; 4]; + reader.read_exact(&mut header_len)?; + let header_len = u32::from_le_bytes(header_len) as u64; + if header_len > MAX_HEADER_BYTES { + bail!("frame header of {header_len} bytes exceeds limit"); + } + let mut header_bytes = vec![0u8; header_len as usize]; + reader.read_exact(&mut header_bytes)?; + let header = serde_json::from_slice(&header_bytes).context("malformed frame header")?; + let mut payload_len = [0u8; 8]; + reader.read_exact(&mut payload_len)?; + let payload_len = u64::from_le_bytes(payload_len); + if payload_len > MAX_SEGMENT_BYTES { + bail!("frame payload of {payload_len} bytes exceeds limit"); + } + let mut payload = vec![0u8; payload_len as usize]; + reader.read_exact(&mut payload)?; + Ok((kind[0], header, payload)) +} + +fn read_frame_expect( + reader: &mut impl Read, + expected_kind: u8, +) -> Result<(T, Vec)> { + let (kind, header, payload) = read_frame(reader)?; + if kind != expected_kind { + bail!("expected frame kind {expected_kind}, got {kind}"); + } + Ok(( + serde_json::from_value(header).context("malformed frame header for expected kind")?, + payload, + )) +} + +fn digest_hex(bytes: &[u8]) -> String { + blake3::hash(bytes).to_hex().to_string() +} diff --git a/crates/skippy-correctness/src/runner/remote_handoff/identity.rs b/crates/skippy-correctness/src/runner/remote_handoff/identity.rs new file mode 100644 index 0000000000..306625c938 --- /dev/null +++ b/crates/skippy-correctness/src/runner/remote_handoff/identity.rs @@ -0,0 +1,173 @@ +use std::path::Path; + +use model_artifact::ModelIdentity; +use skippy_cache::{ExactStateIdentityParams, exact_state_identity}; + +use crate::cli::{FlashAttentionArg, RemoteHandoffArgs}; + +use super::{effective_lane_count, protocol_load_mode}; + +pub(super) fn effective_payload_kind(args: &RemoteHandoffArgs) -> &'static str { + if args.streaming { + "kv-page-stream" + } else { + super::payload_kind_name(args.state_payload_kind) + } +} + +/// Content digest of the served artifact, memoized per path: two harness +/// processes serving different local GGUFs behind the same display model id +/// must never share a state identity. Directories (layer-package refs) are +/// not hashed here; their identity rides the package manifest via the +/// model-identity fields. +fn artifact_sha256_cached(path: &Path) -> Option { + use sha2::Digest as _; + use std::collections::HashMap; + use std::io::Read as _; + use std::sync::{Mutex, OnceLock}; + + static CACHE: OnceLock>>> = OnceLock::new(); + let cache = CACHE.get_or_init(|| Mutex::new(HashMap::new())); + if let Some(cached) = cache + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get(path) + { + return cached.clone(); + } + let digest = (|| -> Option { + if !path.is_file() { + return None; + } + let mut file = std::fs::File::open(path).ok()?; + let mut hasher = sha2::Sha256::new(); + let mut buffer = [0_u8; 64 * 1024]; + loop { + let read = file.read(&mut buffer).ok()?; + if read == 0 { + break; + } + hasher.update(&buffer[..read]); + } + Some(hex::encode(hasher.finalize())) + })(); + cache + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(path.to_path_buf(), digest.clone()); + digest +} + +/// Numerical identity of the state this harness produces or accepts. +pub(super) fn state_identity_for(args: &RemoteHandoffArgs, identity: &ModelIdentity) -> String { + let source_model_sha256 = artifact_sha256_cached(&args.runtime.model); + exact_state_identity(&ExactStateIdentityParams { + model_id: &identity.model_id, + model_revision: identity.source_revision.as_deref(), + model_file: identity.source_file.as_deref(), + manifest_sha256: None, + source_model_sha256: source_model_sha256.as_deref(), + package_ref: None, + load_mode: protocol_load_mode(args.runtime.stage_load_mode), + cache_type_k: "f16", + cache_type_v: "f16", + flash_attn_type: protocol_flash_attn_type(args.runtime.flash_attn), + n_gpu_layers: args.runtime.n_gpu_layers, + backend_device: None, + layer_start: 0, + layer_end: args.runtime.layer_end, + ctx_size: args.runtime.ctx_size, + lane_count: effective_lane_count(args), + payload_kind: effective_payload_kind(args), + }) +} + +fn protocol_flash_attn_type(value: FlashAttentionArg) -> skippy_protocol::FlashAttentionType { + match value { + FlashAttentionArg::Auto => skippy_protocol::FlashAttentionType::Auto, + FlashAttentionArg::Disabled => skippy_protocol::FlashAttentionType::Disabled, + FlashAttentionArg::Enabled => skippy_protocol::FlashAttentionType::Enabled, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::cli::{OutputArgs, RemoteHandoffRole, RuntimeArgs, StageLoadMode, StatePayloadKind}; + + fn args_for(model: std::path::PathBuf) -> RemoteHandoffArgs { + RemoteHandoffArgs { + runtime: RuntimeArgs { + model, + model_id: None, + stage_model: None, + stage_load_mode: StageLoadMode::RuntimeSlice, + layer_end: 28, + ctx_size: 2048, + n_gpu_layers: 99, + n_batch: None, + n_ubatch: None, + prompt: "Hello".to_string(), + flash_attn: FlashAttentionArg::Auto, + }, + output: OutputArgs { report_out: None }, + role: RemoteHandoffRole::Send, + listen: "127.0.0.1:19081".parse().expect("addr"), + peer: None, + state_payload_kind: StatePayloadKind::FullState, + prefix_token_count: None, + decode_tokens: 16, + segment_bytes: 8 * 1024 * 1024, + baseline: false, + runtime_lane_count: None, + handshake_timeout_secs: 600, + accept_count: 1, + store_dir: None, + store_budget_bytes: 0, + manifest: None, + streaming: false, + stream_chunk_tokens: 512, + allow_mismatch: false, + } + } + + #[test] + fn different_file_contents_behind_one_model_id_change_identity() { + let dir = std::env::temp_dir() + .join("skippy-remote-handoff-identity-tests") + .join(std::process::id().to_string()); + std::fs::create_dir_all(&dir).expect("temp dir"); + let first = dir.join("model-a.gguf"); + let second = dir.join("model-b.gguf"); + std::fs::write(&first, b"weights generation one").expect("write first"); + std::fs::write(&second, b"weights generation two").expect("write second"); + let identity = ModelIdentity::from_model_id("org/model:Q4_K_M"); + + let first_identity = state_identity_for(&args_for(first.clone()), &identity); + let second_identity = state_identity_for(&args_for(second), &identity); + assert_ne!(first_identity, second_identity); + assert_eq!( + first_identity, + state_identity_for(&args_for(first), &identity) + ); + } + + #[test] + fn load_mode_changes_handoff_identity() { + let dir = std::env::temp_dir() + .join("skippy-remote-handoff-load-mode-tests") + .join(std::process::id().to_string()); + std::fs::create_dir_all(&dir).expect("temp dir"); + let model = dir.join("model.gguf"); + std::fs::write(&model, b"weights").expect("write model"); + let identity = ModelIdentity::from_model_id("org/model:Q4_K_M"); + let runtime = args_for(model.clone()); + let mut artifact = args_for(model); + artifact.runtime.stage_load_mode = StageLoadMode::ArtifactSlice; + + assert_ne!( + state_identity_for(&runtime, &identity), + state_identity_for(&artifact, &identity) + ); + } +} diff --git a/crates/skippy-correctness/src/runner/stage_execution.rs b/crates/skippy-correctness/src/runner/stage_execution.rs index b7e0ac0632..ca14ec503f 100644 --- a/crates/skippy-correctness/src/runner/stage_execution.rs +++ b/crates/skippy-correctness/src/runner/stage_execution.rs @@ -100,6 +100,13 @@ pub(in crate::runner) struct BinaryStateHandoffConfig { pub(in crate::runner) skip_suffix_prefill_check: bool, pub(in crate::runner) synthetic_input_activation: bool, pub(in crate::runner) binary_control: bool, + pub(in crate::runner) cachegen_gate: bool, + pub(in crate::runner) cache_type_k: u32, + pub(in crate::runner) cache_type_v: u32, + pub(in crate::runner) cachegen_continuation_steps: usize, + pub(in crate::runner) cachegen_min_token_agreement: f64, + pub(in crate::runner) cachegen_max_p99_decode_regression: f64, + pub(in crate::runner) cachegen_max_peak_working_bytes: Option, pub(in crate::runner) child_logs: bool, pub(in crate::runner) startup_timeout_secs: u64, pub(in crate::runner) max_inflight: usize, diff --git a/crates/skippy-correctness/src/runner/state_handoff.rs b/crates/skippy-correctness/src/runner/state_handoff.rs index 8b13af973a..8355ff5283 100644 --- a/crates/skippy-correctness/src/runner/state_handoff.rs +++ b/crates/skippy-correctness/src/runner/state_handoff.rs @@ -9,14 +9,13 @@ use skippy_protocol::binary::{ write_stage_message, }; use skippy_runtime::{ - ActivationFrame, GGML_TYPE_F16, MtpSource, RuntimeConfig, RuntimeKvPageDesc, StageModel, - StageSession, + ActivationFrame, MtpSource, RuntimeConfig, RuntimeKvPageDesc, StageModel, StageSession, }; use crate::{ cli::{StageLoadMode, StateHandoffArgs, StatePayloadKind}, report::{ - StageModelReport, StateHandoffReport, StatePayloadBlockDigestReport, + CacheGenGateReport, StageModelReport, StateHandoffReport, StatePayloadBlockDigestReport, StatePayloadDigestReport, }, support::{ @@ -70,11 +69,12 @@ struct BinaryStateHandoffResult { pub(in crate::runner) restored_output_matches: Option, pub(in crate::runner) suffix_prefill_matches: Option, pub(in crate::runner) cache_hit_matches: bool, + pub(in crate::runner) cachegen_gate: Option, pub(in crate::runner) stage_models: Vec, } #[derive(Clone)] -enum LocalStatePayload { +pub(in crate::runner) enum LocalStatePayload { ResidentKv { cache_seq_id: i32, token_count: u64, @@ -208,6 +208,13 @@ pub fn state_handoff(args: StateHandoffArgs) -> Result<()> { skip_suffix_prefill_check: args.skip_suffix_prefill_check, synthetic_input_activation: args.synthetic_input_activation, binary_control: args.binary_control, + cachegen_gate: args.cachegen_gate, + cache_type_k: args.cache_type_k.ggml_type(), + cache_type_v: args.cache_type_v.ggml_type(), + cachegen_continuation_steps: args.cachegen_continuation_steps, + cachegen_min_token_agreement: args.cachegen_min_token_agreement, + cachegen_max_p99_decode_regression: args.cachegen_max_p99_decode_regression, + cachegen_max_peak_working_bytes: args.cachegen_max_peak_working_bytes, child_logs: args.server.child_logs, startup_timeout_secs: args.server.startup_timeout_secs, max_inflight: args.server.max_inflight, @@ -268,6 +275,7 @@ pub fn state_handoff(args: StateHandoffArgs) -> Result<()> { ), cache_hit_import_ms: handoff.cache_hit_import_ms, cache_hit_decode_ms: handoff.cache_hit_decode_ms, + cachegen_gate: handoff.cachegen_gate, stage_models: handoff.stage_models, }; emit_report(&report, report_out.as_deref())?; @@ -287,6 +295,27 @@ fn run_binary_state_handoff(args: BinaryStateHandoffConfig) -> Result Result Result Result, @@ -704,8 +738,8 @@ fn run_local_state_handoff( checkpoint_quantization: skippy_runtime::CheckpointQuantization::Preserve, checkpoint_imatrix: None, checkpoint_imatrix_sha256: None, - cache_type_k: GGML_TYPE_F16, - cache_type_v: GGML_TYPE_F16, + cache_type_k: args.cache_type_k, + cache_type_v: args.cache_type_v, flash_attn_type: runtime_flash_attn(args.flash_attn), kv_offload: None, kv_unified: None, @@ -762,6 +796,21 @@ fn run_local_state_handoff( .context("local state handoff resident KV size measurement failed")?; let source_guard = (args.state_payload_kind == StatePayloadKind::ResidentKv).then_some(source); + let cachegen_gate = if args.cachegen_gate { + Some( + super::cachegen_gate::run_cachegen_gate( + &model, + args, + &state_payload, + &prefix, + continuation, + ) + .context("CacheGen acceptance gate failed to execute")?, + ) + } else { + None + }; + let ( roundtrip_state_payload, restored_predicted_token, @@ -902,12 +951,14 @@ fn run_local_state_handoff( matches: predicted_token_matches && restored_output_matches && suffix_prefill_matches.unwrap_or(true) - && cache_hit_matches, + && cache_hit_matches + && cachegen_gate.as_ref().is_none_or(|gate| gate.passed), predicted_token_matches, roundtrip_state_matches, restored_output_matches: Some(restored_output_matches), suffix_prefill_matches, cache_hit_matches, + cachegen_gate, stage_models, }) } @@ -1096,6 +1147,7 @@ fn run_local_resident_slot_handoff( restored_output_matches: Some(restored_output_matches), suffix_prefill_matches: Some(suffix_prefill_matches), cache_hit_matches, + cachegen_gate: None, stage_models, }) } @@ -1478,8 +1530,8 @@ fn build_state_handoff_inputs( checkpoint_quantization: skippy_runtime::CheckpointQuantization::Preserve, checkpoint_imatrix: None, checkpoint_imatrix_sha256: None, - cache_type_k: GGML_TYPE_F16, - cache_type_v: GGML_TYPE_F16, + cache_type_k: args.cache_type_k, + cache_type_v: args.cache_type_v, flash_attn_type: runtime_flash_attn(args.flash_attn), kv_offload: None, kv_unified: None, @@ -1506,6 +1558,16 @@ fn build_state_handoff_inputs( Ok((Some(prefill_input), Some(decode_input), prefill_width)) } +fn cache_type_name(value: u32) -> Result<&'static str> { + match value { + skippy_runtime::GGML_TYPE_F16 => Ok("f16"), + skippy_runtime::GGML_TYPE_F32 => Ok("f32"), + skippy_runtime::GGML_TYPE_Q8_0 => Ok("q8_0"), + skippy_runtime::GGML_TYPE_Q4_0 => Ok("q4_0"), + _ => bail!("unsupported state-handoff K/V cache type {value}"), + } +} + fn synthetic_activation_frame( args: &BinaryStateHandoffConfig, token_count: u32, diff --git a/crates/skippy-ffi/src/abi.rs b/crates/skippy-ffi/src/abi.rs index df5e0e9b37..a679e50283 100644 --- a/crates/skippy-ffi/src/abi.rs +++ b/crates/skippy-ffi/src/abi.rs @@ -19,6 +19,7 @@ pub const FEATURE_KV_EVENTS: u64 = 1 << 33; pub const FEATURE_DEVICE_EVENTS: u64 = 1 << 34; pub const FEATURE_DIAGNOSTIC_EVENTS: u64 = 1 << 35; pub const FEATURE_UNLOAD_EVENTS: u64 = 1 << 36; +pub const FEATURE_CACHEGEN_KV_PAGE: u64 = 1 << 37; pub const MODEL_TENSOR_SOURCE_V1_ABI_VERSION: u32 = 1; pub type ModelReadTensorF32Callback = Option< diff --git a/crates/skippy-ffi/src/dynamic.rs b/crates/skippy-ffi/src/dynamic.rs index d2b8492596..c05c8c515e 100644 --- a/crates/skippy-ffi/src/dynamic.rs +++ b/crates/skippy-ffi/src/dynamic.rs @@ -6,15 +6,16 @@ use std::{ use crate::{ ABI_VERSION_MAJOR, ABI_VERSION_MINOR, ABI_VERSION_PATCH, AbiVersion, ActivationBoundaryDesc, - ActivationDesc, BackendDevice, Error, GenerationSignalWindow, IterationRequest, KvPageDesc, - LlamaLogCallback, LlamaModelQuantizeParams, Model, ModelInfo, ModelTensorSourceV1, MtmdBitmap, - MtmdContext, MtmdContextParams, MtmdDecoderPos, MtmdHelperBitmapWrapper, MtmdHelperInitOpt, - MtmdHelperVideo, MtmdInputChunkType, MtmdInputChunks, MtmdInputText, NativeMtpDraft, - NativeRuntimeLoadError, NgramCache, Opaque, RuntimeConfig, SamplingConfig, Session, - SkippyDecodeStepSampledMtpFn, SkippyModelAttachMtpDraftModelFn, SkippyRuntimeEventReporterV1, - SlicePlan, StagePlan, StagePlanDescV1, StagePlanProfileDescV1, StagePlanStateDescV1, - StagePlanStringRefV1, StagePlanValueDescV1, StagePlanValueKind, StagePlanner, - StagePlannerConfigV1, Status, TensorInfo, TokenSignal, runtime_abi_supported, + ActivationDesc, BackendDevice, CacheGenRecordV1, Error, GenerationSignalWindow, + IterationRequest, KvPageDesc, LlamaLogCallback, LlamaModelQuantizeParams, Model, ModelInfo, + ModelTensorSourceV1, MtmdBitmap, MtmdContext, MtmdContextParams, MtmdDecoderPos, + MtmdHelperBitmapWrapper, MtmdHelperInitOpt, MtmdHelperVideo, MtmdInputChunkType, + MtmdInputChunks, MtmdInputText, NativeMtpDraft, NativeRuntimeLoadError, NgramCache, Opaque, + RuntimeConfig, SamplingConfig, Session, SkippyDecodeStepSampledMtpFn, + SkippyModelAttachMtpDraftModelFn, SkippyRuntimeEventReporterV1, SlicePlan, StagePlan, + StagePlanDescV1, StagePlanProfileDescV1, StagePlanStateDescV1, StagePlanStringRefV1, + StagePlanValueDescV1, StagePlanValueKind, StagePlanner, StagePlannerConfigV1, Status, + TensorInfo, TokenSignal, runtime_abi_supported, }; static SYMBOLS: OnceLock = OnceLock::new(); @@ -220,6 +221,7 @@ dynamic_symbols! { skippy_import_full_state(session: *mut Session, layer_start: i32, layer_end: i32, input: *const c_void, input_bytes: usize, out_error: *mut *mut Error) -> Status; skippy_export_kv_page(session: *mut Session, layer_start: i32, layer_end: i32, token_start: u64, token_count: u64, out_desc: *mut KvPageDesc, output: *mut c_void, output_capacity: usize, out_bytes: *mut usize, out_error: *mut *mut Error) -> Status; skippy_import_kv_page(session: *mut Session, desc: *const KvPageDesc, input: *const c_void, input_bytes: usize, out_error: *mut *mut Error) -> Status; + skippy_import_cachegen_kv_page_v1(session: *mut Session, desc: *const KvPageDesc, records: *const CacheGenRecordV1, record_count: usize, out_error: *mut *mut Error) -> Status; skippy_export_recurrent_state(session: *mut Session, output: *mut c_void, output_capacity: usize, out_bytes: *mut usize, out_error: *mut *mut Error) -> Status; skippy_import_recurrent_state(session: *mut Session, input: *const c_void, input_bytes: usize, out_error: *mut *mut Error) -> Status; skippy_session_save_prefix(session: *mut Session, cache_seq_id: i32, token_count: u64, out_error: *mut *mut Error) -> Status; diff --git a/crates/skippy-ffi/src/lib.rs b/crates/skippy-ffi/src/lib.rs index acf9680f45..7b5d77497c 100644 --- a/crates/skippy-ffi/src/lib.rs +++ b/crates/skippy-ffi/src/lib.rs @@ -5,7 +5,7 @@ mod dynamic_library; // without compiling the crate to determine native-runtime compatibility. pub const ABI_VERSION_MAJOR: u32 = 0; pub const ABI_VERSION_MINOR: u32 = 1; -pub const ABI_VERSION_PATCH: u32 = 54; +pub const ABI_VERSION_PATCH: u32 = 55; mod abi; mod activation; @@ -27,18 +27,18 @@ pub use abi::{ AbiVersion, ActivationDType, ActivationLayout, BACKEND_DEVICE_CAP_ASYNC, BACKEND_DEVICE_CAP_BUFFER_FROM_HOST_PTR, BACKEND_DEVICE_CAP_EVENTS, BACKEND_DEVICE_CAP_HOST_BUFFER, BackendDevice, BackendDeviceType, Error, - FEATURE_ACTIVATION_BOUNDARY, FEATURE_BACKEND_DEVICES, FEATURE_DEVICE_EVENTS, - FEATURE_DIAGNOSTIC_EVENTS, FEATURE_INKLING_MTP_MM, FEATURE_ITERATION_BATCH, FEATURE_KV_EVENTS, - FEATURE_MODEL_LOAD_EVENTS_V2, FEATURE_MODEL_SOURCE, FEATURE_NATIVE_MTP_N1, - FEATURE_NGRAM_CACHE_DRAFT, FEATURE_RUNTIME_EVENT_REPORTER, FEATURE_RUNTIME_EVENTS, - FEATURE_STAGE_PLAN, FEATURE_UNLOAD_EVENTS, IterationRequest, LlamaLogCallback, LoadMode, - MODEL_TENSOR_SOURCE_V1_ABI_VERSION, Model, ModelImatrixEntryV1, ModelInfo, - ModelReadTensorF32Callback, ModelTensorSourceV1, MtmdProgressCallback, MtpSource, NgramCache, - Opaque, RuntimeConfig, Session, SkippyDecodeStepSampledMtpFn, SkippyModelAttachMtpDraftModelFn, - SkippyRuntimeEventCallback, SkippyRuntimeEventCategory, SkippyRuntimeEventEmitterKind, - SkippyRuntimeEventFailureCode, SkippyRuntimeEventKind, SkippyRuntimeEventProgressUnit, - SkippyRuntimeEventReporterV1, SkippyRuntimeEventV1, SlicePlan, Status, TRISTATE_AUTO, - TRISTATE_FALSE, TRISTATE_TRUE, TensorRole, runtime_abi_supported, + FEATURE_ACTIVATION_BOUNDARY, FEATURE_BACKEND_DEVICES, FEATURE_CACHEGEN_KV_PAGE, + FEATURE_DEVICE_EVENTS, FEATURE_DIAGNOSTIC_EVENTS, FEATURE_INKLING_MTP_MM, + FEATURE_ITERATION_BATCH, FEATURE_KV_EVENTS, FEATURE_MODEL_LOAD_EVENTS_V2, FEATURE_MODEL_SOURCE, + FEATURE_NATIVE_MTP_N1, FEATURE_NGRAM_CACHE_DRAFT, FEATURE_RUNTIME_EVENT_REPORTER, + FEATURE_RUNTIME_EVENTS, FEATURE_STAGE_PLAN, FEATURE_UNLOAD_EVENTS, IterationRequest, + LlamaLogCallback, LoadMode, MODEL_TENSOR_SOURCE_V1_ABI_VERSION, Model, ModelImatrixEntryV1, + ModelInfo, ModelReadTensorF32Callback, ModelTensorSourceV1, MtmdProgressCallback, MtpSource, + NgramCache, Opaque, RuntimeConfig, Session, SkippyDecodeStepSampledMtpFn, + SkippyModelAttachMtpDraftModelFn, SkippyRuntimeEventCallback, SkippyRuntimeEventCategory, + SkippyRuntimeEventEmitterKind, SkippyRuntimeEventFailureCode, SkippyRuntimeEventKind, + SkippyRuntimeEventProgressUnit, SkippyRuntimeEventReporterV1, SkippyRuntimeEventV1, SlicePlan, + Status, TRISTATE_AUTO, TRISTATE_FALSE, TRISTATE_TRUE, TensorRole, runtime_abi_supported, }; pub use activation::{ ACTIVATION_FLAG_GEMMA3N_ALTUP, ACTIVATION_FLAG_INKLING_MTP_EMBD, ACTIVATION_SIDEBAND_TOKEN_IDS, @@ -71,6 +71,9 @@ pub use stage_plan::{ StagePlannerProfileV1, StagePlannerTensorV1, }; pub use state::{ + CACHEGEN_RECORD_EXACT, CACHEGEN_RECORD_F16, CACHEGEN_RECORD_F16_TRANSPOSED, + CACHEGEN_RECORD_F32, CACHEGEN_RECORD_F32_TRANSPOSED, CACHEGEN_RECORD_Q4_0, + CACHEGEN_RECORD_Q8_0, CACHEGEN_RECORD_V1_ABI_VERSION, CacheGenRecordV1, KV_PAGE_CODEC_ISWA_COMPOSITE_V1, KV_PAGE_CODEC_SINGLE_V1, KV_PAGE_FLAG_HAS_K_IDX, KV_PAGE_FLAG_V_TRANSPOSED, KvPageComponentDesc, KvPageDesc, }; @@ -101,10 +104,10 @@ pub use dynamic::{ skippy_decode_step_frame_sampled_mtp, skippy_decode_step_sampled, skippy_decode_step_sampled_mtp, skippy_decode_step_sampled_mtp_fn, skippy_detokenize, skippy_error_free, skippy_export_full_state, skippy_export_kv_page, - skippy_export_recurrent_state, skippy_export_state, skippy_import_full_state, - skippy_import_kv_page, skippy_import_recurrent_state, skippy_import_state, - skippy_iteration_batch_sampled, skippy_model_attach_mtp_draft_model_fn, skippy_model_free, - skippy_model_info_free, skippy_model_info_open, skippy_model_info_tensor_at, + skippy_export_recurrent_state, skippy_export_state, skippy_import_cachegen_kv_page_v1, + skippy_import_full_state, skippy_import_kv_page, skippy_import_recurrent_state, + skippy_import_state, skippy_iteration_batch_sampled, skippy_model_attach_mtp_draft_model_fn, + skippy_model_free, skippy_model_info_free, skippy_model_info_open, skippy_model_info_tensor_at, skippy_model_info_tensor_count, skippy_model_input_activation_boundary, skippy_model_llama_model, skippy_model_open, skippy_model_open_from_parts, skippy_model_open_from_parts_with_events_fn, skippy_model_open_from_source, @@ -153,31 +156,31 @@ pub use static_bindings::{ skippy_decode_step_frame_sampled_mtp, skippy_decode_step_sampled, skippy_decode_step_sampled_mtp, skippy_detokenize, skippy_error_free, skippy_export_full_state, skippy_export_kv_page, skippy_export_recurrent_state, skippy_export_state, - skippy_import_full_state, skippy_import_kv_page, skippy_import_recurrent_state, - skippy_import_state, skippy_iteration_batch_sampled, skippy_model_attach_mtp_draft_model, - skippy_model_free, skippy_model_info_free, skippy_model_info_open, skippy_model_info_tensor_at, - skippy_model_info_tensor_count, skippy_model_input_activation_boundary, - skippy_model_llama_model, skippy_model_open, skippy_model_open_from_parts, - skippy_model_open_from_source, skippy_model_output_activation_boundary, - skippy_ngram_cache_append, skippy_ngram_cache_create, skippy_ngram_cache_draft, - skippy_ngram_cache_free, skippy_ngram_cache_reset, skippy_parse_chat_response_json, - skippy_prefill_chunk, skippy_prefill_chunk_frame, skippy_prefill_chunk_frame_sampled, - skippy_prefill_chunk_frame_sampled_with_positions, skippy_prefill_chunk_frame_with_positions, - skippy_retire_verify_checkpoint, skippy_session_batch_size, - skippy_session_begin_external_decode, skippy_session_configure_chat_sampling, - skippy_session_copy_output_activation_frame, skippy_session_create, - skippy_session_create_from_resident_prefix, skippy_session_drop_sequence, - skippy_session_end_external_decode, skippy_session_free, skippy_session_last_token_signal, - skippy_session_llama_context, skippy_session_memory_used_cells, skippy_session_position, - skippy_session_reset, skippy_session_restore_prefix, skippy_session_sample_current, - skippy_session_save_prefix, skippy_session_sequence_id, skippy_session_set_position, - skippy_session_signal_window, skippy_slice_plan_add_layer_range, skippy_slice_plan_create, - skippy_slice_plan_free, skippy_stage_plan_describe_v1, skippy_stage_plan_free, - skippy_stage_plan_profile_at_v1, skippy_stage_plan_resident_tensor_at_v1, - skippy_stage_plan_state_at_v1, skippy_stage_plan_string_v1, - skippy_stage_plan_validate_chain_v1, skippy_stage_plan_value_at_v1, - skippy_stage_planner_create_v1, skippy_stage_planner_free, skippy_stage_planner_realize_v1, - skippy_token_is_eog, skippy_tokenize, skippy_trim_session, skippy_verify_tokens, - skippy_verify_tokens_frame_sampled, skippy_write_gguf_from_parts, + skippy_import_cachegen_kv_page_v1, skippy_import_full_state, skippy_import_kv_page, + skippy_import_recurrent_state, skippy_import_state, skippy_iteration_batch_sampled, + skippy_model_attach_mtp_draft_model, skippy_model_free, skippy_model_info_free, + skippy_model_info_open, skippy_model_info_tensor_at, skippy_model_info_tensor_count, + skippy_model_input_activation_boundary, skippy_model_llama_model, skippy_model_open, + skippy_model_open_from_parts, skippy_model_open_from_source, + skippy_model_output_activation_boundary, skippy_ngram_cache_append, skippy_ngram_cache_create, + skippy_ngram_cache_draft, skippy_ngram_cache_free, skippy_ngram_cache_reset, + skippy_parse_chat_response_json, skippy_prefill_chunk, skippy_prefill_chunk_frame, + skippy_prefill_chunk_frame_sampled, skippy_prefill_chunk_frame_sampled_with_positions, + skippy_prefill_chunk_frame_with_positions, skippy_retire_verify_checkpoint, + skippy_session_batch_size, skippy_session_begin_external_decode, + skippy_session_configure_chat_sampling, skippy_session_copy_output_activation_frame, + skippy_session_create, skippy_session_create_from_resident_prefix, + skippy_session_drop_sequence, skippy_session_end_external_decode, skippy_session_free, + skippy_session_last_token_signal, skippy_session_llama_context, + skippy_session_memory_used_cells, skippy_session_position, skippy_session_reset, + skippy_session_restore_prefix, skippy_session_sample_current, skippy_session_save_prefix, + skippy_session_sequence_id, skippy_session_set_position, skippy_session_signal_window, + skippy_slice_plan_add_layer_range, skippy_slice_plan_create, skippy_slice_plan_free, + skippy_stage_plan_describe_v1, skippy_stage_plan_free, skippy_stage_plan_profile_at_v1, + skippy_stage_plan_resident_tensor_at_v1, skippy_stage_plan_state_at_v1, + skippy_stage_plan_string_v1, skippy_stage_plan_validate_chain_v1, + skippy_stage_plan_value_at_v1, skippy_stage_planner_create_v1, skippy_stage_planner_free, + skippy_stage_planner_realize_v1, skippy_token_is_eog, skippy_tokenize, skippy_trim_session, + skippy_verify_tokens, skippy_verify_tokens_frame_sampled, skippy_write_gguf_from_parts, skippy_write_gguf_metadata_from_parts, skippy_write_slice_gguf, }; diff --git a/crates/skippy-ffi/src/state.rs b/crates/skippy-ffi/src/state.rs index d738c0e4bd..232cbf51d9 100644 --- a/crates/skippy-ffi/src/state.rs +++ b/crates/skippy-ffi/src/state.rs @@ -1,3 +1,5 @@ +use std::ffi::c_void; + #[repr(C)] #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct KvPageComponentDesc { @@ -23,6 +25,31 @@ pub const KV_PAGE_CODEC_ISWA_COMPOSITE_V1: u32 = 2; pub const KV_PAGE_FLAG_V_TRANSPOSED: u64 = 1 << 0; pub const KV_PAGE_FLAG_HAS_K_IDX: u64 = 1 << 1; +pub const CACHEGEN_RECORD_V1_ABI_VERSION: u32 = 1; +pub const CACHEGEN_RECORD_F16: u32 = 0; +pub const CACHEGEN_RECORD_EXACT: u32 = 1; +pub const CACHEGEN_RECORD_F16_TRANSPOSED: u32 = 2; +pub const CACHEGEN_RECORD_F32: u32 = 3; +pub const CACHEGEN_RECORD_F32_TRANSPOSED: u32 = 4; +pub const CACHEGEN_RECORD_Q8_0: u32 = 5; +pub const CACHEGEN_RECORD_Q4_0: u32 = 6; + +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct CacheGenRecordV1 { + pub abi_version: u32, + pub kind: u32, + pub element_bytes: u32, + pub reserved0: u32, + pub output_offset: u64, + pub decoded_bytes: u64, + pub token_count: u64, + pub token_start: u64, + pub total_tokens: u64, + pub payload: *const c_void, + pub payload_bytes: usize, +} + #[repr(C)] #[derive(Debug, Clone, Copy, Default)] pub struct KvPageDesc { diff --git a/crates/skippy-ffi/src/static_bindings.rs b/crates/skippy-ffi/src/static_bindings.rs index a40b48ac0c..34e2873fe1 100644 --- a/crates/skippy-ffi/src/static_bindings.rs +++ b/crates/skippy-ffi/src/static_bindings.rs @@ -470,6 +470,14 @@ unsafe extern "C" { out_error: *mut *mut Error, ) -> Status; + pub fn skippy_import_cachegen_kv_page_v1( + session: *mut Session, + desc: *const KvPageDesc, + records: *const crate::CacheGenRecordV1, + record_count: usize, + out_error: *mut *mut Error, + ) -> Status; + pub fn skippy_export_recurrent_state( session: *mut Session, output: *mut c_void, diff --git a/crates/skippy-ffi/src/tests.rs b/crates/skippy-ffi/src/tests.rs index d9e87260f6..1b107ac98d 100644 --- a/crates/skippy-ffi/src/tests.rs +++ b/crates/skippy-ffi/src/tests.rs @@ -2,9 +2,12 @@ use std::mem::{offset_of, size_of}; use crate::{ ABI_VERSION_MAJOR, ABI_VERSION_MINOR, ABI_VERSION_PATCH, AbiVersion, ActivationBoundaryDesc, - StagePlanDescV1, StagePlanProfileDescV1, StagePlanStateDescV1, StagePlanStateKind, - StagePlanStringRefV1, StagePlanValueDescV1, StagePlannerConfigV1, StagePlannerProfileV1, - StagePlannerTensorV1, runtime_abi_supported, + CACHEGEN_RECORD_EXACT, CACHEGEN_RECORD_F16, CACHEGEN_RECORD_F16_TRANSPOSED, + CACHEGEN_RECORD_F32, CACHEGEN_RECORD_F32_TRANSPOSED, CACHEGEN_RECORD_Q4_0, + CACHEGEN_RECORD_Q8_0, CACHEGEN_RECORD_V1_ABI_VERSION, CacheGenRecordV1, StagePlanDescV1, + StagePlanProfileDescV1, StagePlanStateDescV1, StagePlanStateKind, StagePlanStringRefV1, + StagePlanValueDescV1, StagePlannerConfigV1, StagePlannerProfileV1, StagePlannerTensorV1, + runtime_abi_supported, }; #[cfg(target_pointer_width = "64")] @@ -76,6 +79,27 @@ fn activation_boundary_descriptor_matches_native_layout() { assert_eq!(offset_of!(ActivationBoundaryDesc, required_sidebands), 40); } +#[test] +#[cfg(target_pointer_width = "64")] +fn cachegen_record_matches_native_layout() { + assert_eq!(CACHEGEN_RECORD_V1_ABI_VERSION, 1); + assert_eq!(CACHEGEN_RECORD_F16, 0); + assert_eq!(CACHEGEN_RECORD_EXACT, 1); + assert_eq!(CACHEGEN_RECORD_F16_TRANSPOSED, 2); + assert_eq!(CACHEGEN_RECORD_F32, 3); + assert_eq!(CACHEGEN_RECORD_F32_TRANSPOSED, 4); + assert_eq!(CACHEGEN_RECORD_Q8_0, 5); + assert_eq!(CACHEGEN_RECORD_Q4_0, 6); + assert_eq!(size_of::(), 72); + assert_eq!(offset_of!(CacheGenRecordV1, output_offset), 16); + assert_eq!(offset_of!(CacheGenRecordV1, decoded_bytes), 24); + assert_eq!(offset_of!(CacheGenRecordV1, token_count), 32); + assert_eq!(offset_of!(CacheGenRecordV1, token_start), 40); + assert_eq!(offset_of!(CacheGenRecordV1, total_tokens), 48); + assert_eq!(offset_of!(CacheGenRecordV1, payload), 56); + assert_eq!(offset_of!(CacheGenRecordV1, payload_bytes), 64); +} + #[test] #[cfg(target_pointer_width = "64")] fn stage_plan_types_match_native_layout() { diff --git a/crates/skippy-protocol/src/binary/activation_codec.rs b/crates/skippy-protocol/src/binary/activation_codec.rs index 431d41fa5b..59864f329c 100644 --- a/crates/skippy-protocol/src/binary/activation_codec.rs +++ b/crates/skippy-protocol/src/binary/activation_codec.rs @@ -382,72 +382,86 @@ fn f32_to_bf16_bits(value: f32) -> u16 { } fn f32_to_f16_bits(value: f32) -> u16 { - let bits = value.to_bits(); - let sign = ((bits >> 16) & 0x8000) as u16; - let exponent = ((bits >> 23) & 0xff) as i32; - let mantissa = bits & 0x7f_ffff; - - if exponent == 0 { - return sign; - } - if exponent == 0xff { - return sign | 0x7c00 | u16::from(mantissa != 0); - } + f16_bits::f32_to_f16_bits(value) +} - let half_exponent = exponent - 127 + 15; - if half_exponent >= 31 { - return sign | 0x7c00; - } - if half_exponent <= 0 { - if half_exponent < -10 { +/// Shared bit-exact IEEE 754 binary16 conversions. Lived inside this module +/// historically; exposed as a `pub(crate)` submodule so the CacheGen +/// reference in `skippy-cache` reuses the exact conversion instead of +/// duplicating it (a second RNE implementation is a future bit-mismatch). +pub(crate) mod f16_bits { + pub fn f32_to_f16_bits(value: f32) -> u16 { + let bits = value.to_bits(); + let sign = ((bits >> 16) & 0x8000) as u16; + let exponent = ((bits >> 23) & 0xff) as i32; + let mantissa = bits & 0x7f_ffff; + + if exponent == 0 { return sign; } - let mantissa = mantissa | 0x80_0000; - let shift = 14 - half_exponent; - let mut half_mantissa = mantissa >> shift; - let remainder = mantissa & ((1_u32 << shift) - 1); - let halfway = 1_u32 << (shift - 1); - if remainder > halfway || (remainder == halfway && (half_mantissa & 1) != 0) { - half_mantissa += 1; + if exponent == 0xff { + return sign | 0x7c00 | u16::from(mantissa != 0); } - return sign | half_mantissa as u16; - } - let mut half_mantissa = mantissa >> 13; - let remainder = mantissa & 0x1fff; - if remainder > 0x1000 || (remainder == 0x1000 && (half_mantissa & 1) != 0) { - half_mantissa += 1; - if half_mantissa == 0x400 { - let rounded_exponent = half_exponent + 1; - if rounded_exponent >= 31 { - return sign | 0x7c00; + let half_exponent = exponent - 127 + 15; + if half_exponent >= 31 { + return sign | 0x7c00; + } + if half_exponent <= 0 { + if half_exponent < -10 { + return sign; + } + let mantissa = mantissa | 0x80_0000; + let shift = 14 - half_exponent; + let mut half_mantissa = mantissa >> shift; + let remainder = mantissa & ((1_u32 << shift) - 1); + let halfway = 1_u32 << (shift - 1); + if remainder > halfway || (remainder == halfway && (half_mantissa & 1) != 0) { + half_mantissa += 1; } - return sign | ((rounded_exponent as u16) << 10); + return sign | half_mantissa as u16; } + + let mut half_mantissa = mantissa >> 13; + let remainder = mantissa & 0x1fff; + if remainder > 0x1000 || (remainder == 0x1000 && (half_mantissa & 1) != 0) { + half_mantissa += 1; + if half_mantissa == 0x400 { + let rounded_exponent = half_exponent + 1; + if rounded_exponent >= 31 { + return sign | 0x7c00; + } + return sign | ((rounded_exponent as u16) << 10); + } + } + sign | ((half_exponent as u16) << 10) | half_mantissa as u16 + } + + pub fn f16_bits_to_f32(bits: u16) -> f32 { + let sign = (u32::from(bits & 0x8000)) << 16; + let exponent = (bits >> 10) & 0x1f; + let mantissa = u32::from(bits & 0x03ff); + let f32_bits = match exponent { + 0 if mantissa == 0 => sign, + 0 => { + let mut mantissa = mantissa; + let mut exponent = -14_i32; + while (mantissa & 0x0400) == 0 { + mantissa <<= 1; + exponent -= 1; + } + mantissa &= 0x03ff; + sign | (((exponent + 127) as u32) << 23) | (mantissa << 13) + } + 0x1f => sign | 0x7f80_0000 | (mantissa << 13), + _ => sign | ((u32::from(exponent) + 112) << 23) | (mantissa << 13), + }; + f32::from_bits(f32_bits) } - sign | ((half_exponent as u16) << 10) | half_mantissa as u16 } fn f16_bits_to_f32(bits: u16) -> f32 { - let sign = (u32::from(bits & 0x8000)) << 16; - let exponent = (bits >> 10) & 0x1f; - let mantissa = u32::from(bits & 0x03ff); - let f32_bits = match exponent { - 0 if mantissa == 0 => sign, - 0 => { - let mut mantissa = mantissa; - let mut exponent = -14_i32; - while (mantissa & 0x0400) == 0 { - mantissa <<= 1; - exponent -= 1; - } - mantissa &= 0x03ff; - sign | (((exponent + 127) as u32) << 23) | (mantissa << 13) - } - 0x1f => sign | 0x7f80_0000 | (mantissa << 13), - _ => sign | ((u32::from(exponent) + 112) << 23) | (mantissa << 13), - }; - f32::from_bits(f32_bits) + f16_bits::f16_bits_to_f32(bits) } #[cfg(test)] diff --git a/crates/skippy-protocol/src/binary/mod.rs b/crates/skippy-protocol/src/binary/mod.rs index a742055a0d..8185279641 100644 --- a/crates/skippy-protocol/src/binary/mod.rs +++ b/crates/skippy-protocol/src/binary/mod.rs @@ -10,6 +10,7 @@ pub use activation::{ encode_f32_activation_payload_with_state_flags, select_lossless_activation_codec_with_state_flags, }; +pub use activation_codec::f16_bits::{f16_bits_to_f32, f32_to_f16_bits}; pub use codec::{ read_stage_message, read_stage_message_for_codec, read_stage_message_for_codec_policy, recv_ready, recv_reply, send_ready, send_reply_ack, send_reply_ack_with_stats, diff --git a/crates/skippy-protocol/src/lib.rs b/crates/skippy-protocol/src/lib.rs index 1cf0df9f74..6d48d08bb1 100644 --- a/crates/skippy-protocol/src/lib.rs +++ b/crates/skippy-protocol/src/lib.rs @@ -35,8 +35,8 @@ pub use messages::{ StateImportMessage, StopMessage, TokenReplyMessage, }; pub use validation::{ - MAX_STAGE_FRAME_BYTES, MAX_VERIFY_WINDOW_PIPELINE_DEPTH, SCHEMA_VERSION, STAGE_ALPN_V2, - STAGE_PROTOCOL_GENERATION, STAGE_STREAM_ARTIFACT_TRANSFER, STAGE_STREAM_CONTROL, + KV_ALPN_V1, MAX_STAGE_FRAME_BYTES, MAX_VERIFY_WINDOW_PIPELINE_DEPTH, SCHEMA_VERSION, + STAGE_ALPN_V2, STAGE_PROTOCOL_GENERATION, STAGE_STREAM_ARTIFACT_TRANSFER, STAGE_STREAM_CONTROL, STAGE_STREAM_TRANSPORT, STAGE_SUBPROTOCOL_FEATURE_ARTIFACT_TRANSFER, STAGE_SUBPROTOCOL_FEATURE_LOCAL_GGUF_CONTENT_ID_V1, STAGE_SUBPROTOCOL_FEATURE_STAGE_CONTROL, STAGE_SUBPROTOCOL_FEATURE_STAGE_GENERATION, diff --git a/crates/skippy-protocol/src/validation.rs b/crates/skippy-protocol/src/validation.rs index 5b745fb310..93f00653f6 100644 --- a/crates/skippy-protocol/src/validation.rs +++ b/crates/skippy-protocol/src/validation.rs @@ -3,6 +3,9 @@ use crate::proto; pub const SCHEMA_VERSION: u32 = 1; pub const STAGE_ALPN_V2: &[u8] = b"skippy-stage/2"; +/// ALPN for the L3 KV segment fetch subprotocol: pull exact-state manifests +/// and content-addressed segments from a peer's store (`skippy_cache::l3`). +pub const KV_ALPN_V1: &[u8] = b"skippy-kv/1"; pub const STAGE_SUBPROTOCOL_NAME: &str = "skippy-stage"; pub const STAGE_SUBPROTOCOL_MAJOR: u32 = 2; pub const STAGE_SUBPROTOCOL_FEATURE_STAGE_CONTROL: &str = "stage-control"; diff --git a/crates/skippy-runtime/Cargo.toml b/crates/skippy-runtime/Cargo.toml index 792bd94e78..f7bc3d33b9 100644 --- a/crates/skippy-runtime/Cargo.toml +++ b/crates/skippy-runtime/Cargo.toml @@ -16,6 +16,7 @@ crossbeam-queue.workspace = true anyhow.workspace = true hex = "0.4" skippy-ffi = { path = "../skippy-ffi", version = "0.76.1", default-features = false } +skippy-cache = { path = "../skippy-cache", version = "0.76.1" } skippy-model = { path = "../skippy-model", version = "0.76.1" } model-ref = { path = "../model-ref", version = "0.76.1" } serde.workspace = true diff --git a/crates/skippy-runtime/src/kv_pages.rs b/crates/skippy-runtime/src/kv_pages.rs index 3034e517f9..e8a0f90e74 100644 --- a/crates/skippy-runtime/src/kv_pages.rs +++ b/crates/skippy-runtime/src/kv_pages.rs @@ -1,10 +1,41 @@ use std::ptr; use anyhow::{Result, ensure}; -use skippy_ffi::KvPageDesc as RawKvPageDesc; +use skippy_cache::cachegen::archive::{RecordKind, ValidatedArchive, validate_archive}; +use skippy_ffi::{CacheGenRecordV1, KvPageDesc as RawKvPageDesc}; use crate::error::{ensure_ok, free_error}; use crate::session::StageSession; + +fn cachegen_records(validated: &ValidatedArchive<'_>) -> Result> { + validated + .records + .iter() + .map(|record| { + Ok(CacheGenRecordV1 { + abi_version: skippy_ffi::CACHEGEN_RECORD_V1_ABI_VERSION, + kind: match record.kind { + RecordKind::CacheGen => skippy_ffi::CACHEGEN_RECORD_F16, + RecordKind::Exact => skippy_ffi::CACHEGEN_RECORD_EXACT, + RecordKind::CacheGenTransposed => skippy_ffi::CACHEGEN_RECORD_F16_TRANSPOSED, + RecordKind::CacheGenF32 => skippy_ffi::CACHEGEN_RECORD_F32, + RecordKind::CacheGenF32Transposed => skippy_ffi::CACHEGEN_RECORD_F32_TRANSPOSED, + RecordKind::CacheGenQ8_0 => skippy_ffi::CACHEGEN_RECORD_Q8_0, + RecordKind::CacheGenQ4_0 => skippy_ffi::CACHEGEN_RECORD_Q4_0, + }, + element_bytes: record.element_bytes as u32, + reserved0: 0, + output_offset: record.output_offset as u64, + decoded_bytes: record.decoded_len as u64, + token_count: record.token_count as u64, + token_start: record.token_start as u64, + total_tokens: record.total_tokens as u64, + payload: record.payload.as_ptr().cast(), + payload_bytes: record.payload.len(), + }) + }) + .collect() +} use crate::{RuntimeKvPage, RuntimeKvPageDesc, Status}; impl StageSession { @@ -259,6 +290,36 @@ impl StageSession { Ok(()) } + pub fn import_cachegen_kv_page( + &mut self, + desc: &RuntimeKvPageDesc, + archive: &[u8], + ) -> Result<()> { + let raw_len = usize::try_from(desc.payload_bytes)?; + desc.validate_payload(raw_len)?; + if desc.codec == skippy_ffi::KV_PAGE_CODEC_ISWA_COMPOSITE_V1 && self.token_count != 0 { + anyhow::bail!("composite ISWA CacheGen page import requires a fresh session"); + } + let validated = validate_archive(archive, raw_len)?; + let records = cachegen_records(&validated)?; + let raw = desc.as_raw(); + let mut error = ptr::null_mut(); + let status = unsafe { + skippy_ffi::skippy_import_cachegen_kv_page_v1( + self.raw, + &raw, + records.as_ptr(), + records.len(), + &mut error, + ) + }; + ensure_ok(status, error)?; + self.token_count = self + .token_count + .max(desc.token_start.saturating_add(desc.token_count)); + Ok(()) + } + pub fn export_recurrent_state(&mut self) -> Result> { let mut bytes = 0usize; let mut error = ptr::null_mut(); @@ -327,7 +388,103 @@ fn validate_imported_full_state_position(expected: u64, actual: u64) -> Result<( #[cfg(test)] mod tests { - use super::validate_imported_full_state_position; + use skippy_cache::cachegen::archive::{Record, RecordKind, ValidatedArchive}; + + use super::{cachegen_records, validate_imported_full_state_position}; + + #[test] + fn maps_validated_cachegen_records_to_the_native_abi() { + let payload = [1_u8, 2, 3, 4]; + let validated = ValidatedArchive { + raw_len: 16, + records: vec![Record { + kind: RecordKind::CacheGenTransposed, + element_bytes: 2, + output_offset: 8, + decoded_len: 8, + token_count: 2, + token_start: 3, + total_tokens: 5, + payload: &payload, + }], + }; + + let records = cachegen_records(&validated).expect("F16 records map to the native ABI"); + assert_eq!(records.len(), 1); + let record = records[0]; + assert_eq!(record.abi_version, 1); + assert_eq!(record.kind, skippy_ffi::CACHEGEN_RECORD_F16_TRANSPOSED); + assert_eq!(record.element_bytes, 2); + assert_eq!(record.reserved0, 0); + assert_eq!(record.output_offset, 8); + assert_eq!(record.decoded_bytes, 8); + assert_eq!(record.token_count, 2); + assert_eq!(record.token_start, 3); + assert_eq!(record.total_tokens, 5); + assert_eq!(record.payload, payload.as_ptr().cast()); + assert_eq!(record.payload_bytes, payload.len()); + } + + #[test] + fn maps_quantized_records_to_the_native_abi() { + let payload = [1_u8, 2, 3, 4]; + let validated = ValidatedArchive { + raw_len: 16, + records: vec![Record { + kind: RecordKind::CacheGenQ8_0, + element_bytes: 34, + output_offset: 0, + decoded_len: 16, + token_count: 1, + token_start: 0, + total_tokens: 0, + payload: &payload, + }], + }; + + let records = cachegen_records(&validated).expect("Q8_0 records map to the native ABI"); + assert_eq!(records[0].kind, skippy_ffi::CACHEGEN_RECORD_Q8_0); + assert_eq!(records[0].element_bytes, 34); + + let validated = ValidatedArchive { + raw_len: 16, + records: vec![Record { + kind: RecordKind::CacheGenQ4_0, + element_bytes: 18, + output_offset: 0, + decoded_len: 16, + token_count: 1, + token_start: 0, + total_tokens: 0, + payload: &payload, + }], + }; + let records = cachegen_records(&validated).expect("Q4_0 records map to the native ABI"); + assert_eq!(records[0].kind, skippy_ffi::CACHEGEN_RECORD_Q4_0); + assert_eq!(records[0].element_bytes, 18); + } + + #[test] + fn maps_f32_records_to_the_native_abi() { + let payload = [1_u8, 2, 3, 4]; + let validated = ValidatedArchive { + raw_len: 16, + records: vec![Record { + kind: RecordKind::CacheGenF32Transposed, + element_bytes: 4, + output_offset: 8, + decoded_len: 16, + token_count: 2, + token_start: 3, + total_tokens: 5, + payload: &payload, + }], + }; + + let records = cachegen_records(&validated).expect("F32 records map to the native ABI"); + assert_eq!(records[0].kind, skippy_ffi::CACHEGEN_RECORD_F32_TRANSPOSED); + assert_eq!(records[0].element_bytes, 4); + } #[test] fn full_state_import_accepts_the_position_carried_by_native_state() { diff --git a/crates/skippy-runtime/src/lib.rs b/crates/skippy-runtime/src/lib.rs index 774dcb0388..cf1151ff2f 100644 --- a/crates/skippy-runtime/src/lib.rs +++ b/crates/skippy-runtime/src/lib.rs @@ -38,11 +38,11 @@ pub use gguf_writer::{ ModelInfo, SlicePlan, write_gguf_from_parts, write_gguf_metadata_from_parts, }; pub use logging::{ - LLAMA_LOG_LEVEL_DEBUG, NativeLogEvent, NativeLogParserMode, NativeLogParserPolicy, - configure_native_log_parser, disable_verbose_native_logs, enable_verbose_native_logs, - redirect_native_logs_to_file, register_filtered_native_logs, restore_native_logs, - set_filtered_native_logs_enabled, suppress_native_logs, unregister_filtered_native_logs, - write_native_log_note, + LLAMA_LOG_LEVEL_DEBUG, MeasuredNativeBuffers, NativeLogEvent, NativeLogParserMode, + NativeLogParserPolicy, configure_native_log_parser, disable_verbose_native_logs, + enable_verbose_native_logs, measured_native_buffers, redirect_native_logs_to_file, + register_filtered_native_logs, restore_native_logs, set_filtered_native_logs_enabled, + suppress_native_logs, unregister_filtered_native_logs, write_native_log_note, }; pub use native::{StageModel, StageModelReader}; pub use native_mtp::NativeMtpDraft; @@ -65,6 +65,9 @@ pub use skippy_ffi::{ ACTIVATION_FLAG_GEMMA3N_ALTUP, ACTIVATION_SIDEBAND_TOKEN_IDS, ActivationDType as RuntimeActivationDType, ActivationLayout as RuntimeActivationLayout, }; +// KV page descriptor flags. Re-exported so callers can read a page's layout +// without taking a direct dependency on the raw ABI crate. +pub use skippy_ffi::{KV_PAGE_FLAG_HAS_K_IDX, KV_PAGE_FLAG_V_TRANSPOSED}; pub use stage_planning::plan_gguf_stage_resident_tensor_names; pub use types::{ ActivationBoundaryDesc, ActivationDesc, ActivationFrame, ChatReasoningFormat, diff --git a/crates/skippy-runtime/src/logging.rs b/crates/skippy-runtime/src/logging.rs index 5d9322d4de..b621f9c9b1 100644 --- a/crates/skippy-runtime/src/logging.rs +++ b/crates/skippy-runtime/src/logging.rs @@ -150,6 +150,30 @@ impl ProgressTracker { } } +/// Latest measured buffer sizes parsed from native log lines, keyed by the +/// line's kind (compute vs KV). These are what llama.cpp actually allocated +/// during `sched_reserve`, and are the ground truth the memory planner should +/// charge instead of the KV-scaled estimate. +#[derive(Clone, Copy, Debug, Default, PartialEq)] +pub struct MeasuredNativeBuffers { + pub compute_mib: Option, + pub kv_mib: Option, + /// A CPU-resident compute or KV allocation was observed, so the device + /// footprint is incomplete for a capacity pool that includes host RAM. + pub host_memory_observed: bool, +} + +/// Snapshot of the measured native buffer sizes observed so far in this +/// process. The native log callback is synchronous with model open, so by the +/// time `skippy_model_open` returns, the `sched_reserve` buffer lines have +/// already been parsed. The snapshot is returned whenever the aggregator is +/// reachable; its fields stay `None` until a buffer line is observed (e.g. +/// native log forwarding disabled). +pub fn measured_native_buffers() -> Option { + let aggregator = native_log_aggregator().lock().ok()?; + Some(aggregator.measured_snapshot()) +} + #[derive(Debug, Default)] struct ModelMetadataHighlights { architecture: Option, @@ -246,6 +270,19 @@ struct NativeLogAggregator { tensor_groups: Vec<(String, usize)>, tensor_groups_emitted: bool, kv_layers_seen: BTreeSet, + /// Latest measured buffer sizes parsed from native log lines, keyed by + /// backend device name (e.g. `CUDA0`, `Metal`). Updated by the + /// memory/kv_cache arms of `summarize_native_log_line`; read via + /// [`measured_native_buffers`] after model open completes. Multiple + /// reserves on the same device keep the high-water mark; distinct devices + /// are summed by [`measured_native_buffers`] (one buffer line is printed + /// per device, so a plain per-kind max would under-measure multi-GPU by a + /// factor of N). Host-pinned buffers (`CUDA_Host`) and CPU buffers are + /// excluded at record time — they are not device memory and must never be + /// charged against a VRAM budget. + measured_compute_mib: BTreeMap, + measured_kv_mib: BTreeMap, + host_memory_observed: bool, } fn native_log_file() -> &'static Mutex>> { @@ -294,6 +331,20 @@ pub fn configure_native_log_parser(policy: NativeLogParserPolicy) { } impl NativeLogAggregator { + /// Per-device-summed measured buffer snapshot (see + /// [`measured_native_buffers`] for the accounting rules). + fn measured_snapshot(&self) -> MeasuredNativeBuffers { + let compute_mib = (!self.measured_compute_mib.is_empty()) + .then(|| self.measured_compute_mib.values().sum::()); + let kv_mib = + (!self.measured_kv_mib.is_empty()).then(|| self.measured_kv_mib.values().sum::()); + MeasuredNativeBuffers { + compute_mib, + kv_mib, + host_memory_observed: self.host_memory_observed, + } + } + fn reset(&mut self) { *self = Self::default(); } @@ -311,6 +362,12 @@ impl NativeLogAggregator { self.tensor_groups.clear(); self.tensor_groups_emitted = false; self.kv_layers_seen.clear(); + // A new model load invalidates the previous model's measured buffer + // sizes: buffer scales are model/shape-specific, and charging one + // model's HWM against another's budget would be wrong both directions. + self.measured_compute_mib.clear(); + self.measured_kv_mib.clear(); + self.host_memory_observed = false; } fn process_line(&mut self, line: &str) -> Vec { @@ -399,6 +456,8 @@ impl NativeLogAggregator { return events; } + self.record_measured_buffer_size(s); + if let Some(event) = summarize_native_log_line(s) { events.push(event); } @@ -406,6 +465,49 @@ impl NativeLogAggregator { events } + /// Track the largest measured buffer size per kind. A later, smaller line + /// (e.g. a per-graph reserve for a shorter context) must not lower the + /// high-water mark recorded at full context init. CPU-offload lines are + /// skipped: the snapshot feeds VRAM planning, and a CPU-resident buffer + /// larger than the accelerator's must not be charged against VRAM. + fn record_measured_buffer_size(&mut self, line: &str) { + if !line.contains("buffer size") { + return; + } + let Some(device) = buffer_size_device(line) else { + return; + }; + let is_compute = line.contains("compute buffer size"); + let is_kv = line.contains("KV buffer size"); + if !is_compute && !is_kv { + return; + } + if device == "CPU" || device.starts_with("CPU_") { + self.host_memory_observed = true; + return; + } + let Some(mib) = parse_buffer_size_mib(line) else { + return; + }; + // Host-pinned staging buffers (CUDA_Host and friends) are host RAM, + // not device memory — never charge them against a VRAM budget. + if is_host_pinned_device_name(&device) { + return; + } + let field = if is_compute { + &mut self.measured_compute_mib + } else { + &mut self.measured_kv_mib + }; + // One line per device per reserve: keep the high-water mark within a + // device (larger of repeated reserves) so a smaller re-reserve on the + // same device cannot shrink the measured footprint. + let slot = field.entry(device).or_insert(0.0); + if mib > *slot { + *slot = mib; + } + } + fn record_layer_assignment(&mut self, layer_index: usize, device: &str) -> Vec { if self.layer_devices.get(&layer_index).map(String::as_str) != Some(device) { self.layer_devices.insert(layer_index, device.to_string()); @@ -523,6 +625,59 @@ fn should_suppress_native_log_line(line: &str) -> bool { && (line.contains(": filtered") || line.contains(": dev ="))) } +fn parse_buffer_size_mib(line: &str) -> Option { + // Native buffer-size lines print the value with a fixed-width field, e.g. + // `sched_reserve: CUDA0 compute buffer size = 579.83 MiB` or + // `llama_kv_cache: CUDA0 KV buffer size = 1088.00 MiB`. Capture the + // last ` MiB` occurrence on the line. + let rest = line.rfind("MiB")?; + let prefix = line[..rest].trim_end(); + let start = prefix + .rfind(|c: char| !(c.is_ascii_digit() || c == '.')) + .map(|idx| idx + 1) + .unwrap_or(0); + prefix[start..].trim().parse::().ok() +} + +/// Host-pinned buffer names some CUDA backends report (e.g. `CUDA_Host`). +/// Their memory is host RAM pinned for device transfers, not device memory — +/// it must not be charged against a VRAM budget. +fn is_host_pinned_device_name(device: &str) -> bool { + device == "CUDA_Host" || device.ends_with("_Host") +} + +/// Backend device a buffer-size line belongs to, from the buffer name token +/// that precedes `KV buffer size` / `compute buffer size` (e.g. `CUDA0`, +/// `CUDA1`, `CUDA_Host`, `Metal`, `CPU`). `None` when the device cannot be +/// determined. +fn buffer_size_device(line: &str) -> Option { + let marker = if line.contains("KV buffer size") { + "KV buffer size" + } else if line.contains("compute buffer size") { + "compute buffer size" + } else { + return None; + }; + let idx = line.find(marker)?; + let name = line[..idx].trim(); + name.rsplit(' ').next().map(str::to_string) +} + +fn buffer_size_params(line: &str) -> Vec<(String, Value)> { + // Structured facts for memory-planning telemetry: the measured buffer size + // (the number llama.cpp actually allocated) plus the device the line names. + // These are the inputs the topology planner will consume in place of its + // KV-scaled compute-buffer estimate. + let mut params = Vec::new(); + if let Some(mib) = parse_buffer_size_mib(line) { + params.push(("buffer_mib".to_string(), Value::from(mib))); + } + if let Some(device) = buffer_size_device(line) { + params.push(("backend_device".to_string(), Value::String(device))); + } + params +} + fn summarize_native_log_line(line: &str) -> Option { if let Some((category, params)) = cpu_offload_diagnostic_params(line) { return Some(NativeLogEvent { @@ -564,6 +719,19 @@ fn summarize_native_log_line(line: &str) -> Option { }); } + if line.starts_with("llama_context: n_ubatch") || line.starts_with("llama_context: flash_attn") + { + // Forward the resolved micro-batch size and flash-attention mode so a live + // deployment can prove which values the runtime actually constructed with. + // These lines come from the llama_context parameter dump + // (llama-context.cpp, `n_ubatch = ...` / `flash_attn = ...`). + return Some(NativeLogEvent { + message: line.to_string(), + category: "runtime", + params: Vec::new(), + }); + } + if line.contains("VRAM") || line.contains("vram") || line.contains("mem_alloc") @@ -572,20 +740,30 @@ fn summarize_native_log_line(line: &str) -> Option { || line.contains("compute buffer size") || line.contains("scratch buffer") { + let params = if line.contains("buffer size") { + buffer_size_params(line) + } else { + Vec::new() + }; return Some(NativeLogEvent { message: line.to_string(), category: "memory", - params: Vec::new(), + params, }); } if line.starts_with("llama_kv_cache:") && (line.contains("buffer size") || line.contains("size = ") || line.contains("attn_rot")) { + let params = if line.contains("buffer size") { + buffer_size_params(line) + } else { + Vec::new() + }; return Some(NativeLogEvent { message: line.to_string(), category: "kv_cache", - params: Vec::new(), + params, }); } @@ -1139,6 +1317,37 @@ mod tests { ); } + #[test] + fn aggregator_forwards_llama_context_config_lines() { + let mut aggregator = NativeLogAggregator::default(); + assert_eq!( + aggregator.process_line("llama_context: n_ubatch = 512"), + vec![NativeLogEvent { + message: "llama_context: n_ubatch = 512".to_string(), + category: "runtime", + params: Vec::new(), + }] + ); + assert_eq!( + aggregator.process_line("llama_context: flash_attn = enabled"), + vec![NativeLogEvent { + message: "llama_context: flash_attn = enabled".to_string(), + category: "runtime", + params: Vec::new(), + }] + ); + assert!( + aggregator + .process_line("llama_context: n_ctx = 8192") + .is_empty() + ); + assert!( + aggregator + .process_line("llama_context: causal_attn = 1") + .is_empty() + ); + } + #[test] fn aggregator_ignores_non_backend_cuda_mentions() { let mut aggregator = NativeLogAggregator::default(); @@ -1246,6 +1455,207 @@ mod tests { })); } + #[test] + fn aggregator_records_measured_buffers_for_snapshot_api() { + let mut aggregator = NativeLogAggregator::default(); + aggregator.process_line("sched_reserve: CUDA0 compute buffer size = 579.83 MiB"); + aggregator.process_line("llama_kv_cache: CUDA0 KV buffer size = 1088.00 MiB"); + assert_eq!( + aggregator.measured_snapshot(), + MeasuredNativeBuffers { + compute_mib: Some(579.83), + kv_mib: Some(1088.00), + host_memory_observed: false, + } + ); + + // A later, smaller reserve for a shorter context must not lower the + // recorded high-water mark for either kind. + aggregator.process_line("sched_reserve: CUDA0 compute buffer size = 512.00 MiB"); + aggregator.process_line("llama_kv_cache: CUDA0 KV buffer size = 1024.00 MiB"); + assert_eq!( + aggregator.measured_snapshot(), + MeasuredNativeBuffers { + compute_mib: Some(579.83), + kv_mib: Some(1088.00), + host_memory_observed: false, + } + ); + + // CPU-offloaded buffers stay excluded from device totals, but make the + // device-only snapshot ineligible for reuse against a mixed pool. + aggregator.process_line("load_tensors: CPU_Mapped model buffer size = 2048.00 MiB"); + aggregator.process_line("llama_kv_cache: CPU KV buffer size = 4096.00 MiB"); + aggregator.process_line("llama_kv_cache: CPU compute buffer size = 8192.00 MiB"); + assert_eq!( + aggregator.measured_snapshot(), + MeasuredNativeBuffers { + compute_mib: Some(579.83), + kv_mib: Some(1088.00), + host_memory_observed: true, + } + ); + + // Model-agnostic summary lines carry no buffer size to record. + aggregator.process_line("VRAM used: 12.4 GB"); + assert_eq!( + aggregator.measured_snapshot(), + MeasuredNativeBuffers { + compute_mib: Some(579.83), + kv_mib: Some(1088.00), + host_memory_observed: true, + } + ); + + // register/unregister reset clears the snapshot for the next model. + aggregator.reset(); + assert_eq!( + aggregator.measured_snapshot(), + MeasuredNativeBuffers { + compute_mib: None, + kv_mib: None, + host_memory_observed: false, + } + ); + } + + #[test] + fn aggregator_sums_measured_buffers_across_devices() { + // One buffer line is printed per backend device: on multi-GPU the + // measured footprint must SUM across devices (a per-kind max would + // under-measure by the device count and the planner would buy + // roughly N x too much context). + let mut aggregator = NativeLogAggregator::default(); + aggregator.process_line("sched_reserve: CUDA0 compute buffer size = 544.00 MiB"); + aggregator.process_line("sched_reserve: CUDA1 compute buffer size = 544.00 MiB"); + aggregator.process_line("llama_kv_cache: CUDA0 KV buffer size = 544.00 MiB"); + aggregator.process_line("llama_kv_cache: CUDA1 KV buffer size = 544.00 MiB"); + assert_eq!( + aggregator.measured_snapshot(), + MeasuredNativeBuffers { + compute_mib: Some(1088.00), + kv_mib: Some(1088.00), + host_memory_observed: false, + } + ); + } + + #[test] + fn aggregator_excludes_host_pinned_buffers() { + // CUDA_Host is host RAM pinned for device transfers, not device + // memory — charging it against a VRAM budget would over-reserve. + let mut aggregator = NativeLogAggregator::default(); + aggregator.process_line("sched_reserve: CUDA0 compute buffer size = 400.00 MiB"); + aggregator.process_line("sched_reserve: CUDA_Host compute buffer size = 128.00 MiB"); + aggregator.process_line("llama_kv_cache: CUDA_Host KV buffer size = 64.00 MiB"); + assert_eq!( + aggregator.measured_snapshot(), + MeasuredNativeBuffers { + compute_mib: Some(400.00), + kv_mib: None, + host_memory_observed: false, + } + ); + } + + #[test] + fn aggregator_reset_clears_stale_measured_buffers_on_model_load() { + // Review blocker (PR #1719): the model-load "loaded meta data" line + // fires reset_model_loading_state mid-open. A new model load + // invalidates the previous model's measured buffer sizes (buffer + // scales are model/shape-specific), so the reset clears them; buffer + // lines emitted after the reset are measured normally, and the host + // plan tuple (model/context/lanes) lives outside the aggregator + // entirely so it is untouched by the reset. + let mut aggregator = NativeLogAggregator::default(); + aggregator.process_line("sched_reserve: CUDA0 compute buffer size = 579.83 MiB"); + // Use a fully parsable line so `process_line` actually invokes + // reset_model_loading_state (see parse_loaded_metadata_counts). + aggregator.process_line( + "llama_model_loader: loaded meta data with 26 key-value pairs and 291 tensors from model.gguf (version GGUF V3)", + ); + // The pre-reset measurement belongs to the previous model and must be + // cleared, not stranded into the new model's footprint. + assert_eq!( + aggregator.measured_snapshot(), + MeasuredNativeBuffers { + compute_mib: None, + kv_mib: None, + host_memory_observed: false, + } + ); + aggregator.process_line("llama_kv_cache: CUDA0 KV buffer size = 1088.00 MiB"); + assert_eq!( + aggregator.measured_snapshot(), + MeasuredNativeBuffers { + compute_mib: None, + kv_mib: Some(1088.00), + host_memory_observed: false, + } + ); + } + + #[test] + fn aggregator_parses_measured_compute_buffer_size() { + let mut aggregator = NativeLogAggregator::default(); + assert_eq!( + aggregator + .process_line("sched_reserve: CUDA0 compute buffer size = 579.83 MiB"), + vec![NativeLogEvent { + message: "sched_reserve: CUDA0 compute buffer size = 579.83 MiB" + .to_string(), + category: "memory", + params: vec![ + ("buffer_mib".to_string(), Value::from(579.83_f64)), + ( + "backend_device".to_string(), + Value::String("CUDA0".to_string()) + ), + ], + }] + ); + } + + #[test] + fn aggregator_parses_measured_kv_buffer_size() { + let mut aggregator = NativeLogAggregator::default(); + assert_eq!( + aggregator.process_line("llama_kv_cache: CUDA0 KV buffer size = 1088.00 MiB"), + vec![NativeLogEvent { + message: "llama_kv_cache: CUDA0 KV buffer size = 1088.00 MiB".to_string(), + category: "kv_cache", + params: vec![ + ("buffer_mib".to_string(), Value::from(1088.00_f64)), + ( + "backend_device".to_string(), + Value::String("CUDA0".to_string()) + ), + ], + }] + ); + } + + #[test] + fn aggregator_parses_metal_compute_buffer_size() { + let mut aggregator = NativeLogAggregator::default(); + assert_eq!( + aggregator + .process_line("sched_reserve: Metal compute buffer size = 312.50 MiB"), + vec![NativeLogEvent { + message: "sched_reserve: Metal compute buffer size = 312.50 MiB" + .to_string(), + category: "memory", + params: vec![ + ("buffer_mib".to_string(), Value::from(312.5_f64)), + ( + "backend_device".to_string(), + Value::String("Metal".to_string()) + ), + ], + }] + ); + } + #[test] fn aggregator_preserves_memory_summary_lines() { let mut aggregator = NativeLogAggregator::default(); diff --git a/crates/skippy-server/src/binary_transport/binary_messaging.rs b/crates/skippy-server/src/binary_transport/binary_messaging.rs index 7d15876de8..ece1445ca2 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging.rs @@ -277,6 +277,7 @@ fn run_binary_stage( native_mtp_enabled, continuous_batching, openai, + l3_manager, } = options; let native_mtp_enabled = native_mtp_enabled && config.native_mtp_enabled; validate_config(&config, topology.as_ref())?; @@ -346,9 +347,10 @@ fn run_binary_stage( telemetry.clone(), ) .map_err(|error| anyhow!("create binary iteration scheduler: {error}"))?; - let kv = KvStageIntegration::from_loaded_model( + let kv = KvStageIntegration::from_loaded_model_with_l3_manager( &config, loaded_model_state_kind(Some(&runtime)), + l3_manager.clone(), None, )? .map(Arc::new); @@ -416,6 +418,7 @@ fn run_binary_stage( openai_guardrails: Some( frontend::OpenAiGuardrailsConfig::disabled_for_skippy(), ), + l3_manager, }, openai_iteration_scheduler, ) diff --git a/crates/skippy-server/src/binary_transport/options.rs b/crates/skippy-server/src/binary_transport/options.rs index 376e8b3f2e..caed5aba1b 100644 --- a/crates/skippy-server/src/binary_transport/options.rs +++ b/crates/skippy-server/src/binary_transport/options.rs @@ -32,6 +32,8 @@ pub struct BinaryStageOptions { /// the stage-control load request. pub continuous_batching: bool, pub openai: Option, + /// Shared node owner for the durable disk tier. + pub l3_manager: Option, } #[derive(Clone)] @@ -149,6 +151,7 @@ impl BinaryStageOptions { native_mtp_enabled, continuous_batching: true, openai, + l3_manager: None, }) } diff --git a/crates/skippy-server/src/frontend/generation/queue.rs b/crates/skippy-server/src/frontend/generation/queue.rs index 3ead8ccaca..af96b5905c 100644 --- a/crates/skippy-server/src/frontend/generation/queue.rs +++ b/crates/skippy-server/src/frontend/generation/queue.rs @@ -804,10 +804,11 @@ pub(in crate::frontend) fn prewarm_generation_sessions( event_name: &'static str, ) -> Result<()> { let timer = PhaseTimer::start(); - let sessions = runtime + let mut runtime = runtime .lock() - .map_err(|_| anyhow!("runtime lock poisoned"))? - .prewarm_idle_sessions(generation_concurrency)?; + .map_err(|_| anyhow!("runtime lock poisoned"))?; + let generation_graph_warmed = runtime.warmup_generation_graph()?; + let sessions = runtime.prewarm_idle_sessions(generation_concurrency)?; let mut attrs = lifecycle_attrs(config); attrs.insert( "llama_stage.generation_concurrency".to_string(), @@ -825,6 +826,10 @@ pub(in crate::frontend) fn prewarm_generation_sessions( "llama_stage.runtime_sessions_idle".to_string(), json!(sessions.idle_sessions), ); + attrs.insert( + "llama_stage.generation_graph_warmed".to_string(), + json!(generation_graph_warmed), + ); attrs.insert( "llama_stage.elapsed_ms".to_string(), json!(timer.elapsed_ms()), diff --git a/crates/skippy-server/src/frontend/generation/server.rs b/crates/skippy-server/src/frontend/generation/server.rs index ab0a5b9b4a..c17c23770b 100644 --- a/crates/skippy-server/src/frontend/generation/server.rs +++ b/crates/skippy-server/src/frontend/generation/server.rs @@ -250,6 +250,9 @@ pub struct EmbeddedOpenAiArgs { pub linear_proposal_ingress: Option, pub openai_guardrails: Option, pub kv_lifecycle_observer: Option>, + /// Node-scoped durable disk-cache owner supplied by the embedding host. + /// `None` keeps standalone and cache-disabled launches in-memory only. + pub l3_manager: Option, } #[derive(Clone, Debug, Default, PartialEq)] @@ -501,9 +504,10 @@ fn embedded_openai_backend_with_scheduler( "stage.openai_runtime_prewarm", ) .context("prewarm embedded OpenAI runtime sessions")?; - let kv = KvStageIntegration::from_loaded_model( + let kv = KvStageIntegration::from_loaded_model_with_l3_manager( &args.config, loaded_model_state_kind(Some(&args.runtime)), + args.l3_manager.clone(), args.kv_lifecycle_observer.clone(), )? .map(Arc::new); diff --git a/crates/skippy-server/src/frontend/local_generation/token_generation/kv_restore.rs b/crates/skippy-server/src/frontend/local_generation/token_generation/kv_restore.rs index c7e1e56f20..41d9c2cd95 100644 --- a/crates/skippy-server/src/frontend/local_generation/token_generation/kv_restore.rs +++ b/crates/skippy-server/src/frontend/local_generation/token_generation/kv_restore.rs @@ -37,6 +37,20 @@ impl StageOpenAiBackend { "skippy.exact_cache.restored_tokens".to_string(), json!(restored.token_count), ); + attrs.insert( + "skippy.exact_cache.source".to_string(), + json!(restored.source), + ); + if restored.source == "l3" { + attrs.insert( + "skippy.exact_cache.fill_ms".to_string(), + json!(restored.fill_ms), + ); + attrs.insert( + "skippy.exact_cache.rewarm_enqueued".to_string(), + json!(restored.rewarm_enqueued), + ); + } attrs.insert( "skippy.kv.matched_prefix_tokens".to_string(), json!(restored.token_count), diff --git a/crates/skippy-server/src/frontend/tests/multimodal.rs b/crates/skippy-server/src/frontend/tests/multimodal.rs index e82231c1d8..b1347774a7 100644 --- a/crates/skippy-server/src/frontend/tests/multimodal.rs +++ b/crates/skippy-server/src/frontend/tests/multimodal.rs @@ -469,6 +469,7 @@ async fn real_multimodal_split_smoke_when_fixture_is_set() -> Result<()> { native_mtp_enabled: true, continuous_batching: true, openai: None, + l3_manager: None, }); let ready = connect_endpoint_ready(&stage1_addr.to_string(), 120); if let Err(error) = ready { diff --git a/crates/skippy-server/src/kv_integration/config.rs b/crates/skippy-server/src/kv_integration/config.rs index c82905807f..968a4e5231 100644 --- a/crates/skippy-server/src/kv_integration/config.rs +++ b/crates/skippy-server/src/kv_integration/config.rs @@ -6,11 +6,12 @@ use std::{ use anyhow::Result; use mesh_llm_events::OutputEvent; use skippy_cache::{ - CacheBlobStore, ResidentActivationCache, ResidentCacheConfig, SparseCheckpointPolicy, - UnifiedRadixCache, + CacheBlobStore, GeometryBlock, GeometryKind, L3CacheManager, L3Tier, PayloadGeometry, + ResidentActivationCache, ResidentCacheConfig, SparseCheckpointPolicy, StoreLimits, + UnifiedRadixCache, exact_state_identity_for_stage, numerical_model_identity_for_stage, }; use skippy_protocol::{StageConfig, StageKvCacheConfig, StageKvCacheMode, StageKvCachePayload}; -use skippy_runtime::ModelStateKind; +use skippy_runtime::{ModelStateKind, RuntimeKvPageDesc}; use super::{ EXACT_STATE_RECORD_CAPACITY, ExactStateByteLimits, KvLifecycleEvent, KvLifecycleObserver, @@ -65,6 +66,32 @@ impl KvStageIntegration { config: &StageConfig, model_state_kind: Option, observer: Option>, + ) -> Result> { + Self::from_loaded_model_with_l3( + config, + model_state_kind, + || l3_manager_from_env(config), + observer, + ) + } + + /// Construct a stage with an explicitly injected node cache manager. + /// Embedders that own several stages can use this path without consulting + /// process environment or opening the root again. + pub fn from_loaded_model_with_l3_manager( + config: &StageConfig, + model_state_kind: Option, + manager: Option, + observer: Option>, + ) -> Result> { + Self::from_loaded_model_with_l3(config, model_state_kind, || Ok(manager), observer) + } + + fn from_loaded_model_with_l3( + config: &StageConfig, + model_state_kind: Option, + manager: impl FnOnce() -> Result>, + observer: Option>, ) -> Result> { let Some(mut cache_config) = effective_cache_config(config) else { return Ok(None); @@ -86,6 +113,7 @@ impl KvStageIntegration { if payload == StagePrefixCachePayload::Disabled { return Ok(None); } + let dense_without_recurrent = matches!(model_capability, ModelKvCapability::KnownDense); if matches!(model_capability, ModelKvCapability::KnownRecurrent) && matches!(payload, StagePrefixCachePayload::ResidentKv) { @@ -104,6 +132,21 @@ impl KvStageIntegration { ); return Ok(None); } + let l3_manager = manager()?; + let durable_payload = l3_manager.as_ref().map(|_| { + if payload == StagePrefixCachePayload::ResidentKv && dense_without_recurrent { + // Resident KV stays the in-process fast path. Dense families + // export KV pages with an empty recurrent snapshot for L3; + // the known-dense capability makes that empty component valid. + StagePrefixCachePayload::KvRecurrent + } else { + payload + } + }); + let l3 = l3_manager + .zip(durable_payload) + .map(|(manager, payload)| l3_tier_for_manager(config, payload, manager)) + .transpose()?; // FullState is architecture-neutral: the native runtime serializes the // complete session state for both dense and recurrent model families. if matches!(model_capability, ModelKvCapability::KnownRecurrent) { @@ -134,8 +177,19 @@ impl KvStageIntegration { std::sync::mpsc::sync_channel::(EXACT_STATE_RECORD_CAPACITY); let worker_radix = radix.clone(); let worker_exact_blobs = exact_blobs.clone(); - let inflight_records = Arc::new(Mutex::new(BTreeSet::new())); + let worker_l3 = l3.clone(); + let inflight_records: Arc>> = l3.as_ref().map_or_else( + || Arc::new(Mutex::new(BTreeSet::new())), + |tier| tier.manager().record_claims(tier.state_identity()), + ); let worker_inflight_records = inflight_records.clone(); + let inflight_fills: Arc>> = l3.as_ref().map_or_else( + || Arc::new(Mutex::new(BTreeSet::new())), + |tier| tier.manager().fill_claims(), + ); + let worker_inflight_fills = inflight_fills.clone(); + let exact_state_record_queue_bytes = Arc::new(std::sync::atomic::AtomicU64::new(0)); + let worker_exact_state_record_queue_bytes = exact_state_record_queue_bytes.clone(); let exact_state_records_queued = Arc::new(std::sync::atomic::AtomicU64::new(0)); let exact_state_records_dropped = Arc::new(std::sync::atomic::AtomicU64::new(0)); let worker_exact_state_records_dropped = exact_state_records_dropped.clone(); @@ -146,15 +200,17 @@ impl KvStageIntegration { let exact_state_record_worker_panics = Arc::new(std::sync::atomic::AtomicU64::new(0)); let worker_exact_state_record_worker_panics = exact_state_record_worker_panics.clone(); let worker_observer = observer.clone(); - let spawned = std::thread::Builder::new() + let exact_state_record_task = std::thread::Builder::new() .name(format!("skippy-exact-cache-{}", config.stage_id)) .spawn(move || { while let Ok(pending) = exact_state_record_rx.recv() { + let fill_claim = pending.l3_fill_claim.clone(); super::run_exact_state_record_job( super::ExactStateWorkerHandles { inflight_records: &worker_inflight_records, dropped: &worker_exact_state_records_dropped, pending_count: &worker_exact_state_records_pending, + queue_bytes: &worker_exact_state_record_queue_bytes, worker_healthy: &worker_exact_state_record_worker_healthy, worker_panics: &worker_exact_state_record_worker_panics, }, @@ -166,24 +222,39 @@ impl KvStageIntegration { &worker_exact_blobs, exact_max_entries, exact_byte_limits, + worker_l3.as_deref(), pending, ) }, ); + if let Some(fill_claim) = fill_claim { + // The filled entry is radix-resident now (or the + // insert failed, and a re-fill is the right call + // either way): release the claim. + worker_inflight_fills + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&fill_claim); + } } }); - if spawned.is_err() + if exact_state_record_task.is_err() && let Some(observer) = observer.as_ref() { observer.observe(KvLifecycleEvent::KvInitFailed); } - spawned?; + let exact_state_record_task = exact_state_record_task?; if let Some(observer) = observer.as_ref() { observer.observe(KvLifecycleEvent::KvInitCompleted); } + let exact_state_record_worker = Arc::new(super::ExactStateRecordWorker::new( + exact_state_record_tx, + exact_state_record_task, + )); Ok(Some(Self { mode, payload, + durable_payload, correctness_mode: false, trust_local_writes: true, checkpoint_policy, @@ -198,7 +269,7 @@ impl KvStageIntegration { exact_blobs, exact_max_entries, exact_byte_limits, - exact_state_record_tx, + exact_state_record_worker, exact_state_records_queued, exact_state_records_dropped, exact_state_records_pending, @@ -208,6 +279,10 @@ impl KvStageIntegration { output_tokens: Arc::new(Mutex::new(OutputTokenCache::new(exact_max_entries))), split_prefill_tokens: Arc::new(Mutex::new(BTreeMap::new())), kv_lifecycle_observer: observer, + exact_state_record_queue_bytes, + l3, + inflight_fills, + dense_without_recurrent, })) } @@ -220,6 +295,221 @@ impl KvStageIntegration { } } +/// Open the durable L3 tier when `SKIPPY_L3_DIR` is set. +/// +/// Experimental, environment-only plumbing: the public `[runtime.kv_cache.disk]` +/// configuration replaces it and this reader becomes a one-release +/// compatibility shim with a deprecation warning. The tier identity is the +/// radix cache's own numerical namespace, so restarts reuse it and a change +/// in weights, cache dtypes, layout or platform refuses stale state. +const DEFAULT_L3_BUDGET_BYTES: u64 = 32 * 1024 * 1024 * 1024; +const DEFAULT_L3_MINIMUM_FREE_BYTES: u64 = 16 * 1024 * 1024 * 1024; +const L3_SEGMENT_BYTES: usize = 8 * 1024 * 1024; +/// Cap on rows per segment. +/// +/// Windows are the dedupe granularity: a turn leaves a partial window in every +/// run, and those rows are rewritten next turn, so amplification is roughly +/// `1 + (tokens_so_far mod window) / new_tokens`. Measured on an M4 mini with +/// 2000-token base and 300-token turns: 512 rows gives 1.92x over the soak +/// (worst turn 2.55x) and misses §13.4's 1.2x gate; 128 rows gives 1.18x with +/// no margin; 64 rows gives 1.07x (worst turn 1.17x). The cost of the smaller +/// window is segment count, which is why it is capped rather than shrunk +/// further. +const L3_MAX_WINDOW_ROWS: u64 = 64; + +fn l3_manager_from_env(config: &StageConfig) -> Result> { + let Ok(root) = std::env::var("SKIPPY_L3_DIR") else { + return Ok(None); + }; + if root.trim().is_empty() { + return Ok(None); + } + let budget_bytes = match std::env::var("SKIPPY_L3_BUDGET_BYTES") + .ok() + .and_then(|value| value.parse::().ok()) + { + // There is no unbounded mode. Zero used to mean "no cap"; it is now + // rejected rather than silently reinterpreted. + Some(0) => { + let _ = mesh_llm_events::emit_event(OutputEvent::Warning { + message: "SKIPPY_L3_BUDGET_BYTES=0 is not unbounded; using the default budget" + .to_string(), + context: Some(format!("budget_bytes={DEFAULT_L3_BUDGET_BYTES}")), + }); + DEFAULT_L3_BUDGET_BYTES + } + Some(bytes) => bytes, + None => DEFAULT_L3_BUDGET_BYTES, + }; + let manager = match L3CacheManager::acquire( + &root, + StoreLimits::new(budget_bytes, DEFAULT_L3_MINIMUM_FREE_BYTES), + ) { + Ok(manager) => manager, + Err(error) => { + // A tier that cannot open is a visible decline, not a crash: + // serving continues cache-off for this stage. + let _ = mesh_llm_events::emit_event(OutputEvent::Warning { + message: "Skippy L3 disk cache disabled for this model stage".to_string(), + context: Some(format!( + "stage_id={} directory={root} reason={error:#}", + config.stage_id + )), + }); + return Ok(None); + } + }; + Ok(Some(manager)) +} + +fn l3_tier_for_manager( + config: &StageConfig, + payload: StagePrefixCachePayload, + manager: L3CacheManager, +) -> Result> { + let identity = exact_state_identity_for_stage(config, l3_payload_kind(payload)); + let model_identity = numerical_model_identity_for_stage(config); + let root = manager.root().display().to_string(); + let tier = manager.tier_for_model(model_identity, identity, L3_SEGMENT_BYTES); + // Warm state must never be invisible: say what the tier can restore the + // moment the stage comes up. + match tier.status() { + Ok(status) => { + let _ = mesh_llm_events::emit_event(OutputEvent::Info { + message: "Skippy L3 disk cache open".to_string(), + context: Some(format!( + "stage_id={} directory={root} restorable_manifests={} restorable_tokens={} used_bytes={} budget_bytes={}", + config.stage_id, + status.restorable_manifests, + status.restorable_tokens, + status.usage.used_bytes, + status.usage.budget_bytes + )), + }); + } + Err(error) => { + let _ = mesh_llm_events::emit_event(OutputEvent::Warning { + message: "Skippy L3 disk cache open; status unavailable".to_string(), + context: Some(format!("stage_id={} reason={error:#}", config.stage_id)), + }); + } + } + Ok(Arc::new(tier)) +} + +fn l3_payload_kind(payload: StagePrefixCachePayload) -> &'static str { + match payload { + StagePrefixCachePayload::KvRecurrent => "kv-recurrent", + StagePrefixCachePayload::FullState => "full-state", + StagePrefixCachePayload::ResidentKv | StagePrefixCachePayload::Disabled => "unsupported", + } +} + +/// Describe an exported KV page so the store can cut segments where a growing +/// prefix keeps its bytes still. +/// +/// The runtime writes every selected layer's K rows, then every layer's V rows, +/// then the indexer rows, each run holding one row per token in token order +/// (`llama_kv_cache::stage_export_kv_page`). Cutting that on fixed byte offsets +/// re-writes the whole payload every turn, because adding tokens shifts every +/// run after the first — measured at 8x on an M4 mini. Cutting per run into +/// fixed token windows writes only the new rows. +/// +/// Returns `None` when the layout is not one this mapping can state exactly; +/// the store then falls back to fixed-size cutting, which is correct but not +/// cheap. +fn kv_page_geometry(desc: &RuntimeKvPageDesc, payload_bytes: u64) -> Option { + // A composite (ISWA) page is two independently-shaped components; its + // geometry is not this single-run description. + if desc.component_count != 0 || desc.token_count == 0 || desc.layer_count == 0 { + return None; + } + let rows = desc.token_count; + let layers = desc.layer_count; + let k_stride = u64::from(desc.k_row_bytes); + if k_stride == 0 { + return None; + } + let mut blocks = Vec::new(); + for layer in 0..layers { + blocks.push(GeometryBlock { + stride: k_stride, + kind: GeometryKind::Key, + layer, + column: 0, + }); + } + let transposed = desc.flags & skippy_runtime::KV_PAGE_FLAG_V_TRANSPOSED != 0; + if transposed { + // Transposed V is stored column-major, but each column is still one + // contiguous run of one element per token, so it windows the same way. + // The column count is not in the descriptor; derive it from the bytes + // the K and indexer runs do not claim. + let element_bytes = u64::from(desc.v_element_bytes); + let k_idx_stride = u64::from(desc.k_idx_row_bytes); + let claimed = u64::from(layers) + .checked_mul(rows)? + .checked_mul(k_stride.checked_add(k_idx_stride)?)?; + let v_bytes = payload_bytes.checked_sub(claimed)?; + let per_layer = v_bytes.checked_div(u64::from(layers))?; + let column_bytes = rows.checked_mul(element_bytes)?; + if element_bytes == 0 || column_bytes == 0 || per_layer % column_bytes != 0 { + return None; + } + let columns = u32::try_from(per_layer / column_bytes).ok()?; + for layer in 0..layers { + for column in 0..columns { + blocks.push(GeometryBlock { + stride: element_bytes, + kind: GeometryKind::Value, + layer, + column, + }); + } + } + } else if desc.v_row_bytes > 0 { + for layer in 0..layers { + blocks.push(GeometryBlock { + stride: u64::from(desc.v_row_bytes), + kind: GeometryKind::Value, + layer, + column: 0, + }); + } + } + if desc.k_idx_row_bytes > 0 { + for layer in 0..layers { + blocks.push(GeometryBlock { + stride: u64::from(desc.k_idx_row_bytes), + kind: GeometryKind::KeyIndex, + layer, + column: 0, + }); + } + } + // The window must depend only on the model's shape, never on this entry's + // token count, or the boundaries move between turns and nothing dedupes. + let widest = blocks.iter().map(|block| block.stride).max()?; + let window_rows = (L3_SEGMENT_BYTES as u64 / widest.max(1)) + .clamp(1, L3_MAX_WINDOW_ROWS) + .next_power_of_two() + .min(L3_MAX_WINDOW_ROWS); + let geometry = PayloadGeometry { + blocks, + rows, + window_rows, + // A recurrent snapshot rides after the KV page and has no row + // structure; it is whatever the payload has left. + tail_bytes: 0, + }; + let described = geometry.total_bytes(); + let tail = payload_bytes.checked_sub(described)?; + Some(PayloadGeometry { + tail_bytes: tail, + ..geometry + }) +} + fn emit_cache_disabled_warning(config: &StageConfig, reason: &str) { let _ = mesh_llm_events::emit_event(OutputEvent::Warning { message: "Skippy KV cache disabled for this model stage".to_string(), @@ -230,13 +520,59 @@ fn emit_cache_disabled_warning(config: &StageConfig, reason: &str) { }); } +fn emit_l3_state_transitions(l3: &L3Tier) { + for transition in l3.manager().take_state_transitions() { + let context = serde_json::to_string(&transition).ok(); + let _ = mesh_llm_events::emit_event(OutputEvent::Info { + message: "Skippy L3 disk cache state changed".to_string(), + context, + }); + } +} + fn store_exact_radix_record( radix: &Mutex>, blobs: &Mutex, max_entries: usize, limits: ExactStateByteLimits, + l3: Option<&L3Tier>, pending: PendingExactStateRecord, ) -> Result<()> { + // Write through to the durable tier before the payload is deduplicated + // into blocks, while its bytes are still contiguous. Best-effort: a full + // or failing disk must not fail the in-memory record. The refusal reason + // lands in the tier's status; one warning per process keeps a full disk + // from flooding the log. + if let Some(l3) = l3 { + let kv_desc_json = pending + .extra + .kv_desc + .as_ref() + .and_then(|desc| serde_json::to_string(desc).ok()); + let geometry = pending + .extra + .kv_desc + .as_ref() + .and_then(|desc| kv_page_geometry(desc, pending.payload.byte_len())); + let spill = l3.spill( + &pending.namespace, + &pending.token_ids, + &pending.payload, + kv_desc_json, + geometry.as_ref(), + ); + emit_l3_state_transitions(l3); + if let Err(error) = spill { + static WARNED: std::sync::atomic::AtomicBool = + std::sync::atomic::AtomicBool::new(false); + if !WARNED.swap(true, std::sync::atomic::Ordering::AcqRel) { + let _ = mesh_llm_events::emit_event(OutputEvent::Warning { + message: "Skippy L3 disk cache write refused; see kv-cache status".to_string(), + context: Some(format!("page_id={} reason={error:#}", pending.page_id)), + }); + } + } + } let logical_bytes = pending.payload.byte_len(); let (payload, _) = pending.payload.dedupe_into( &mut blobs @@ -496,6 +832,7 @@ mod tests { extra: super::super::ExactStateExtra::default(), namespace: "model".to_string(), token_ids: tokens.to_vec(), + l3_fill_claim: None, } } @@ -509,6 +846,7 @@ mod tests { &blobs, 1, limits(0, 0), + None, pending("first", &[1, 2], b"aaaabbbb"), ) .unwrap(); @@ -517,6 +855,7 @@ mod tests { &blobs, 1, limits(0, 0), + None, pending("second", &[1, 3], b"aaaacccc"), ) .unwrap(); @@ -546,6 +885,7 @@ mod tests { &blobs, 1, limits(0, 0), + None, pending("empty", &[], b"aaaabbbb"), ) .unwrap_err(); @@ -558,6 +898,166 @@ mod tests { assert_eq!(radix.lock().unwrap().stats().recurrent_entries, 0); } + #[test] + fn exact_records_write_through_to_l3_and_survive_radix_eviction() { + let radix = Mutex::new(UnifiedRadixCache::new()); + let blobs = Mutex::new(CacheBlobStore::new(4)); + let root = std::env::temp_dir() + .join("skippy-server-l3-tests") + .join(format!("write-through-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&root); + let tier = L3Tier::open(&root, 0, "blake3:test-tier".to_string(), 4096).unwrap(); + + // Two records with max_entries = 1: the first is evicted from RAM. + store_exact_radix_record( + &radix, + &blobs, + 1, + limits(0, 0), + Some(&tier), + pending("first", &[1, 2], b"first-exact-state"), + ) + .unwrap(); + store_exact_radix_record( + &radix, + &blobs, + 1, + limits(0, 0), + Some(&tier), + pending("second", &[1, 3], b"second-exact-state"), + ) + .unwrap(); + assert!( + radix + .lock() + .unwrap() + .lookup_recurrent("model", &[1, 2]) + .is_none(), + "first record must be evicted from RAM" + ); + + // The evicted prefix still fills from the durable tier, including + // for a longer query that extends the recorded path. + let fill = tier + .fill_longest("model", &[1, 2, 7, 8], 64) + .unwrap() + .expect("evicted record must remain in L3"); + assert_eq!(fill.token_count, 2); + assert_eq!( + fill.payload + .full_state_bytes_timed() + .unwrap() + .0 + .into_owned(), + b"first-exact-state".to_vec() + ); + let status = tier.status().unwrap(); + assert_eq!(status.activity.writes, 2); + assert_eq!(status.activity.fills, 1); + } + + fn kv_desc(layers: u32, tokens: u64, k_row: u32, v_row: u32) -> RuntimeKvPageDesc { + let mut desc = RuntimeKvPageDesc { + version: 1, + layer_start: 0, + layer_end: layers as i32, + token_start: 0, + token_count: tokens, + layer_count: layers, + k_type: 1, + v_type: 1, + k_row_bytes: k_row, + v_row_bytes: v_row, + v_element_bytes: 2, + k_idx_row_bytes: 0, + payload_bytes: 0, + flags: 0, + codec: 0, + component_count: 0, + components: Default::default(), + }; + desc.payload_bytes = u64::from(layers) * tokens * (u64::from(k_row) + u64::from(v_row)); + desc + } + + #[test] + fn kv_page_geometry_describes_the_runtime_export_layout() { + // Every layer's K rows, then every layer's V rows: the order + // `stage_export_kv_page` writes them in. + let desc = kv_desc(4, 2048, 1024, 1024); + let geometry = + kv_page_geometry(&desc, desc.payload_bytes).expect("dense page must be describable"); + + assert_eq!(geometry.blocks.len(), 8); + assert_eq!(geometry.rows, 2048); + assert!( + geometry.blocks[..4] + .iter() + .all(|block| block.kind == GeometryKind::Key) + ); + assert!( + geometry.blocks[4..] + .iter() + .all(|block| block.kind == GeometryKind::Value) + ); + assert_eq!(geometry.total_bytes(), desc.payload_bytes); + assert!(geometry.matches(desc.payload_bytes)); + } + + #[test] + fn kv_page_geometry_windows_are_stable_as_the_prefix_grows() { + // The property the whole fix rests on: the same model must produce the + // same window size at every length, or turn N+1's cuts land elsewhere + // and nothing is reused. + let short = kv_desc(4, 2048, 1024, 1024); + let long = kv_desc(4, 4096, 1024, 1024); + let short_geometry = kv_page_geometry(&short, short.payload_bytes).expect("short"); + let long_geometry = kv_page_geometry(&long, long.payload_bytes).expect("long"); + assert_eq!(short_geometry.window_rows, long_geometry.window_rows); + assert_eq!(short_geometry.blocks, long_geometry.blocks); + } + + #[test] + fn kv_page_geometry_accounts_for_a_recurrent_tail() { + let desc = kv_desc(2, 512, 512, 512); + let tail = 4096; + let geometry = kv_page_geometry(&desc, desc.payload_bytes + tail).expect("hybrid page"); + assert_eq!(geometry.tail_bytes, tail); + assert!(geometry.matches(desc.payload_bytes + tail)); + } + + #[test] + fn kv_page_geometry_declines_what_it_cannot_state_exactly() { + // A composite ISWA page is two differently-shaped components. + let mut composite = kv_desc(4, 512, 1024, 1024); + composite.component_count = 2; + assert!(kv_page_geometry(&composite, composite.payload_bytes).is_none()); + + // Bytes that the described runs cannot account for. + let desc = kv_desc(4, 512, 1024, 1024); + assert!(kv_page_geometry(&desc, desc.payload_bytes - 1).is_none()); + } + + #[test] + fn kv_page_geometry_windows_transposed_value_columns() { + // Transposed V is column-major, but each column is one contiguous run + // of one element per token, so it windows like any other run. + let mut desc = kv_desc(2, 256, 1024, 0); + desc.flags = skippy_runtime::KV_PAGE_FLAG_V_TRANSPOSED; + desc.v_element_bytes = 2; + let columns = 512u64; + let payload = desc.payload_bytes + u64::from(desc.layer_count) * columns * 256 * 2; + let geometry = kv_page_geometry(&desc, payload).expect("transposed page"); + + let value_blocks = geometry + .blocks + .iter() + .filter(|block| block.kind == GeometryKind::Value) + .count(); + assert_eq!(value_blocks as u64, u64::from(desc.layer_count) * columns); + assert!(geometry.matches(payload)); + } + #[test] fn oversized_exact_payloads_retain_a_reusable_working_set() { let radix = Mutex::new(UnifiedRadixCache::new()); @@ -575,6 +1075,7 @@ mod tests { &blobs, 8, limits(4, 1024), + None, pending(page_id, &tokens, bytes), ) .unwrap(); @@ -605,6 +1106,7 @@ mod tests { &blobs, 2, limits(4, 1_024), + None, pending(page_id, &tokens, bytes), ) .unwrap(); @@ -632,6 +1134,7 @@ mod tests { &blobs, 8, limits(4, 8), + None, pending(page_id, &tokens, bytes), ) .unwrap(); @@ -657,6 +1160,7 @@ mod tests { &blobs, 8, limits(2, 4), + None, pending("checkpoint", &[1, 2], b"aaaabbbb"), ) .unwrap(); @@ -809,6 +1313,95 @@ mod tests { assert_eq!(kv.exact_max_entries, RECURRENT_CACHE_MAX_ENTRIES); } + #[test] + fn injected_manager_is_shared_across_placement_equivalent_stages() { + let root = std::env::temp_dir() + .join("skippy-server-l3-manager-tests") + .join(format!("shared-stages-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&root); + let manager = L3CacheManager::acquire(&root, StoreLimits::new(1_000_000, 0)).unwrap(); + let mut first = enabled_auto_config("future/model"); + first.kv_cache.as_mut().unwrap().payload = StageKvCachePayload::FullState; + let second = StageConfig { + stage_id: "replica-stage".to_string(), + stage_index: 7, + topology_id: "other-topology".to_string(), + run_id: "other-run".to_string(), + ..first.clone() + }; + + let first = KvStageIntegration::from_loaded_model_with_l3_manager( + &first, + Some(ModelStateKind::Dense), + Some(manager.clone()), + None, + ) + .unwrap() + .unwrap(); + let second = KvStageIntegration::from_loaded_model_with_l3_manager( + &second, + Some(ModelStateKind::Dense), + Some(manager), + None, + ) + .unwrap() + .unwrap(); + let first_l3 = first.l3.as_ref().unwrap(); + let second_l3 = second.l3.as_ref().unwrap(); + + assert!(first_l3.manager().shares_root_with(second_l3.manager())); + assert_eq!(first_l3.state_identity(), second_l3.state_identity()); + assert_eq!(first_l3.model_identity(), second_l3.model_identity()); + assert!(Arc::ptr_eq(&first.inflight_fills, &second.inflight_fills)); + assert!(Arc::ptr_eq( + &first.inflight_records, + &second.inflight_records + )); + assert!(first.try_begin_record("shared-page")); + assert!( + !second.try_begin_record("shared-page"), + "placement replicas performed duplicate record work" + ); + first.finish_record("shared-page"); + assert!(second.try_begin_record("shared-page")); + second.finish_record("shared-page"); + } + + #[test] + fn dense_disk_cache_preserves_resident_fast_path_and_exports_exact_state() { + let root = std::env::temp_dir() + .join("skippy-server-l3-manager-tests") + .join(format!("dense-export-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&root); + let manager = L3CacheManager::acquire(&root, StoreLimits::new(1_000_000, 0)).unwrap(); + let config = enabled_auto_config("future/model"); + + let kv = KvStageIntegration::from_loaded_model_with_l3_manager( + &config, + Some(ModelStateKind::Dense), + Some(manager), + None, + ) + .unwrap() + .expect("dense disk cache should remain enabled"); + + assert_eq!(kv.payload, StagePrefixCachePayload::ResidentKv); + assert_eq!( + kv.durable_payload, + Some(StagePrefixCachePayload::KvRecurrent) + ); + assert_eq!( + kv.exact_state_payload(), + Some(StagePrefixCachePayload::KvRecurrent) + ); + assert!(kv.l3.is_some()); + + let mut invalid = kv; + invalid.payload = StagePrefixCachePayload::Disabled; + invalid.durable_payload = Some(StagePrefixCachePayload::ResidentKv); + assert_eq!(invalid.exact_state_payload(), None); + } + #[test] fn parses_cache_mode_and_payload_aliases() { assert_eq!( diff --git a/crates/skippy-server/src/kv_integration/exact_state.rs b/crates/skippy-server/src/kv_integration/exact_state.rs index 31ffee3eb3..1d5047cef6 100644 --- a/crates/skippy-server/src/kv_integration/exact_state.rs +++ b/crates/skippy-server/src/kv_integration/exact_state.rs @@ -11,6 +11,36 @@ use super::{ records::add_reconstruct_stats, }; +fn l3_fill_claim_key(l3: &skippy_cache::L3Tier, location: &skippy_cache::L3Location) -> String { + format!("{}:{}", l3.state_identity(), location.manifest_key) +} + +fn resident_prefix_is_complete(matched_tokens: usize, requested_tokens: usize) -> bool { + matched_tokens >= requested_tokens +} + +fn preflight_native_kv_location( + location: &skippy_cache::L3Location, +) -> Result> { + if !location.native_kv_passthrough { + return Ok(None); + } + let json = location + .kv_desc_json + .as_deref() + .context("native KV manifest has no runtime page descriptor")?; + let desc: skippy_runtime::RuntimeKvPageDesc = + serde_json::from_str(json).context("native KV manifest has an invalid page descriptor")?; + let kv_bytes = + usize::try_from(location.kv_bytes).context("native KV payload length exceeds usize")?; + desc.validate_payload(kv_bytes) + .context("native KV manifest page descriptor is incompatible")?; + if desc.token_start != 0 || desc.token_count != location.token_count { + anyhow::bail!("native KV manifest page descriptor does not cover the located prefix"); + } + Ok(Some(desc)) +} + impl KvStageIntegration { pub fn restore_exact_state( &self, @@ -29,7 +59,21 @@ impl KvStageIntegration { session_id: &str, identities: &[PrefillKvIdentity], ) -> Result> { - if !self.should_lookup() || !self.payload.is_exact_state() { + if !self.should_lookup() || self.exact_state_payload().is_none() { + return Ok(None); + } + // Dense L3 uses serialized exact state only as the durable floor. + // Prefer a native resident-prefix hit whenever one is already warm; + // importing the serialized snapshot would otherwise make enabling L3 + // slower than the ordinary L1 path on every repeated request. + if self.payload == StagePrefixCachePayload::ResidentKv + && identities.iter().any(|identity| { + self.probe_resident_prefix(identity) + .is_some_and(|resident| { + resident_prefix_is_complete(resident.token_count, identity.token_ids.len()) + }) + }) + { return Ok(None); } for identity in identities { @@ -44,6 +88,14 @@ impl KvStageIntegration { (lookup, entries) }; let Some(lookup) = lookup else { + // Radix miss: the durable tier may still hold this prefix. + // Runs inside the restore transaction, so a failed import + // rolls the lane back exactly as a radix restore would. + if let Some(restored) = + self.restore_from_l3(runtime, session_id, identity, lookup_started)? + { + return Ok(Some(restored)); + } continue; }; let lease = ExactStateLease { @@ -132,7 +184,7 @@ impl KvStageIntegration { .map_err(|error| { mark_deterministic_failure(&mut deterministic_failure, error) })?; - if recurrent.is_empty() { + if recurrent.is_empty() && !self.dense_without_recurrent { deterministic_failure = true; return Err(anyhow::anyhow!("cached recurrent-state payload is empty")); } @@ -143,11 +195,17 @@ impl KvStageIntegration { stats, ); let import_started = Instant::now(); - runtime.import_recurrent_state_for_token_count( - session_id, - recurrent.as_ref(), - token_count, - )?; + if recurrent.is_empty() { + // Known-dense model: there is no snapshot to + // import, only a position to finalize. + runtime.set_session_position(session_id, token_count)?; + } else { + runtime.import_recurrent_state_for_token_count( + session_id, + recurrent.as_ref(), + token_count, + )?; + } recurrent_import_ms = import_started.elapsed().as_secs_f64() * 1000.0; } _ => return Ok(false), @@ -196,6 +254,9 @@ impl KvStageIntegration { lookup_ms, kv_import_ms, recurrent_import_ms, + source: "radix", + fill_ms: 0.0, + rewarm_enqueued: false, }; drop(lease); return Ok(Some(restored)); @@ -232,7 +293,10 @@ impl KvStageIntegration { session_id: &str, identity: &PrefillKvIdentity, ) -> Result> { - if !self.should_record() || !self.payload.is_exact_state() { + let Some(exact_state_payload) = self.exact_state_payload() else { + return Ok(None); + }; + if !self.should_record() { return Ok(None); } let token_count = identity.identity.token_count; @@ -267,7 +331,7 @@ impl KvStageIntegration { self.finish_record(&identity.page_id); return Ok(None); } - let exported = match self.payload { + let exported = match exact_state_payload { StagePrefixCachePayload::FullState => { runtime.export_full_state(session_id).map(|state| { ( @@ -282,7 +346,17 @@ impl KvStageIntegration { Err(error) if is_native_kv_unavailable(&error) => None, Err(error) => return Err(error), }; - let recurrent = runtime.export_recurrent_state(session_id)?; + let recurrent = match runtime.export_recurrent_state(session_id) { + Ok(recurrent) => recurrent, + // A known-dense model has no recurrent memory to export; + // its snapshot is legitimately empty. + Err(error) + if self.dense_without_recurrent && is_recurrent_unavailable(&error) => + { + Vec::new() + } + Err(error) => return Err(error), + }; Ok(( ExactStatePayload::kv_recurrent( kv.as_ref().map(|kv| kv.payload.clone()).unwrap_or_default(), @@ -305,6 +379,13 @@ impl KvStageIntegration { return Err(error); } }; + if payload.byte_len() == 0 { + // A dense model whose native KV export was unavailable has no + // state component at all. Recording it would later restore as a + // bare position advance over missing attention state. + self.finish_record(&identity.page_id); + return Ok(None); + } let payload_kind = payload.kind(); let logical_bytes = payload.byte_len(); match self.enqueue_exact_state_record(PendingExactStateRecord { @@ -313,6 +394,7 @@ impl KvStageIntegration { extra, namespace: identity.namespace.clone(), token_ids: identity.token_ids.clone(), + l3_fill_claim: None, }) { ExactStateRecordAdmission::Queued => { // Recording owns the radix/blob locks while it hashes a potentially @@ -350,6 +432,215 @@ impl KvStageIntegration { } } +impl KvStageIntegration { + /// Fill a radix miss from the durable tier with the longest recorded + /// prefix of the query, import it, and enqueue a radix re-warm so the + /// next lookup hits RAM. `None` when there is no tier, nothing usable is + /// stored, or another fill of the same entry is in flight: concurrent + /// misses must not each read the entry from disk, so the loser prefills + /// normally while the winner warms the radix for everyone. + fn restore_from_l3( + &self, + runtime: &mut RuntimeState, + session_id: &str, + identity: &PrefillKvIdentity, + lookup_started: Instant, + ) -> Result> { + const MAX_PREFIX_PROBES: usize = 64; + let Some(l3) = &self.l3 else { + return Ok(None); + }; + // Locate first (cheap index probes), then single-flight the expensive + // load on the located entry's manifest key: same-length queries for + // different prefixes never suppress each other, and different-length + // queries resolving to one entry never load it twice. + let location = + match l3.locate_longest(&identity.namespace, &identity.token_ids, MAX_PREFIX_PROBES) { + Ok(Some(location)) => location, + // Nothing stored, or a corrupt / identity-mismatched entry. + // Either way the miss path is the safe one; the tier has + // recorded the reason for the status surface. + Ok(None) | Err(_) => return Ok(None), + }; + // Segment and manifest digests intentionally deduplicate bytes across + // numerical states. A fill claim must not: one state's fill cannot + // warm another state's radix namespace, even when their payload bytes + // happen to be identical. + let fill_claim = l3_fill_claim_key(l3, &location); + { + let mut inflight = self + .inflight_fills + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if !inflight.insert(fill_claim.clone()) { + return Ok(None); + } + } + let outcome = + self.fill_and_import(runtime, session_id, identity, lookup_started, l3, &location); + // On success the claim travels with the re-warm record and the worker + // releases it once the entry is radix-resident. On any other outcome + // release it here. + let handed_to_worker = matches!(&outcome, Ok(Some(restored)) if restored.rewarm_enqueued); + if !handed_to_worker { + self.inflight_fills + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&fill_claim); + } + outcome + } + + fn fill_and_import( + &self, + runtime: &mut RuntimeState, + session_id: &str, + identity: &PrefillKvIdentity, + lookup_started: Instant, + l3: &std::sync::Arc, + location: &skippy_cache::L3Location, + ) -> Result> { + // Native page capability is checked from manifest metadata before the + // tier reads any segment bytes. Runtime ABI, platform and numerical + // mode are already bound by the tier's exact-state identity; the page + // descriptor completes the representation check. + let Ok(native_kv_desc) = preflight_native_kv_location(location) else { + return Ok(None); + }; + let fill_started = Instant::now(); + // A load failure (corrupt segment, now quarantined) is a miss, not a + // request failure. Import failures below do propagate: the transaction + // rolls the lane back and the caller falls back to cold prefill. + let Ok(fill) = l3.load(location) else { + return Ok(None); + }; + if fill.payload.byte_len() == 0 { + return Ok(None); + } + if location.native_kv_passthrough + && (fill.token_count != location.token_count + || fill.kv_desc_json != location.kv_desc_json) + { + return Ok(None); + } + let fill_ms = fill_started.elapsed().as_secs_f64() * 1000.0; + let token_count = fill.token_count; + let kv_desc: Option = native_kv_desc.or_else(|| { + fill.kv_desc_json + .as_deref() + .and_then(|json| serde_json::from_str(json).ok()) + }); + let lookup_ms = lookup_started.elapsed().as_secs_f64() * 1000.0; + let mut kv_import_ms = 0.0; + let mut recurrent_import_ms = 0.0; + match fill.payload.kind().into() { + StagePrefixCachePayload::FullState => { + let (full_state, _) = fill + .payload + .full_state_bytes_timed() + .context("reconstruct L3 full-state payload")?; + if full_state.is_empty() { + return Ok(None); + } + let import_started = Instant::now(); + runtime.import_full_state_for_token_count( + session_id, + full_state.as_ref(), + token_count, + )?; + kv_import_ms = import_started.elapsed().as_secs_f64() * 1000.0; + } + StagePrefixCachePayload::KvRecurrent => { + // Every check runs before the first import. Once bytes have + // gone into the session, the only acceptable exit is `Err`, + // which the transaction rolls back; an `Ok(None)` after a + // partial import would hand a dirty lane to cold prefill. + let kv = fill + .payload + .kv_bytes() + .context("reconstruct L3 KV payload")?; + let recurrent = fill + .payload + .recurrent_state_bytes() + .context("reconstruct L3 recurrent payload")?; + if recurrent.is_empty() && !self.dense_without_recurrent { + return Ok(None); + } + let kv_page = match (kv.as_ref(), kv_desc.as_ref()) { + (Some(kv), Some(desc)) => { + // Same fail-closed checks as a radix restore: a + // descriptor that does not describe these bytes, or a + // page that is not the whole prefix, is a miss. + if desc.validate_payload(kv.len()).is_err() + || desc.token_start != 0 + || desc.token_count != token_count + { + return Ok(None); + } + Some((kv, desc)) + } + (Some(kv), None) if !kv.is_empty() => return Ok(None), + _ => None, + }; + + if let Some((kv, desc)) = kv_page { + let import_started = Instant::now(); + runtime.import_kv_page(session_id, desc, kv.as_ref())?; + kv_import_ms = import_started.elapsed().as_secs_f64() * 1000.0; + } + let import_started = Instant::now(); + if recurrent.is_empty() { + runtime.set_session_position(session_id, token_count)?; + } else { + runtime.import_recurrent_state_for_token_count( + session_id, + recurrent.as_ref(), + token_count, + )?; + } + recurrent_import_ms = import_started.elapsed().as_secs_f64() * 1000.0; + } + _ => return Ok(None), + } + let logical_bytes = fill.payload.byte_len(); + let payload_kind = fill.payload.kind(); + // Re-warm the RAM tier off the request path. A drop is fine: the + // disk copy stays authoritative. The fill claim rides along so the + // worker releases it only once the entry is radix-resident. + let admission = self.enqueue_exact_state_record(PendingExactStateRecord { + page_id: identity.page_id.clone(), + payload: fill.payload, + extra: ExactStateExtra { kv_desc }, + namespace: identity.namespace.clone(), + token_ids: identity.token_ids[..token_count as usize].to_vec(), + l3_fill_claim: Some(l3_fill_claim_key(l3, location)), + }); + let rewarm_enqueued = matches!(admission, ExactStateRecordAdmission::Queued); + Ok(Some(ExactStateRestore { + page_id: identity.page_id.clone(), + token_count: token_count as usize, + payload_kind, + logical_bytes, + entries: 0, + reconstruct_ms: 0.0, + reconstruct_bytes: 0, + reconstruct_blocks: 0, + lookup_ms, + kv_import_ms, + recurrent_import_ms, + source: "l3", + fill_ms, + rewarm_enqueued, + })) + } +} + +fn is_recurrent_unavailable(error: &anyhow::Error) -> bool { + error + .chain() + .any(|cause| cause.to_string().contains("no recurrent memory")) +} + struct ExactStateLease { radix: std::sync::Arc< std::sync::Mutex< @@ -441,15 +732,22 @@ mod tests { time::{Duration, Instant}, }; - use skippy_cache::UnifiedRadixCache; + use skippy_cache::{L3Location, UnifiedRadixCache}; - use super::try_touch_exact_state; + use super::{preflight_native_kv_location, resident_prefix_is_complete, try_touch_exact_state}; type TestRadix = UnifiedRadixCache< crate::kv_integration::RadixResidentEntry, crate::kv_integration::RadixExactEntry, >; + #[test] + fn only_complete_resident_prefixes_skip_exact_restore() { + assert!(resident_prefix_is_complete(4_000, 4_000)); + assert!(resident_prefix_is_complete(4_001, 4_000)); + assert!(!resident_prefix_is_complete(200, 4_000)); + } + #[test] fn busy_exact_state_lock_skips_touch_without_waiting() { let cache = Arc::new(Mutex::new(TestRadix::new())); @@ -492,4 +790,57 @@ mod tests { Some(false) ); } + + fn native_location(desc: &skippy_runtime::RuntimeKvPageDesc) -> L3Location { + L3Location { + namespace_key: "namespace".to_string(), + prefix_key: "prefix".to_string(), + token_count: desc.token_count, + manifest_key: "manifest".to_string(), + kv_desc_json: Some(serde_json::to_string(desc).unwrap()), + kv_bytes: desc.payload_bytes, + native_kv_passthrough: true, + } + } + + fn native_desc() -> skippy_runtime::RuntimeKvPageDesc { + skippy_runtime::RuntimeKvPageDesc { + version: 1, + layer_start: 0, + layer_end: 1, + token_start: 0, + token_count: 8, + layer_count: 1, + k_type: skippy_runtime::GGML_TYPE_Q8_0, + v_type: skippy_runtime::GGML_TYPE_Q8_0, + k_row_bytes: 16, + v_row_bytes: 16, + v_element_bytes: 2, + k_idx_row_bytes: 0, + payload_bytes: 256, + flags: 0, + codec: 0, + component_count: 0, + components: Box::new([Default::default(); 2]), + } + } + + #[test] + fn native_kv_descriptor_is_validated_before_segment_load() { + let desc = native_desc(); + assert_eq!( + preflight_native_kv_location(&native_location(&desc)).unwrap(), + Some(desc) + ); + + let mut wrong_length = native_desc(); + wrong_length.payload_bytes += 1; + let mut location = native_location(&wrong_length); + location.kv_bytes -= 1; + assert!(preflight_native_kv_location(&location).is_err()); + + let mut wrong_prefix = native_desc(); + wrong_prefix.token_start = 1; + assert!(preflight_native_kv_location(&native_location(&wrong_prefix)).is_err()); + } } diff --git a/crates/skippy-server/src/kv_integration/mod.rs b/crates/skippy-server/src/kv_integration/mod.rs index b6012e513a..93db8aa401 100644 --- a/crates/skippy-server/src/kv_integration/mod.rs +++ b/crates/skippy-server/src/kv_integration/mod.rs @@ -5,6 +5,7 @@ use std::{ atomic::{AtomicBool, AtomicU64, AtomicUsize}, mpsc::{SyncSender, TrySendError}, }, + thread::JoinHandle, }; use anyhow::{Result, bail}; @@ -150,7 +151,13 @@ pub(crate) struct ExactStateByteLimits { #[derive(Clone)] pub struct KvStageIntegration { pub(crate) mode: StageKvMode, + /// The in-process cache representation. Dense models keep native resident + /// KV here even when a durable tier is configured, so enabling disk does + /// not replace the fast warm path with serialized state import. pub(crate) payload: StagePrefixCachePayload, + /// Exportable representation written to and restored from L3. This is + /// separate from `payload` because resident KV is native and borrow-only. + pub(crate) durable_payload: Option, pub(crate) correctness_mode: bool, pub(crate) trust_local_writes: bool, pub(crate) checkpoint_policy: SparseCheckpointPolicy, @@ -163,7 +170,7 @@ pub struct KvStageIntegration { pub(crate) exact_blobs: Arc>, pub(crate) exact_max_entries: usize, pub(crate) exact_byte_limits: ExactStateByteLimits, - pub(crate) exact_state_record_tx: SyncSender, + pub(crate) exact_state_record_worker: Arc, pub(crate) exact_state_records_queued: Arc, pub(crate) exact_state_records_dropped: Arc, pub(crate) exact_state_records_pending: Arc, @@ -173,6 +180,23 @@ pub struct KvStageIntegration { pub(crate) output_tokens: Arc>, pub(crate) split_prefill_tokens: Arc>>>, pub(crate) kv_lifecycle_observer: Option>, + /// Payload bytes held by records queued for the worker but not yet + /// stored. The queue is bounded in bytes, not entries: an entry bound + /// lets one multi-GiB export sit next to another and doubles the RAM the + /// cache can pin behind a request. + pub(crate) exact_state_record_queue_bytes: Arc, + /// Durable L3 floor under the radix cache: exact-state records write + /// through to it on the worker, and radix misses fill back from it. + pub(crate) l3: Option>, + /// Manifest keys with an L3 fill in flight. Concurrent misses on one + /// stored prefix must not each read it from disk: the loser prefills + /// normally while the winner re-warms the radix for everyone. + pub(crate) inflight_fills: Arc>>, + /// The model has attention KV but no recurrent memory, so an exact-state + /// entry legitimately carries an empty recurrent snapshot. Restores set + /// the position directly instead of importing one. Never true for a + /// recurrent family, where an empty snapshot is corruption. + pub(crate) dense_without_recurrent: bool, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -183,7 +207,15 @@ pub enum StagePrefixCachePayload { FullState, } -pub(crate) const EXACT_STATE_RECORD_CAPACITY: usize = 1; +/// Entry backstop on the record queue. The binding limit is +/// [`EXACT_STATE_RECORD_QUEUE_BYTES`]; this only caps bookkeeping. +pub(crate) const EXACT_STATE_RECORD_CAPACITY: usize = 8; + +/// Payload bytes the record queue may hold. A record that does not fit is +/// dropped, never delayed: recording is optional and inference is not. One +/// record larger than the whole bound is still admitted when the queue is +/// empty, or large models could never record at all. +pub(crate) const EXACT_STATE_RECORD_QUEUE_BYTES: u64 = 4 * 1024 * 1024 * 1024; #[derive(Debug)] pub(crate) struct PendingExactStateRecord { @@ -192,6 +224,54 @@ pub(crate) struct PendingExactStateRecord { pub(crate) extra: ExactStateExtra, pub(crate) namespace: String, pub(crate) token_ids: Vec, + /// When this record re-warms the radix after an L3 fill, the fill's + /// claim key. The worker releases it only once the radix insert lands, + /// so requests arriving during the asynchronous re-warm prefill normally + /// instead of duplicating the disk read. + pub(crate) l3_fill_claim: Option, +} + +#[derive(Debug)] +pub(crate) struct ExactStateRecordWorker { + sender: Mutex>>, + task: Mutex>>, +} + +impl ExactStateRecordWorker { + pub(crate) fn new(sender: SyncSender, task: JoinHandle<()>) -> Self { + Self { + sender: Mutex::new(Some(sender)), + task: Mutex::new(Some(task)), + } + } + + fn with_sender( + &self, + use_sender: impl FnOnce(Option<&SyncSender>) -> T, + ) -> T { + let sender = self + .sender + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + use_sender(sender.as_ref()) + } +} + +impl Drop for ExactStateRecordWorker { + fn drop(&mut self) { + self.sender + .get_mut() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take(); + if let Some(task) = self + .task + .get_mut() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take() + { + let _ = task.join(); + } + } } #[derive(Debug, Clone)] @@ -306,20 +386,34 @@ pub(crate) enum ExactStateRecordAdmission { WorkerStopped, } -fn has_exact_state_record_capacity(pending_count: &AtomicUsize) -> bool { +fn has_exact_state_record_capacity( + pending_count: &AtomicUsize, + queue_bytes: &AtomicU64, + queue_bytes_cap: u64, +) -> bool { pending_count.load(std::sync::atomic::Ordering::Acquire) < EXACT_STATE_RECORD_CAPACITY + && queue_bytes.load(std::sync::atomic::Ordering::Acquire) < queue_bytes_cap +} + +/// Whether `bytes` more may join a queue already holding `held` bytes under +/// `cap`. A single record over the cap is admitted only into an empty queue. +fn record_fits_queue(held: u64, bytes: u64, cap: u64) -> bool { + held == 0 || held.saturating_add(bytes) <= cap } fn finish_exact_state_record( inflight_records: &Mutex>, pending_count: &AtomicUsize, + queue_bytes: &AtomicU64, page_id: &str, + bytes: u64, ) { inflight_records .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .remove(page_id); pending_count.fetch_sub(1, std::sync::atomic::Ordering::Release); + queue_bytes.fetch_sub(bytes, std::sync::atomic::Ordering::Release); } /// Shared bookkeeping handles for the exact-state-record background worker, @@ -331,6 +425,7 @@ struct ExactStateWorkerHandles<'a> { inflight_records: &'a Mutex>, dropped: &'a AtomicU64, pending_count: &'a AtomicUsize, + queue_bytes: &'a AtomicU64, worker_healthy: &'a AtomicBool, worker_panics: &'a AtomicU64, } @@ -347,6 +442,7 @@ fn run_exact_state_record_job( } }; let page_id = pending.page_id.clone(); + let bytes = pending.payload.byte_len(); if !handles .worker_healthy .load(std::sync::atomic::Ordering::Acquire) @@ -355,7 +451,13 @@ fn run_exact_state_record_job( .dropped .fetch_add(1, std::sync::atomic::Ordering::Relaxed); notify(KvLifecycleEvent::ExactStateRecordFailed); - finish_exact_state_record(handles.inflight_records, handles.pending_count, &page_id); + finish_exact_state_record( + handles.inflight_records, + handles.pending_count, + handles.queue_bytes, + &page_id, + bytes, + ); return; } @@ -380,15 +482,24 @@ fn run_exact_state_record_job( notify(KvLifecycleEvent::ExactStateRecordFailed); } } - finish_exact_state_record(handles.inflight_records, handles.pending_count, &page_id); + finish_exact_state_record( + handles.inflight_records, + handles.pending_count, + handles.queue_bytes, + &page_id, + bytes, + ); } +#[allow(clippy::too_many_arguments)] fn enqueue_exact_state_record( sender: &SyncSender, inflight_records: &Mutex>, queued: &AtomicU64, dropped: &AtomicU64, pending_count: &AtomicUsize, + queue_bytes: &AtomicU64, + queue_bytes_cap: u64, worker_healthy: &AtomicBool, pending: PendingExactStateRecord, ) -> ExactStateRecordAdmission { @@ -400,6 +511,19 @@ fn enqueue_exact_state_record( dropped.fetch_add(1, std::sync::atomic::Ordering::Relaxed); return ExactStateRecordAdmission::WorkerStopped; } + let bytes = pending.payload.byte_len(); + // Claim the bytes before the send so two producers cannot both see room. + // Released on every non-queued path below, and by the worker on finish. + let held = queue_bytes.fetch_add(bytes, std::sync::atomic::Ordering::AcqRel); + if !record_fits_queue(held, bytes, queue_bytes_cap) { + queue_bytes.fetch_sub(bytes, std::sync::atomic::Ordering::Release); + inflight_records + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&pending.page_id); + dropped.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + return ExactStateRecordAdmission::DroppedFull; + } pending_count.fetch_add(1, std::sync::atomic::Ordering::Release); match sender.try_send(pending) { Ok(()) => { @@ -408,6 +532,7 @@ fn enqueue_exact_state_record( } Err(TrySendError::Full(pending)) => { pending_count.fetch_sub(1, std::sync::atomic::Ordering::Release); + queue_bytes.fetch_sub(bytes, std::sync::atomic::Ordering::Release); inflight_records .lock() .expect("kv inflight record lock poisoned") @@ -417,6 +542,7 @@ fn enqueue_exact_state_record( } Err(TrySendError::Disconnected(pending)) => { pending_count.fetch_sub(1, std::sync::atomic::Ordering::Release); + queue_bytes.fetch_sub(bytes, std::sync::atomic::Ordering::Release); inflight_records .lock() .expect("kv inflight record lock poisoned") @@ -473,7 +599,16 @@ impl KvStageIntegration { } pub(crate) fn payload_is_exact_state(&self) -> bool { - self.payload.is_exact_state() + self.exact_state_payload().is_some() + } + + pub(crate) fn exact_state_payload(&self) -> Option { + self.payload + .is_exact_state() + .then_some(self.payload) + .or(self + .durable_payload + .filter(|payload| payload.is_exact_state())) } pub fn should_lookup(&self) -> bool { @@ -529,22 +664,48 @@ impl KvStageIntegration { pub(crate) fn has_exact_state_record_capacity(&self) -> bool { self.exact_state_record_worker_healthy .load(std::sync::atomic::Ordering::Acquire) - && has_exact_state_record_capacity(&self.exact_state_records_pending) + && has_exact_state_record_capacity( + &self.exact_state_records_pending, + &self.exact_state_record_queue_bytes, + EXACT_STATE_RECORD_QUEUE_BYTES, + ) + } + + /// Payload bytes waiting for the worker: the status contract's + /// `write_queue_bytes`. + pub fn exact_state_record_queue_bytes(&self) -> u64 { + self.exact_state_record_queue_bytes + .load(std::sync::atomic::Ordering::Acquire) + } + + /// The durable tier, when one is open. + pub fn l3(&self) -> Option<&Arc> { + self.l3.as_ref() } pub(crate) fn enqueue_exact_state_record( &self, pending: PendingExactStateRecord, ) -> ExactStateRecordAdmission { - let admission = enqueue_exact_state_record( - &self.exact_state_record_tx, - &self.inflight_records, - &self.exact_state_records_queued, - &self.exact_state_records_dropped, - &self.exact_state_records_pending, - &self.exact_state_record_worker_healthy, - pending, - ); + let admission = self.exact_state_record_worker.with_sender(|sender| { + let Some(sender) = sender else { + self.finish_record(&pending.page_id); + self.exact_state_records_dropped + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + return ExactStateRecordAdmission::WorkerStopped; + }; + enqueue_exact_state_record( + sender, + &self.inflight_records, + &self.exact_state_records_queued, + &self.exact_state_records_dropped, + &self.exact_state_records_pending, + &self.exact_state_record_queue_bytes, + EXACT_STATE_RECORD_QUEUE_BYTES, + &self.exact_state_record_worker_healthy, + pending, + ) + }); if matches!( admission, ExactStateRecordAdmission::DroppedFull | ExactStateRecordAdmission::WorkerStopped @@ -958,8 +1119,9 @@ mod exact_state_record_queue_tests { use super::{ BTreeSet, EXACT_STATE_RECORD_CAPACITY, ExactStateExtra, ExactStateRecordAdmission, - ExactStateWorkerHandles, KvLifecycleEvent, KvLifecycleObserver, PendingExactStateRecord, - enqueue_exact_state_record, has_exact_state_record_capacity, run_exact_state_record_job, + ExactStateRecordWorker, ExactStateWorkerHandles, KvLifecycleEvent, KvLifecycleObserver, + PendingExactStateRecord, enqueue_exact_state_record, has_exact_state_record_capacity, + run_exact_state_record_job, }; struct RecordingObserver(Arc>>); @@ -977,16 +1139,173 @@ mod exact_state_record_queue_tests { extra: ExactStateExtra::default(), namespace: "test".to_string(), token_ids: vec![1], + l3_fill_claim: None, } } + fn pending_with_bytes(page_id: &str, bytes: usize) -> PendingExactStateRecord { + PendingExactStateRecord { + payload: ExactStatePayload::full_state(vec![1; bytes]), + ..pending(page_id) + } + } + + const CAP: u64 = 1024; + + #[test] + fn final_worker_owner_drains_queued_records_before_drop_returns() { + let (sender, receiver) = sync_channel(2); + let completed = Arc::new(AtomicUsize::new(0)); + let worker_completed = completed.clone(); + let task = std::thread::spawn(move || { + while receiver.recv().is_ok() { + worker_completed.fetch_add(1, Ordering::Release); + } + }); + let worker = Arc::new(ExactStateRecordWorker::new(sender, task)); + worker.with_sender(|sender| sender.unwrap().send(pending("latest")).unwrap()); + + drop(worker); + + assert_eq!(completed.load(Ordering::Acquire), 1); + } + #[test] fn pending_capacity_signal_rejects_work_before_export() { let pending_count = AtomicUsize::new(0); - assert!(has_exact_state_record_capacity(&pending_count)); + let queue_bytes = AtomicU64::new(0); + assert!(has_exact_state_record_capacity( + &pending_count, + &queue_bytes, + CAP + )); pending_count.store(EXACT_STATE_RECORD_CAPACITY, Ordering::Release); - assert!(!has_exact_state_record_capacity(&pending_count)); + assert!(!has_exact_state_record_capacity( + &pending_count, + &queue_bytes, + CAP + )); + + pending_count.store(0, Ordering::Release); + queue_bytes.store(CAP, Ordering::Release); + assert!( + !has_exact_state_record_capacity(&pending_count, &queue_bytes, CAP), + "a byte-full queue must refuse before the export is paid for" + ); + } + + #[test] + fn queue_is_bounded_in_bytes_not_entries() { + let (sender, _receiver) = sync_channel(EXACT_STATE_RECORD_CAPACITY); + let inflight = Mutex::new(BTreeSet::from(["a".to_string(), "b".to_string()])); + let queued = AtomicU64::new(0); + let dropped = AtomicU64::new(0); + let pending_count = AtomicUsize::new(0); + let queue_bytes = AtomicU64::new(0); + let worker_healthy = AtomicBool::new(true); + + assert_eq!( + enqueue_exact_state_record( + &sender, + &inflight, + &queued, + &dropped, + &pending_count, + &queue_bytes, + CAP, + &worker_healthy, + pending_with_bytes("a", 700), + ), + ExactStateRecordAdmission::Queued + ); + assert_eq!(queue_bytes.load(Ordering::Relaxed), 700); + + // Plenty of entry slots left; the bytes are what is full. + assert_eq!( + enqueue_exact_state_record( + &sender, + &inflight, + &queued, + &dropped, + &pending_count, + &queue_bytes, + CAP, + &worker_healthy, + pending_with_bytes("b", 700), + ), + ExactStateRecordAdmission::DroppedFull + ); + assert_eq!( + queue_bytes.load(Ordering::Relaxed), + 700, + "a dropped record left bytes claimed" + ); + assert!(!inflight.lock().unwrap().contains("b")); + assert_eq!(dropped.load(Ordering::Relaxed), 1); + assert_eq!(pending_count.load(Ordering::Relaxed), 1); + } + + #[test] + fn one_oversized_record_is_admitted_into_an_empty_queue() { + let (sender, _receiver) = sync_channel(EXACT_STATE_RECORD_CAPACITY); + let inflight = Mutex::new(BTreeSet::from(["huge".to_string()])); + let queued = AtomicU64::new(0); + let dropped = AtomicU64::new(0); + let pending_count = AtomicUsize::new(0); + let queue_bytes = AtomicU64::new(0); + let worker_healthy = AtomicBool::new(true); + + assert_eq!( + enqueue_exact_state_record( + &sender, + &inflight, + &queued, + &dropped, + &pending_count, + &queue_bytes, + CAP, + &worker_healthy, + pending_with_bytes("huge", 4096), + ), + ExactStateRecordAdmission::Queued, + "a large model could never record if one export over the cap were refused" + ); + assert_eq!(queue_bytes.load(Ordering::Relaxed), 4096); + assert!(!has_exact_state_record_capacity( + &pending_count, + &queue_bytes, + CAP + )); + } + + #[test] + fn finishing_a_record_releases_its_bytes() { + let inflight = Mutex::new(BTreeSet::from(["done".to_string()])); + let dropped = AtomicU64::new(0); + let pending_count = AtomicUsize::new(1); + let queue_bytes = AtomicU64::new(300); + let worker_healthy = AtomicBool::new(true); + let worker_panics = AtomicU64::new(0); + + run_exact_state_record_job( + ExactStateWorkerHandles { + inflight_records: &inflight, + dropped: &dropped, + pending_count: &pending_count, + queue_bytes: &queue_bytes, + worker_healthy: &worker_healthy, + worker_panics: &worker_panics, + }, + None, + pending_with_bytes("done", 300), + |_| Ok(()), + ); + + assert_eq!(queue_bytes.load(Ordering::Relaxed), 0); + assert_eq!(pending_count.load(Ordering::Relaxed), 0); + assert!(inflight.lock().unwrap().is_empty()); + assert_eq!(dropped.load(Ordering::Relaxed), 0); } #[test] @@ -997,6 +1316,7 @@ mod exact_state_record_queue_tests { let queued = AtomicU64::new(0); let dropped = AtomicU64::new(0); let pending_count = AtomicUsize::new(0); + let queue_bytes = AtomicU64::new(0); let worker_healthy = AtomicBool::new(true); assert_eq!( @@ -1006,6 +1326,8 @@ mod exact_state_record_queue_tests { &queued, &dropped, &pending_count, + &queue_bytes, + CAP, &worker_healthy, pending("dropped"), ), @@ -1025,6 +1347,7 @@ mod exact_state_record_queue_tests { let queued = AtomicU64::new(0); let dropped = AtomicU64::new(0); let pending_count = AtomicUsize::new(0); + let queue_bytes = AtomicU64::new(0); let worker_healthy = AtomicBool::new(true); assert_eq!( @@ -1034,6 +1357,8 @@ mod exact_state_record_queue_tests { &queued, &dropped, &pending_count, + &queue_bytes, + CAP, &worker_healthy, pending("orphaned"), ), @@ -1053,6 +1378,7 @@ mod exact_state_record_queue_tests { let dropped = AtomicU64::new(0); let pending_count = Arc::new(AtomicUsize::new(0)); let worker_pending_count = pending_count.clone(); + let queue_bytes = AtomicU64::new(0); let worker_healthy = AtomicBool::new(true); assert_eq!( @@ -1062,6 +1388,8 @@ mod exact_state_record_queue_tests { &queued, &dropped, &pending_count, + &queue_bytes, + CAP, &worker_healthy, pending("page"), ), @@ -1088,6 +1416,7 @@ mod exact_state_record_queue_tests { let inflight = Mutex::new(BTreeSet::from(["written".to_string()])); let dropped = AtomicU64::new(0); let pending_count = AtomicUsize::new(1); + let queue_bytes = AtomicU64::new(1); let worker_healthy = AtomicBool::new(true); let worker_panics = AtomicU64::new(0); let events: Arc>> = Arc::default(); @@ -1098,6 +1427,7 @@ mod exact_state_record_queue_tests { inflight_records: &inflight, dropped: &dropped, pending_count: &pending_count, + queue_bytes: &queue_bytes, worker_healthy: &worker_healthy, worker_panics: &worker_panics, }, @@ -1118,6 +1448,7 @@ mod exact_state_record_queue_tests { let inflight = Mutex::new(BTreeSet::from(["broken".to_string()])); let dropped = AtomicU64::new(0); let pending_count = AtomicUsize::new(1); + let queue_bytes = AtomicU64::new(1); let worker_healthy = AtomicBool::new(true); let worker_panics = AtomicU64::new(0); let events: Arc>> = Arc::default(); @@ -1128,6 +1459,7 @@ mod exact_state_record_queue_tests { inflight_records: &inflight, dropped: &dropped, pending_count: &pending_count, + queue_bytes: &queue_bytes, worker_healthy: &worker_healthy, worker_panics: &worker_panics, }, @@ -1150,6 +1482,7 @@ mod exact_state_record_queue_tests { let queued = AtomicU64::new(0); let dropped = AtomicU64::new(0); let pending_count = AtomicUsize::new(1); + let queue_bytes = AtomicU64::new(1); let worker_healthy = AtomicBool::new(true); let worker_panics = AtomicU64::new(0); @@ -1158,6 +1491,7 @@ mod exact_state_record_queue_tests { inflight_records: &inflight, dropped: &dropped, pending_count: &pending_count, + queue_bytes: &queue_bytes, worker_healthy: &worker_healthy, worker_panics: &worker_panics, }, @@ -1180,6 +1514,8 @@ mod exact_state_record_queue_tests { &queued, &dropped, &pending_count, + &queue_bytes, + CAP, &worker_healthy, pending("later"), ), diff --git a/crates/skippy-server/src/kv_integration/records.rs b/crates/skippy-server/src/kv_integration/records.rs index 54721d8204..8b8f113ed3 100644 --- a/crates/skippy-server/src/kv_integration/records.rs +++ b/crates/skippy-server/src/kv_integration/records.rs @@ -84,6 +84,14 @@ pub struct ExactStateRestore { pub lookup_ms: f64, pub kv_import_ms: f64, pub recurrent_import_ms: f64, + /// Where the state came from: `"radix"` (RAM) or `"l3"` (disk). + pub source: &'static str, + /// Store-side cost of an L3 fill (read, digest verification, assembly), + /// kept apart from the runtime import above so restore thresholds come + /// from real numbers. Zero for a radix hit. + pub fill_ms: f64, + /// Whether an L3 fill's radix re-warm record was accepted by the worker. + pub rewarm_enqueued: bool, } #[derive(Debug, Clone)] diff --git a/crates/skippy-server/src/runtime_state/lane_lifecycle.rs b/crates/skippy-server/src/runtime_state/lane_lifecycle.rs index 7fa4df3b4d..3f77b221ce 100644 --- a/crates/skippy-server/src/runtime_state/lane_lifecycle.rs +++ b/crates/skippy-server/src/runtime_state/lane_lifecycle.rs @@ -17,6 +17,24 @@ impl RuntimeState { Ok(self.session_stats()) } + pub(crate) fn warmup_generation_graph(&self) -> Result { + if self.model.input_activation_boundary().is_some() + || self.model.output_activation_boundary().is_some() + { + return Ok(false); + } + let token_id = self + .model + .tokenize("", true)? + .into_iter() + .next() + .unwrap_or(0); + let mut session = self.model.create_session()?; + session.decode_step(token_id)?; + session.reset()?; + Ok(true) + } + /// Release the session slot identified by `session_id`. /// /// This is the cleanup path called at the end of every chat diff --git a/crates/skippy-server/src/runtime_state/state_transfer.rs b/crates/skippy-server/src/runtime_state/state_transfer.rs index a62988fdf9..36f66c2b1e 100644 --- a/crates/skippy-server/src/runtime_state/state_transfer.rs +++ b/crates/skippy-server/src/runtime_state/state_transfer.rs @@ -119,6 +119,22 @@ impl RuntimeState { import_result } + /// Finalize a session position after page imports for a model with no + /// recurrent memory. Pure-attention state has no recurrent snapshot, so + /// the tracked token count follows the native position directly. + pub fn set_session_position(&mut self, session_id: &str, token_count: u64) -> Result<()> { + let result = self.session(session_id)?.set_position(token_count); + if result.is_ok() { + record_restored_session_token_count( + &mut self.session_token_counts, + session_id, + token_count, + ); + } + self.notify_import_outcome(&result); + result + } + /// Reports the real result of any runtime-state export call (`export_state`, /// `export_full_state`, `export_recurrent_state`) to the attached observer. /// Never called before the native call returns, so success is only ever diff --git a/crates/skippy-topology/src/lib.rs b/crates/skippy-topology/src/lib.rs index 826a871cb0..26cec2f737 100644 --- a/crates/skippy-topology/src/lib.rs +++ b/crates/skippy-topology/src/lib.rs @@ -3,6 +3,7 @@ use serde::{Deserialize, Serialize}; mod artifact_diagnostics; mod edge_order; mod family_capability; +pub mod phase_placement; mod planning; mod validation; @@ -20,6 +21,7 @@ pub use family_capability::{ reviewed_capability_for_identity, reviewed_capability_records, rwkv6_capability, rwkv7_capability, }; +pub use phase_placement::{HandoffCostModel, PhaseAssignment, PhaseCandidate, assign_phase_roles}; pub use planning::{ classify_layers, plan_contiguous_with_splits, plan_even_contiguous, plan_package_aware_contiguous, plan_package_aware_contiguous_with_signals, diff --git a/crates/skippy-topology/src/phase_placement.rs b/crates/skippy-topology/src/phase_placement.rs new file mode 100644 index 0000000000..39dd55c33e --- /dev/null +++ b/crates/skippy-topology/src/phase_placement.rs @@ -0,0 +1,249 @@ +//! Cost-based prefill/decode phase placement. +//! +//! Prefill is compute-bound and decode is memory-bandwidth-bound, so on a +//! heterogeneous pair the compute-strong node should prefill and the +//! bandwidth-strong node should decode — but only when the handoff pays: +//! moving continuation state costs a fixed floor (recurrent/SSM snapshot for +//! hybrid families) plus per-token attention KV, so short prompts and poor +//! links must prefill in place. These are pure functions over capability +//! signals nodes already gossip (`compute_tflops_fp16`, +//! `mem_bandwidth_gbps` in `PeerAnnouncement`); the planner and the router +//! call them, they do not read the mesh themselves. + +/// Capability signals for one candidate node, as gossiped. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct PhaseCandidate { + pub compute_tflops_fp16: f64, + pub mem_bandwidth_gbps: f64, +} + +/// A phase role assignment for a pair of nodes. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PhaseAssignment { + /// Index of the node that should prefill. + pub prefill: usize, + /// Index of the node that should decode. + pub decode: usize, +} + +/// Assign prefill/decode roles across a candidate pair. +/// +/// The binding constraint decides: place the scarce resource where it is the +/// bottleneck. When the pair differs meaningfully in compute but not in +/// bandwidth (the M3 Ultra / M1 Ultra shape: ~2x compute apart, ~2% +/// bandwidth apart), the compute-strong node prefills and decode loses +/// almost nothing. When the pair differs in both, the ratio test below +/// picks the split that maximises the product of (prefill compute) and +/// (decode bandwidth) — equivalent to comparing the two assignments' +/// bottleneck utilisation. +pub fn assign_phase_roles(a: PhaseCandidate, b: PhaseCandidate) -> PhaseAssignment { + // Compare: a prefills (a.compute * b.bandwidth) vs b prefills + // (b.compute * a.bandwidth). Guard degenerate zero signals by treating + // them as equal, which falls through to a-prefills for determinism. + let a_prefills = a.compute_tflops_fp16.max(0.0) * b.mem_bandwidth_gbps.max(0.0); + let b_prefills = b.compute_tflops_fp16.max(0.0) * a.mem_bandwidth_gbps.max(0.0); + if b_prefills > a_prefills { + PhaseAssignment { + prefill: 1, + decode: 0, + } + } else { + PhaseAssignment { + prefill: 0, + decode: 1, + } + } +} + +/// Inputs to the per-request disaggregation cost gate. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct HandoffCostModel { + /// Attention-KV bytes per prompt token (measured, e.g. from a + /// remote-handoff report's `state_bytes_per_prompt_token` minus the + /// fixed floor). + pub state_bytes_per_token: f64, + /// Fixed handoff bytes independent of prompt length — the + /// recurrent/SSM + conv snapshot for hybrid families, 0 for pure + /// attention. + pub fixed_state_bytes: f64, + /// Usable link throughput between the pair, bytes/second (measured, + /// not line rate). + pub link_bytes_per_second: f64, + /// Prefill throughput on the decode node, tokens/second. + pub local_prefill_tokens_per_second: f64, + /// Prefill throughput on the prefill node, tokens/second. + pub remote_prefill_tokens_per_second: f64, + /// Fraction of transfer hidden behind prefill compute by page + /// streaming, in [0, 1]. 0 models a flat (non-overlapped) transfer; + /// measured runs on the lab pair should calibrate this. + pub transfer_overlap_fraction: f64, +} + +impl HandoffCostModel { + /// TTFT cost (seconds) of disaggregating a prompt of `tokens`. + pub fn disaggregated_seconds(&self, tokens: u64) -> f64 { + let tokens = tokens as f64; + let remote_rate = self.remote_prefill_tokens_per_second.max(f64::EPSILON); + let link = self.link_bytes_per_second.max(f64::EPSILON); + let prefill = tokens / remote_rate; + let transfer_bytes = tokens * self.state_bytes_per_token.max(0.0); + let hidden = self.transfer_overlap_fraction.clamp(0.0, 1.0); + // Per-token KV can hide behind prefill; the fixed snapshot is only + // final after the last chunk and is always exposed. + let exposed_transfer = + (transfer_bytes * (1.0 - hidden) + self.fixed_state_bytes.max(0.0)) / link; + prefill + exposed_transfer + } + + /// TTFT cost (seconds) of prefilling in place on the decode node. + pub fn local_seconds(&self, tokens: u64) -> f64 { + tokens as f64 / self.local_prefill_tokens_per_second.max(f64::EPSILON) + } + + /// Whether a prompt of `tokens` should hand off. + pub fn should_disaggregate(&self, tokens: u64) -> bool { + self.disaggregated_seconds(tokens) < self.local_seconds(tokens) + } + + /// Smallest prompt length at which handoff wins, or `None` if it never + /// wins (searched up to `max_tokens`). The router caches this per pair + /// and compares incoming prompt lengths against it. + pub fn break_even_tokens(&self, max_tokens: u64) -> Option { + // Both cost curves are affine in `tokens`, so the crossover is + // where the per-token slopes and the fixed offsets balance: + // tokens/local = tokens/remote + tokens*exposed_per_token/link + fixed/link + // Solve directly, then round outward and verify against the model + // to stay robust to degenerate inputs. + let local_rate = self.local_prefill_tokens_per_second.max(f64::EPSILON); + let remote_rate = self.remote_prefill_tokens_per_second.max(f64::EPSILON); + let link = self.link_bytes_per_second.max(f64::EPSILON); + let hidden = self.transfer_overlap_fraction.clamp(0.0, 1.0); + let slope_local = 1.0 / local_rate; + let slope_disaggregated = + 1.0 / remote_rate + self.state_bytes_per_token.max(0.0) * (1.0 - hidden) / link; + if slope_disaggregated >= slope_local { + // Handoff never catches up: the per-token cost alone is worse. + return None; + } + let fixed = self.fixed_state_bytes.max(0.0) / link; + let crossover = fixed / (slope_local - slope_disaggregated); + let candidate = crossover.ceil().max(1.0) as u64; + (candidate <= max_tokens && self.should_disaggregate(candidate)).then_some(candidate) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The lab shape: M3 Ultra ≈ 2x the compute of the M1 Ultra with + /// near-equal bandwidth. The compute-strong box must prefill. + #[test] + fn near_equal_bandwidth_pairs_split_on_compute() { + let m3_ultra = PhaseCandidate { + compute_tflops_fp16: 57.0, + mem_bandwidth_gbps: 819.0, + }; + let m1_ultra = PhaseCandidate { + compute_tflops_fp16: 42.0, + mem_bandwidth_gbps: 800.0, + }; + assert_eq!( + assign_phase_roles(m3_ultra, m1_ultra), + PhaseAssignment { + prefill: 0, + decode: 1 + } + ); + assert_eq!( + assign_phase_roles(m1_ultra, m3_ultra), + PhaseAssignment { + prefill: 1, + decode: 0 + } + ); + } + + /// The DGX-Spark shape: huge compute, modest bandwidth, paired with a + /// bandwidth-rich Mac. Compute prefills, bandwidth decodes. + #[test] + fn compute_heavy_bandwidth_light_node_prefills() { + let spark = PhaseCandidate { + compute_tflops_fp16: 250.0, + mem_bandwidth_gbps: 273.0, + }; + let mac = PhaseCandidate { + compute_tflops_fp16: 57.0, + mem_bandwidth_gbps: 819.0, + }; + assert_eq!( + assign_phase_roles(spark, mac), + PhaseAssignment { + prefill: 0, + decode: 1 + } + ); + } + + fn lab_model() -> HandoffCostModel { + // Calibrated to the remote-handoff report shape: ~115 KiB/token + // attention KV, no recurrent floor (dense), 1 GB/s usable link, + // decode node prefills at 2k tok/s, prefill node at 4k tok/s. + HandoffCostModel { + state_bytes_per_token: 115.0 * 1024.0, + fixed_state_bytes: 0.0, + link_bytes_per_second: 1e9, + local_prefill_tokens_per_second: 2000.0, + remote_prefill_tokens_per_second: 4000.0, + transfer_overlap_fraction: 0.9, + } + } + + #[test] + fn fast_link_with_overlap_favors_handoff_beyond_break_even() { + let model = lab_model(); + let break_even = model.break_even_tokens(65_536); + // No fixed floor and a winning slope: handoff wins from the start. + assert_eq!(break_even, Some(1)); + assert!(model.should_disaggregate(4096)); + } + + /// A hybrid family's fixed recurrent floor (~160 MiB) pushes the break + /// even point out: short prompts must prefill in place. + #[test] + fn fixed_recurrent_floor_rejects_short_prompts() { + let model = HandoffCostModel { + fixed_state_bytes: 160.0 * 1024.0 * 1024.0, + ..lab_model() + }; + assert!(!model.should_disaggregate(256)); + let break_even = model + .break_even_tokens(65_536) + .expect("must eventually win"); + assert!(break_even > 256, "break even was {break_even}"); + assert!(model.should_disaggregate(break_even)); + assert!(!model.should_disaggregate(break_even.saturating_sub(64))); + } + + /// A slow link with no overlap can make handoff lose at every length. + #[test] + fn slow_flat_link_never_disaggregates() { + let model = HandoffCostModel { + link_bytes_per_second: 50e6, + transfer_overlap_fraction: 0.0, + ..lab_model() + }; + assert_eq!(model.break_even_tokens(1_000_000), None); + assert!(!model.should_disaggregate(32_768)); + } + + /// No compute advantage and a real transfer cost: prefill in place. + #[test] + fn equal_nodes_prefill_in_place() { + let model = HandoffCostModel { + remote_prefill_tokens_per_second: 2000.0, + ..lab_model() + }; + assert_eq!(model.break_even_tokens(1_000_000), None); + } +} diff --git a/docs/CONFIGURATION_PR8_CLOSEOUT_AUDIT.md b/docs/CONFIGURATION_PR8_CLOSEOUT_AUDIT.md index 5a37b93d37..99727ab065 100644 --- a/docs/CONFIGURATION_PR8_CLOSEOUT_AUDIT.md +++ b/docs/CONFIGURATION_PR8_CLOSEOUT_AUDIT.md @@ -55,7 +55,9 @@ schema row or stale manifest row from passing review: `runtime.activity.advertisement`, `runtime.reconcile_model_targets`, `runtime.reconcile_model_target_demand_upgrades`, `runtime.native_runtime.mesh_version`, `runtime.native_runtime.skippy_abi`, -`runtime.native_runtime.selection`, `runtime.model_target_demand_upgrade_min_requests`, +`runtime.native_runtime.selection`, `runtime.kv_cache.disk.mode`, +`runtime.kv_cache.disk.directory`, `runtime.kv_cache.disk.budget_mib`, +`runtime.kv_cache.disk.minimum_free_mib`, `runtime.model_target_demand_upgrade_min_requests`, `runtime.model_target_demand_upgrade_max_age_secs`, `advanced.server.alias`, `model`, `hardware.model_path`, `hardware.hf_repo`, `hardware.hf_file`, `model_fit.ctx_size`, `model_fit.batch`, `model_fit.ubatch`, diff --git a/docs/README.md b/docs/README.md index 238750fe06..0f02077dd6 100644 --- a/docs/README.md +++ b/docs/README.md @@ -30,6 +30,7 @@ Use this hub to find project guides that are not owned by a single Rust crate. | [skippy/TOPOLOGY_PLANNER.md](skippy/TOPOLOGY_PLANNER.md) | Stage topology planning behavior | | [skippy/CONFIGURATION.md](skippy/CONFIGURATION.md) | Authoritative operator matrix for Skippy config keys and rejection boundaries | | [skippy/PROMPT_CACHE.md](skippy/PROMPT_CACHE.md) | OpenAI prompt-prefix cache behavior, defaults, telemetry, and benchmark flow | +| [skippy/KV_CACHE_DISK.md](skippy/KV_CACHE_DISK.md) | Operator guide for the node-local disk prompt cache (L3): config, modes, status, prune/clear, corruption handling | | [skippy/PIPELINED_VERIFY_WINDOW.md](skippy/PIPELINED_VERIFY_WINDOW.md) | Native MTP, anchored N-gram extension, VerifyWindow protocol, pipeline behavior, and telemetry | | [skippy/SUFFIX_NGRAM_PROPOSER.md](skippy/SUFFIX_NGRAM_PROPOSER.md) | Long exact-suffix proposer design, invariants, telemetry, and benchmark contract | | [skippy/DATA_FLOW.md](skippy/DATA_FLOW.md) | Stage data flow and transport details | diff --git a/docs/skippy/CACHEGEN_BACKEND_PLAN.md b/docs/skippy/CACHEGEN_BACKEND_PLAN.md new file mode 100644 index 0000000000..54880cfdd8 --- /dev/null +++ b/docs/skippy/CACHEGEN_BACKEND_PLAN.md @@ -0,0 +1,370 @@ +# CacheGen Backend Qualification (#1652) + +Status: **LMCache-compatible direct Metal restore passes the local F32/F32 and +F32+F16 gates; F16 and the lower-width sampled types remain stopped on local +latency, and CacheGen remains opt-in pending selection-path wiring**. Owner: +jian yang. +Reviewed against: #1652 scope, scama's directives of 2026-09-10 (v4 +contract, CPU+Metal parity, six measurements, stop rule). + +## Where each backend stands + +| Backend | Kernel | Status | Evidence | +|---|---|---|---| +| CPU (reference) | scalar Rust | **Correctness reference only** — portable F32, F16, Q8_0, Q4_0, and mixed K/V adapters are implemented | Python/Rust fixtures pin LMCache revision `b5d109e`; typed archive fixtures cover every current user-selectable runtime K/V type | +| Metal (Apple GPU) | native MSL | **All typed restores implemented; F32/F32 and F32+F16 clear the local gate** — F16, F32, Q8_0, and Q4_0 match native fixture layouts; Q8_0/F32 straddles the latency boundary across repeats, while F16 and the lower-width sampled cases remain stopped locally | Apple M1 Ultra device fixture and typed 19K gates below | +| CUDA (NVIDIA) | shared native CUDA/HIP source | **Typed kernels landed; F16 fixture qualified** — the real RTX 5080 fixture matches F16 bytes, while F32, quantized, and matched typed 19K runtime gates remain | Real NVIDIA F16 fixture; remaining typed hardware gates pending | +| HIP/ROCm (AMD) | shared native CUDA/HIP source | **Typed kernels landed; compile/package qualified** — no real AMD runtime claim yet | Real AMD fixture and typed 19K gates remain | + +Nothing may be marked implemented until it runs on real hardware and +matches the CPU reference bit-for-bit. Compile-only checks prove the +kernel lowers; they say nothing about the hardware. + +Native Metal fixture builds must set `GGML_CCACHE=OFF`. The embedded source is +included through a generated assembly `.incbin`; an assembly cache can otherwise +reuse an object after the Metal source changes and run a stale kernel. + +## The six spike measurements (2026-09-11 re-run, M2 Max, 4096x128 tile, +## 524,288 f16 values = 1,048,576 raw bytes, release build, synchronized +## stages, bitwise equality, transfers timed) + +Every timed stage ends with `client.sync()` inside the timer: the +numbers cover real completion, not launch enqueue. H2D and D2H are +measured to completion as well. Cold JIT is the first synchronized +launch of the kernel specialization in the process (cubecl 0.10.0 has +no on-disk kernel cache on this path — the only one, SPIR-V, is +Vulkan-only and not enabled — so per-process first launch is the true +cold path for both backends). Equality is exact `==` on symbols and f32 +bits, with mismatch counts printed. + +| Metric | cubecl-cpu | wgpu (Metal) | +|---|---|---| +| Cold JIT/compile (quantize+delta; undelta+dequantize) | ~38-42 ms / ~13 ms | ~8-9 ms / ~4 ms | +| Warm dispatch, synchronized (avg per launch) | ~570-640 us | ~3.3-4.6 ms | +| H2D bytes (f32 tile + 8 B calibration), timed | 2,097,160 in ~12-18 ms | 2,097,160 in ~0.9-1.4 ms | +| D2H bytes (u32 symbols + f32 rebuilt), timed | 4,194,304 in ~36-44 us | 4,194,304 in ~3.0-4.1 ms | +| Live device buffer peak (tile + calibration + both outputs) | 6,291,464 | 6,291,712 | +| Encoded-size ratio (rANS over device symbols / raw) | **0.075 (13.3x)** | **0.075 (13.3x)** | +| Output equality vs CPU reference (bitwise) | symbols + values exact, 0/524,288 mismatches | symbols + values exact, 0/524,288 mismatches | + +The live peak is what the harness actually holds while kernels run +(2,097,152-byte tile + 8-byte calibration + 2,097,152-byte symbol +buffer + 2,097,152-byte rebuilt buffer); Metal's allocator rounds its +copies slightly differently, hence the 248-byte difference. Warm-dispatch +convergence was checked across iteration counts (5/50/200/500). These +numbers replace both the 2026-09-10 enqueue-only measurements (5/4 us) +and the 2026-09-11 first re-run's output-only peak: the completion wait +and the input buffers are now inside the reported figures. Correctness +claims are unchanged: both backends were already symbol-exact, and the +values are proven bit-equal rather than within 1e-6. + +Caveats, stated rather than buried: warm dispatch at 4096x128 is now +dominated by the synchronization round-trip plus launch overhead, not +bandwidth; the per-column scan is the correct bit-exact baseline, not +the fastest shape (a parallel scan is the follow-up). The 2x H2D cost +versus raw f16 bytes exists because the spike uploads f32; shipping f16 +halves it and is a trivial follow-up. + +## Buffer interop and the stop rule + +The store's exported segments are plain byte buffers. CubeCL's runtime +consumes host buffers via `create_from_slice` and returns via +`read_one`; no zero-copy path into the store's packed segment files +exists today, so the spike measured the honest version: one H2D upload, +one D2H return. On this tile that is 2 MB up / 4 MB down against a +0.075-ratio encoded payload — the copy cost is real and is the thing the +later quality/performance gate must beat on the ~19K acceptance +workload. Per the agreed stop rule this evidence comes back before any +expansion: CubeCL is **not** a committed dependency, the spike lives +behind the `cachegen-spike` feature, and nothing in the library links it. + +## Native exact control + +The exact control arm uses `native-kv-page/1` per-segment identity. KV bytes +exported by the active runtime are written and restored verbatim, including +F32, F16, Q8_0, and Q4_0 layouts supported by that runtime; there is no storage +transcode. Mixed KV plus recurrent payloads cut at the representation boundary, +so auxiliary continuation state remains exact `raw/1`. Runtime page metadata is +validated before segment reads, while the existing exact-state identity binds +the runtime ABI, platform, model, layer range, and KV configuration. This is the +baseline every CacheGen result must beat end to end. + +## Capability failure policy + +A backend that cannot run a codec fails explicitly through the v4 +per-segment identity gate (`SegmentCodecIdentity::is_supported` / +negotiation naming the segment index). There is no hidden fallback: an +unsupported `cachegen/1` segment is a clean miss with a named reason, +never a silent decode on another backend or a raw reinterpretation. +Lossy entries additionally carry `calibration_digest`; a lookup matches +only identically-calibrated entries and can never satisfy an exact +lookup. + +## LMCache-compatible reference port + +The active gate now uses a Rust port of LMCache revision +`b5d109ea99a89b4d8a670ee4fc2e8cb76411ee5c`. The source records the exact +upstream Python and CUDA files and retains Apache-2.0 attribution. It reproduces +the parts that determine reconstruction quality and wire size: + +- per-token maximum-magnitude scaling across channels; +- LMCache's generic 32/16-bin K and V layer schedule; +- the normalized 33-entry CDF for every layer/channel pair; +- the CUDA implementation's 32-bit arithmetic coder and 256-token chunk limit; +- token-major K/V handling, including Skippy's transposed-V page layout. + +Skippy uses a bounded portable envelope instead of LMCache's Python pickle +container. Fixtures generated by the independent Python scalar transcription +pin both 16-bin and 32-bin streams, and normal Rust tests require byte-for-byte +encoder agreement and decoder agreement with those fixtures. + +This remains an opt-in correctness path and is not selected by storage or the +request path. The matched 19K result below proves the reference recovers +continuation quality, while also proving that scalar arithmetic decode cannot +meet the restore-to-first-token gate. + +## Sequencing after the direct Metal gate + +1. Preserve the scalar implementation and its pinned fixtures as the + deterministic oracle for device kernels. +2. Keep the completed Skippy-owned compressed-page transaction and optional + backend-registry hook as the integration boundary. Metal and CUDA/HIP decode + validated records directly into allocated resident cells; capability or + execution failure rolls back without a scalar fallback. +3. Qualify the shared CUDA/HIP decoder against the scalar fixtures and matched + 19K gate on real NVIDIA and AMD hardware. +4. Add independently decodable substreams or an equivalent parallel entropy + layout before retrying local-tier Metal promotion. Preserve the current + scalar stream as the compatibility oracle for the new revision. +5. Add encode from resident K/V storage and copy only the compact archive back + to the persistence layer. +6. Run separate quality and latency gates for Q8_0/Q8_0, Q4_0/Q4_0, and the + supported mixed K/V combinations now that every typed record is wired into + the native backend contract. + +## Typed portable record boundary + +The portable page archive uses one unchanged LMCache-compatible F16 entropy +segment. Record kinds describe the runtime edge adapter: F32 and F16 scalar +rows, Q8_0 and Q4_0 block rows, plus F32/F16 transposed V. Encode converts each +native row into the shared F16 stream; decode reconstructs that stream and +packs the selected native row type. K and V carry independent kinds, so mixed +selections do not require a second container or codec revision. + +The pure-Rust fixtures cover F32/F32, F16/F16, Q8_0/Q8_0, Q4_0/Q4_0, +Q8_0/Q4_0, and transposed F32 V. Metal, CUDA, and ROCm share a typed native +decode contract for direct F16 and F32 resident writes, including transposed V, +and fused Q8_0/Q4_0 block quantization into row-major resident storage. The +former F16-only backend symbol was removed rather than retained as a +compatibility alias. The Metal device fixture compares both quantized outputs +byte-for-byte with ggml's native quantizer. No scalar restore fallback is +permitted. Quantized destinations remain unqualified until matched end-to-end +quality gates measure the combined CacheGen and native repacking loss. + +## Native runtime integration boundary + +The existing native page API is host-buffer oriented. Rust passes a `&[u8]` +to `skippy_import_kv_page`; the C ABI receives a `const void *`; and +`llama_kv_cache::stage_import_kv_page` allocates cells before copying each run +with `ggml_backend_tensor_set`. A device decoder above this API would have to +materialize the complete decoded page in host memory and upload it again. That +would erase the main benefit on discrete CUDA and ROCm devices. + +CacheGen therefore belongs inside the Skippy state-transfer transaction while +its portable envelope and scalar oracle remain in `skippy-cache`. The public +runtime accepts a compressed portable page. The native implementation validates +all records and destination coverage, allocates all target cell runs, resolves +the codec hook for every owning backend, dispatches decode into the resident K +and V tensors, synchronizes, and only then commits the session position. A +validation, capability, launch, or synchronization failure restores the prior +cell state. Unsupported backends return an explicit unsupported result; the +performance gate cannot silently select the scalar oracle. + +The hook is an optional function obtained through +`ggml_backend_reg_get_proc_address`, following the extension mechanism already +used by Metal backend tuning. CacheGen state transfer is intentionally not a +new global ggml graph operation: it runs outside model execution, and making it +an op would also require global enum, scheduler, graph-identity, slice-planning, +shape, and backend-support changes. The registry hook keeps the patch local to +the state capability and the backends that implement it while still receiving +the active backend context needed for ordered execution. + +CUDA and ROCm share the `ggml-cuda` source path, which llama.cpp already builds +through CUDA or HIP. Metal implements the same contract in MSL. Each launch +batches many 256-token arithmetic streams: one thread serially decodes at most +256 symbols for one channel while hundreds of thousands of independent channel +streams run in parallel. Destination metadata maps each stream to a K/V tensor, +allocated cell run, row stride, and optional transposed-V stride. This avoids a +launch per tile and permits dequantization and final layout writes in one pass. + +## 19K direct Metal result (2026-09-11): QUALITY PASS, LOCAL LATENCY STOP + +The direct-device gate was run from exact commit +`c86cf848b0fa183f2c2e594fe1e51bf1024d7185` on an Apple M1 Ultra (128 GiB, +Metal) with the same pinned Qwen3 0.6B Q8_0 model and 19,000-token workload as +the scalar result below. Native and CacheGen continuations run in isolated +sessions. The scalar decoder still runs as an independently timed correctness +oracle, but its 41.44 seconds are excluded from the direct-device TTFT. The +compact result is +[`cachegen-metal-device-qwen3-0.6b-19k-summary.json`](cachegen-metal-device-qwen3-0.6b-19k-summary.json). + +| Metric | Native | Direct Metal CacheGen | Decision | +|---|---:|---:|---| +| Persisted bytes | 2,179,072,000 | 446,903,003 | 20.51% of native (4.88x smaller) | +| Persist path | 424.19 ms | 21,332.42 ms including encode | Fail for synchronous persistence | +| Read | 296.25 ms | 54.21 ms | CacheGen saves 242.04 ms | +| Resident import | 56.69 ms | 540.74 ms | CacheGen spends 484.05 ms more reconstructing K/V | +| Restore to first token | 385.41 ms | 604.84 ms | Fail (1.57x slower) | +| Scalar oracle decode | — | 41,436.07 ms, excluded from device TTFT | Correctness-only reference | +| Continuation throughput | 122.34 tok/s | 128.07 tok/s | No steady-state regression | +| p99 decode | 32.47 ms | 9.88 ms | Within the 5% regression budget | +| Greedy-token agreement | 64/64 control | 64/64 (100%) | Pass versus 95% gate | +| Estimated codec working bytes | — | 2,627,965,638 | Reported; no memory cap was supplied | + +The Metal decoder is correct after aligning every staged tile before typed +metadata reads; the regression test covers two differently sized consecutive +tiles. Replacing the 64-bit arithmetic division with a float reciprocal estimate +and exact integer correction, plus a binary CDF search, reduced the same 19K +resident import from 744.95 ms to 540.74 ms. It still loses the local gate +because every restore reconstructs about 1.09 billion F16 values through +per-channel arithmetic streams of up to 256 symbols. On unified-memory M1 Ultra, +the native 2.18 GB copy is unusually fast: CacheGen's 242.04 ms read saving does +not recover its 484.05 ms reconstruction penalty. The result does not decide a +remote tier, where transfer time and pipelining differ; that tier needs its own +matched end-to-end gate under #1427. + +## 19K typed Metal matrix (2026-09-12): F32 PASS, OTHER LOCAL PROMOTION STOPS + +The typed gate was run from exact commits +`2e26d46ea87e1b8ee783460998e703f669513f91` (quantized and mixed rows) and +`74b4f60719d09dee7d9579c1365ac6f95d14c20c` (F32/F32) on the same Apple M1 +Ultra, pinned Qwen3 0.6B Q8_0 model, 19,000-token prefix, and 64-step +continuation. The F32 run also caught and fixed a native config defect where +GGML enum value zero was interpreted as an unset cache type and silently +replaced with F16. The native regression now pins both the F16 default and an +explicit F32 request. +The F32/F16 crossover probes were run from exact commit +`ddddf34aa5e64262c9e113d07f9dcb506e5f7aab`. +The gate now accepts independent `--cache-type-k` and `--cache-type-v` values +and records them in its report. The compact matrix is +[`cachegen-metal-typed-qwen3-0.6b-19k-summary.json`](cachegen-metal-typed-qwen3-0.6b-19k-summary.json). + +| K/V type | Native bytes | CacheGen bytes | CacheGen/native | Agreement | Native TTFT | CacheGen TTFT | Decision | +|---|---:|---:|---:|---:|---:|---:|---| +| F32/F32 | 4,358,144,000 | 446,903,003 | 10.25% | 64/64 (100%) | 1,199.13 ms | 766.12 ms | Pass: quality, size, local latency, and p99 | +| F32/F16 | 3,268,608,000 | 446,903,003 | 13.67% | 64/64 (100%) | 878.61 ms | 686.15 ms | Pass: quality, size, local latency, and p99 | +| F16/F32 | 3,268,608,000 | 446,903,003 | 13.67% | 64/64 (100%) | 840.32 ms | 671.93 ms | Pass: quality, size, local latency, and p99 | +| Q8_0/F32 | 2,757,888,000 | 565,441,003 | 20.50% | 64/64 (100%) | 800.57 / 727.17 ms | 796.94 / 792.13 ms | Unstable boundary: one pass, one latency stop | +| Q8_0/Q8_0 | 1,157,632,000 | 589,803,358 | 50.95% | 64/64 (100%) | 333.36 ms | 650.64 ms | Quality and size pass; local latency fails | +| Q4_0/Q4_0 | 612,864,000 | 635,037,607 | 103.62% | 61/64 (95.31%) | 189.84 ms | 699.35 ms | Quality floor passes; size and local latency fail | +| Q8_0/F16 | 1,668,352,000 | 565,441,003 | 33.89% | 64/64 (100%) | 558.28 ms | 792.05 ms | Quality and size pass; local latency fails | +| Q4_0/F16 | 1,395,968,000 | 610,269,862 | 43.72% | 61/64 (95.31%) | 483.98 ms | 807.03 ms | Quality and size pass; local latency fails | + +These are representative runtime-valid layouts, not qualification of every +pairing. F32/F32 and both F32+F16 directions clear the matched local gate because +reading the compact archive saves enough time against native pages of at least +3.27 GB to absorb device reconstruction. At 2.76 GB, Q8_0/F32 is inside run +variance: it won once by 3.63 ms and lost once by 64.96 ms, so it remains stopped. +The nearly format-independent 521-586 ms CacheGen import times, compared with +native import scaling from 80 to 495 ms with page width, identify arithmetic +reconstruction as the next Metal target. F32/Q8_0 is not a valid control on this +model: llama.cpp requires Flash Attention for quantized V, while that mixed +cache pair does not have a compatible Flash Attention kernel. + +An exact-head Metal staging optimization at +`e2bdc935e24cbed3bff767762547098db6028459` removes the second host copy that +previously occurred when each validated job was accumulated in `NSMutableData` +and then copied into a shared `MTLBuffer`. The backend now sizes and validates a +job before writing its payload, tile descriptors, and stream prefixes directly +into the final shared buffers. Q8_0/F32 import fell from 580.45 ms across the two +boundary runs above to 491.92 ms across two post-change runs, a 15.25% reduction. +F16/F16 import fell from 540.74 ms to 451.82 ms, a 16.44% reduction. Both cases +remain stopped on local TTFT: the optimized F16/F16 run took 514.34 ms versus +343.16 ms native, and the optimized Q8_0/F32 repeats took 713.51/701.97 ms versus +571.88/547.78 ms native. The remaining gap is device arithmetic decode rather +than redundant host staging. + +A diagnostic 128-row archive doubled tile count from 4,200 to 8,344, expanded +Q8_0/F32 storage from 20.50% to 30.84% of native, and increased import to +727.74 ms. Smaller independently decoded tiles therefore do not solve this +format's Metal crossover; shared calibration with multiple arithmetic +substreams would require a separate format change. + +The quantized K cases expose a consistent quality split: Q8_0 preserves all 64 +greedy continuation tokens, while Q4_0 first diverges at step 15 and finishes at +61/64, barely above the declared 95% floor. Every completed CacheGen continuation +stays within the p99 decode-regression budget after restore. F16 and the +lower-width sampled cases remain stopped for the local tier; Q4_0/Q4_0 also has +no storage benefit. The remaining runtime-valid mixed permutations retain +fixture-level coverage and need separate 19K runs before any broader typed +qualification claim. + +## 19K LMCache-compatible CPU result (2026-09-11): QUALITY PASS, LATENCY STOP + +The opt-in gate was run from exact commit +`4bf0865be6e7a13dbf7159b0fb6e6c41f29db72c` on an Apple M1 Ultra (128 GiB, +Metal) with the pinned Qwen3 0.6B Q8_0 model +(`sha256:12fae8b8f78f0360b498d04c8db7d33aff29ab7d8080231f93a17c18119e6735`), +a 19,000-token prefix, F16 K/V, 256-row LMCache chunks, and 64 teacher-forced +continuation steps. The compact result is +[`cachegen-lmcache-qwen3-0.6b-19k-summary.json`](cachegen-lmcache-qwen3-0.6b-19k-summary.json). + +| Metric | Native | LMCache-compatible CPU | Decision | +|---|---:|---:|---| +| Persisted bytes | 2,179,072,000 | 446,903,003 | 20.51% of native (4.88x smaller) | +| Persist path | 2,402.91 ms | 21,426.43 ms including encode | Fail | +| Read | 281.06 ms | 58.83 ms | Encoded path wins bytes/read time | +| Decode codec | — | 41,285.73 ms | Fail | +| Restore to first token | 374.38 ms | 41,411.13 ms | Fail (110.61x slower) | +| Continuation throughput | 119.85 tok/s | 124.17 tok/s | No steady-state regression | +| p99 decode | 33.24 ms | 8.59 ms | Within the 5% regression budget | +| Greedy-token agreement | 64/64 control | 64/64 (100%) | Pass versus 95% gate | +| Estimated codec working bytes | — | 2,627,965,638 | Reported; no memory cap was supplied | + +This result resolves the earlier 18.75% agreement as a defect in the +simplified prototype rather than a limitation of CacheGen. The faithful +quantization schedule recovers all 64 continuation tokens. Scalar arithmetic +coding is still far outside the restore budget, so the CPU path remains an +oracle and no request-path wiring is allowed. A parallel device implementation +must pass the same gate before promotion. + + +## Historical 19K simplified-prototype result (2026-09-11): STOP + +The implementation measured here is a Mesh-owned prototype inspired by +CacheGen. It uses per-segment min/max affine 4-bit calibration. The published +CacheGen design instead calibrates per model and applies mixed quantization +across tensor dimensions. This gate therefore rejects the current prototype; +it is not evidence that a paper-faithful CacheGen implementation fails. + +The opt-in gate in `skippy-correctness state-handoff --cachegen-gate` was run +from exact commit `2677ad62e5295f6da2ac72ae7b8c978f87753d11` on an Apple M1 Ultra +(128 GiB, Metal) with the Qwen3 0.6B Q8_0 model +(`sha256:9465e63a22add5354d9bb4b99e90117043c7124007664907259bd16d043bb031`), +a 19,000-token prefix, F16 K/V, 4,096-row codec tiles, and 64 +teacher-forced continuation steps. The compact machine-readable result is +[`cachegen-quality-gate-qwen3-0.6b-19k-summary.json`](cachegen-quality-gate-qwen3-0.6b-19k-summary.json). + +| Metric | Native | CacheGen | Decision | +|---|---:|---:|---| +| Persisted bytes | 2,179,072,000 | 202,373,739 | CacheGen is 9.287% of native (10.77x smaller) | +| Persist path | 1,369.84 ms | 15,508.82 ms including encode | Fail | +| Read | 239.14 ms | 22.29 ms | CacheGen wins bytes/read time | +| Decode codec | — | 23,521.69 ms | Fail | +| Restore to first token | 311.61 ms | 23,604.40 ms | Fail (75.75x slower) | +| Continuation throughput | 120.23 tok/s | 121.25 tok/s | No steady-state regression | +| p99 decode | 18.86 ms | 9.30 ms | Within the 5% regression budget | +| Greedy-token agreement | 64/64 control | 12/64 (18.75%) | Fail versus 95% gate | +| First mismatch | — | step 1 | Fail | +| Estimated codec working bytes | — | 2,412,371,750 | Reported; no memory cap was supplied | + +Writes call `sync_all`; the same-run reads may still be page-cache warm, so the +read figures are not a cold-device bandwidth claim. That limitation cannot +reverse this decision: CacheGen's 23.52-second CPU decode alone is more than +75 times the complete native restore-to-first-token path, and continuation +quality fails independently. + +Per the issue's stop rule, this result ended production work on that simplified +prototype. The LMCache-compatible port above replaces it in the opt-in gate; +its results must be measured separately. Native exact `native-kv-page/1` +remains the selected representation until the new result passes. diff --git a/docs/skippy/CONFIGURATION.md b/docs/skippy/CONFIGURATION.md index 33765b1b1c..a210406aa2 100644 --- a/docs/skippy/CONFIGURATION.md +++ b/docs/skippy/CONFIGURATION.md @@ -69,6 +69,15 @@ matrix. `crates/mesh-llm-config/src/website_docs_parity.rs` separately asserts that every key path documented here also appears in the public website configuration reference, with the same `Wiring status`. +## Node-local disk prompt cache + +| Report section | Report setting name | Config key path | Priority | Owner module | Translation target | Supported modes | Live-apply behavior | Default source | Validation rule | Docs anchor | Test evidence | Notes | Wiring status | +|---|---|---|---|---|---|---|---|---|---|---|---|---|---| +| #1576 | Disk cache mode | `runtime.kv_cache.disk.mode` | P0 | `mesh-llm-config` / host runtime | node-scoped `L3CacheManager` | single-stage, staged | process restart | `off` | enum `off`, `auto`, or `fixed` | `#node-local-disk-prompt-cache` | `mesh-llm-config` schema and validation tests | Turning off never deletes existing cache data | wired | +| #1576 | Disk cache directory | `runtime.kv_cache.disk.directory` | P0 | host runtime | node-scoped cache root | single-stage, staged | process restart | `$MESH_LLM_HOME/kv-cache` | explicit values must be absolute | `#node-local-disk-prompt-cache` | `mesh-llm-config` validation tests | Directory changes do not migrate or delete old data | wired | +| #1576 | Fixed disk budget | `runtime.kv_cache.disk.budget_mib` | P0 | host runtime / `skippy-cache` | manager hard byte budget | single-stage, staged | applies dynamically | unset | required and > 0 only in fixed mode | `#node-local-disk-prompt-cache` | config precedence and manager limit tests | One shared physical cap per node root; never unbounded | wired | +| #1576 | Minimum free storage | `runtime.kv_cache.disk.minimum_free_mib` | P0 | host runtime / `skippy-cache` | manager free-space reserve | single-stage, staged | applies dynamically | `16384` MiB | at least `1024` MiB | `#node-local-disk-prompt-cache` | config precedence and low-space transition tests | Writes decline while fills remain available | wired | + ## Model fit, context, and KV cache diff --git a/docs/skippy/KV_CACHE_DISK.md b/docs/skippy/KV_CACHE_DISK.md new file mode 100644 index 0000000000..e22d2c24db --- /dev/null +++ b/docs/skippy/KV_CACHE_DISK.md @@ -0,0 +1,250 @@ +# KV-Cache Disk Tier — Operator Guide + +The node-local **disk prompt cache** is the durable tier under Skippy's radix +cache: exported continuation (KV) state is cut into content-addressed segments +and committed under manifests so a later run can restore an exact prefix from +disk instead of paying cold prefill. This guide covers operating the tier that +ships today (the exact L1/L3 stack): configuration, defaults, status, restart +behavior, prune/clear, corruption handling, and troubleshooting. + +This is a different cache from the in-memory OpenAI prompt-prefix cache +documented in [PROMPT_CACHE.md](PROMPT_CACHE.md). That one lives in the serving +process; this one is the on-disk L3 store described here. + +**Fail-closed configuration, fail-open runtime.** Invalid disk-cache +*configuration* fails **closed**: config validation and resolution reject bad +values — fixed mode without a budget, a non-IEC or zero size, a relative +directory, or a minimum-free reserve below 1 GiB — by returning an error, so a +node refuses to start on a broken cache setting rather than silently ignoring +it. A *valid* configuration is then **fail-open for inference**: if the store +cannot open, reaches low space, or cannot admit a write, the node logs a warning +and serves with cold prefill. Valid runtime unavailability never blocks +generation. + +## Modes and safe defaults + +The tier has three modes, selected by `mode` / `--kv-cache-disk`: + +| Mode | Meaning | +|---|---| +| `off` | Disabled. **This is the default** — no disk cache unless you opt in. | +| `auto` | Enabled with an automatically computed budget from live filesystem free space (see [Auto budget](#auto-budget)). | +| `fixed` | Enabled with an explicit hard byte cap you set. A budget is **required** in this mode. | + +Because the default is `off`, an out-of-the-box node writes nothing to disk for +the prompt cache. Turn it on deliberately. + +## Configuration surface + +The same four settings are expressible through a config file, environment +variables, and CLI flags. Sizes everywhere use **explicit IEC suffixes** +(`KiB`, `MiB`, `GiB`, `TiB`) on a positive whole number — for example `32GiB`. +Bare numbers and decimal/`GB`-style units are rejected. + +| Setting | Config (`[runtime.kv_cache.disk]`) | Environment | CLI flag | +|---|---|---|---| +| Mode / fixed budget | `mode` (`off`/`auto`/`fixed`) + `budget_mib` | `MESH_LLM_KV_CACHE_DISK` (`off`/`auto`/`SIZE`) | `--kv-cache-disk off\|auto\|SIZE` | +| Directory | `directory` (absolute) | `MESH_LLM_KV_CACHE_DISK_DIR` (absolute) | `--kv-cache-disk-dir ABSOLUTE_PATH` | +| Minimum free reserve | `minimum_free_mib` | `MESH_LLM_KV_CACHE_MIN_FREE` (SIZE) | `--kv-cache-min-free SIZE` | + +Notes: + +- In the config file, `budget_mib` is a plain MiB integer and is **only valid + when `mode = "fixed"`**; it is **required** there. Setting it under `off`/`auto` + is a config error. +- On the CLI and in the environment, mode and budget share one value: + `--kv-cache-disk 32GiB` selects fixed mode at 32 GiB; `--kv-cache-disk auto` + and `--kv-cache-disk off` select those modes. +- Directories must be absolute. Relative paths (including bare-drive forms like + `C:\cache` on non-Windows hosts) fail closed rather than resolving under the + working directory. + +### Precedence + +Each of the four settings is resolved **independently, field by field**, with +later sources overriding earlier ones: + +```text +CLI flag > MESH_LLM_KV_CACHE_* env > config file > built-in default +``` + +Legacy `SKIPPY_L3_*` variables only fill a field for which none of those public +sources supplied a value; they never override a public config field. + +So you can pin the directory in the config file and still override just the +budget with `--kv-cache-disk`, without disturbing the other fields. `mesh-llm +kv-cache status` reports the winning source per field (`default`, `config`, +`environment`, `cli`, or `legacy_environment`). + +### Legacy environment variables (deprecated) + +`SKIPPY_L3_DIR` and `SKIPPY_L3_BUDGET_BYTES` are still honored but **only as a +field-level fallback** when no public setting is present, and they emit a +deprecation warning. Migrate to `[runtime.kv_cache.disk]` or the +`MESH_LLM_KV_CACHE_*` variables. Two behaviors to know: + +- Presence of `SKIPPY_L3_DIR` (without a public mode) implies `fixed` mode, and + the budget defaults to the legacy **32 GiB** when unset. +- `SKIPPY_L3_BUDGET_BYTES=0` **no longer means unbounded** — it is treated as + the 32 GiB legacy default, with a warning. `SKIPPY_L3_BUDGET_BYTES` is ignored + entirely if `SKIPPY_L3_DIR` is not set. + +## Directory, budget, and minimum-free behavior + +**Directory.** When unset, the root resolves to `$MESH_LLM_HOME/kv-cache`, or +`~/.mesh-llm/kv-cache` when `MESH_LLM_HOME` is unset. On startup the node +creates the root and its store subdirectories (`segments/`, `manifests/`, +`prefixes/`, and the packed-store dirs; `quarantine/` is used when an object +fails verification), restricts them to owner-only permissions (`0700`), and +takes an **exclusive lock** on `.owner.lock` in the root. One process owns a cache root at a time; a second node pointed at the same +root fails to acquire the lock rather than corrupting it. The root must not +contain symlinks — a symlinked entry under the root is refused. + +**Minimum free reserve.** `minimum_free` is the free space the store preserves +for everything else on the filesystem. Default **16 GiB** +(`DEFAULT_KV_DISK_MINIMUM_FREE_MIB`); the floor is **1 GiB** +(`MIN_KV_DISK_MINIMUM_FREE_MIB`) and a smaller value is rejected. When the +filesystem sits at the reserve, the store goes **read-only**: existing entries +still serve restores, but new writes are refused (`read_only_low_space`) so the +cache never eats into the reserve. + +**Fixed budget.** The hard whole-node cap. Eviction runs oldest-manifest-first +and evicts to ~85% of the budget (a low-water margin that amortizes the +O(manifests×segments) eviction scan); the newest manifest is never evicted. + +**Auto budget.** In `auto` mode the budget is resolved +from live filesystem facts immediately before the root opens, as the **minimum** +of: + +- 20% of the capacity basis (current filesystem available + bytes already + managed under this root), +- what is actually allocatable after honoring `minimum_free`, and +- a hard **64 GiB** ceiling. + +If that resolves to `0` (for example, the filesystem is already at or below the +minimum-free reserve), the disk cache stays disabled and the node logs a warning +— again, cold prefill still serves. + +## Status and observability + +Inspect a node with: + +```bash +mesh-llm kv-cache status # human-readable +mesh-llm kv-cache status --json # machine-readable +``` + +Without `--endpoint`, the command talks to the local node's loopback control +API (default port `3131`; the control endpoint is loopback-only). To inspect +nodes you own remotely, pass one or more `--endpoint ` values (repeatable). + +The status payload reports: + +- `configured` — the resolved `mode`, `directory`, `budget_bytes`, + `minimum_free_bytes`, and the winning **source** for each field. +- `effective.state` — one of: + - `off` — mode is off. + - `active` — store is open and admitting writes. + - `read_only_low_space` — at the minimum-free reserve; reads serve, writes + refused. + - `degraded` — configured on but no manager: `reason` is + `storage_unavailable` (couldn't open the root), + `budget_below_entry_floor` (auto budget resolved to zero), or a runtime + storage error. +- `usage` — `budget_bytes`, `used_bytes`, `reserved_inflight_bytes`, + `filesystem_available_bytes`, `minimum_free_bytes`, `manifests`, + `unique_segments`, `evicted_manifests`, and `quarantined_objects`. +- `activity`, `reconciliation` (see below), and `inventory` (per-model entries). + +At **startup**, any resolution warnings (deprecated legacy vars, zero auto +budget, unavailable store) are emitted as `Warning` events in the node log. +Check the log first when the cache "isn't caching." + +## Restart vs. live-apply semantics + +On a config reload, disk-cache changes split into two classes: + +| Change | Applied | +|---|---| +| `budget_mib`, `minimum_free_mib` | **Live.** Limits update in place; shrinking evicts inactive entries immediately, pinned entries stay valid, and writes stay refused until usage fits. | +| `mode`, `directory` | **Restart required.** These are preserved across a live reload and only take effect when the node restarts. | + +If a reload changes only `mode`/`directory`, those fields are held at their +previous values (and the mode-coupled budget with them) until restart. Plan a +node restart when you move the cache directory or turn the tier on/off via +config reload. + +## Prune, clear, and shutdown + +Both operations act on **inactive** entries only — pinned/in-use state is never +removed — and both are gated behind an explicit confirmation (or `--yes`): + +```bash +# Evict least-recently-used inactive entries, optionally down to a target size +# and/or scoped to one exact model identity. +mesh-llm kv-cache prune [--target 16GiB] [--model-identity ] [--yes] + +# Remove inactive entries entirely (inference falls back to cold prefill), +# optionally scoped to one model identity; omit the filter to clear the root. +mesh-llm kv-cache clear [--model-identity ] [--yes] +``` + +- `--model-identity` takes an **exact numerical model identity**; display names + are not accepted, and a blank/whitespace filter is rejected (it would match + nothing and report a no-op success). +- `prune` without `--target` trims toward the low-water margin of the budget. + On an uncapped (legacy budget `0`) store this default would remove everything, + so that case is refused and steered to `clear`. +- Both accept `--endpoint`/`--port` for owner-controlled remote nodes and + `--json`. + +**Shutdown.** There is no flush step to run. Manifests only commit after every +referenced segment is present and the assembled payload digest matches, so a +process that stops mid-write leaves temp files, never a partial cache entry. +The root lock is released on exit; leftover temp files are reconciled on the +next start (below). A hard kill is safe — you lose only in-flight writes. + +## Corruption, quarantine, and cold fallback + +Integrity is content-addressed end to end. Every segment is addressed by the +BLAKE3 digest of its bytes and reads verify that digest, so corruption is +**detected, never silently imported**. A manifest is only loadable once all its +segments are present and the reassembled payload digest matches; partial state +is unreadable by construction. + +On startup the node **reconciles** the root before serving and records what it +repaired in `reconciliation`: + +- `removed_temporary_files` — abandoned in-flight temp files. +- `quarantined_manifests` — manifests that failed verification, moved aside into + `quarantine/`. +- `removed_prefix_links` — dangling prefix index links. +- `removed_orphan_bytes` — unreferenced segment bytes garbage-collected. + +Objects that fail verification while serving are moved to `quarantine/` and +counted in `usage.quarantined_objects`. A quarantined or missing entry is simply +a cache miss: the request falls back to cold prefill. Rising +`quarantined_objects` points at underlying storage trouble (bad disk, truncation, +external tampering) — investigate the filesystem; the cache itself stays safe. + +## Troubleshooting + +| Symptom | Likely cause | What to do | +|---|---|---| +| No disk caching at all | Mode is `off` (the default) | Set `--kv-cache-disk auto` (or `fixed SIZE`), or `[runtime.kv_cache.disk] mode`. | +| Node logs "disk prompt cache is unavailable; inference will use cold prefill" | Root couldn't be opened (permissions, missing parent, symlink in root, lock held by another process) | Check the directory exists and is writable, contains no symlinks, and no other node owns `.owner.lock`. `status` shows `degraded`/`storage_unavailable`. | +| `auto` mode caches nothing | Auto budget resolved to `0` — filesystem at/below `minimum_free` | Free space or lower `--kv-cache-min-free` (floor 1 GiB). `status` shows `degraded`/`budget_below_entry_floor`. | +| Writes refused, reads still work | At the minimum-free reserve | `status` state `read_only_low_space`; free disk, `prune`, or lower `minimum_free`. | +| Startup fails to open the cache | Another process holds the root lock | Only one node per cache root; point the second node at a different `directory`. | +| Node refuses to start with a config/resolution error (e.g. "must use an explicit IEC suffix", "must be an absolute path", "fixed disk prompt-cache mode requires a positive budget", minimum-free below 1 GiB) | **Invalid configuration** — rejected up front, not a runtime fallback | Fix the value: use `32GiB` (not `32`/`32GB`), absolute directories, a positive `budget_mib` under `fixed`, and `minimum_free ≥ 1GiB`. These are config errors, distinct from the store being unavailable at runtime. | +| Deprecation warning about `SKIPPY_L3_*` | Legacy env vars in use | Migrate to `[runtime.kv_cache.disk]` or `MESH_LLM_KV_CACHE_*`. | +| Config reload didn't move the cache / toggle the mode | `mode`/`directory` are restart-only | Restart the node; only `budget`/`minimum_free` apply live. | +| `quarantined_objects` climbing | Segments failing digest verification | Inspect the underlying storage; entries fall back to cold prefill, cache stays safe. | +| `prune` reports it freed nothing | Wrong/blank `--model-identity`, or only pinned entries present | Use the exact numerical model identity; pinned/in-use entries are never pruned. | + +## Reference + +- Config resolution and precedence: `crates/mesh-llm-host-runtime/src/runtime/kv_disk_config.rs` +- Config schema and constants: `crates/mesh-llm-config/src/model.rs`, `crates/mesh-llm-config/src/validate.rs` +- L3 store, integrity, reconciliation, status contract: `crates/skippy-cache/src/l3.rs` +- CLI and control API: `crates/mesh-llm-cli/src/parser/commands.rs`, `crates/mesh-llm-commands/src/kv_cache.rs`, `crates/mesh-llm-host-runtime/src/api/routes/kv_cache.rs` diff --git a/docs/skippy/PD_DISAGGREGATION_PLAN.md b/docs/skippy/PD_DISAGGREGATION_PLAN.md new file mode 100644 index 0000000000..b64b152fda --- /dev/null +++ b/docs/skippy/PD_DISAGGREGATION_PLAN.md @@ -0,0 +1,284 @@ +# Prefill/Decode Disaggregation — Implementation Plan + +Status: draft for review. Companion to #1427 and the #skippy-radix-L1-L2-L3 workstream. + +Progress (2026-08-28) — implemented, tested (752 tests across the four +crates), verified end-to-end on loopback: + +- **L3 substrate** — `skippy-cache::l3`: content-addressed exact-state + segment store (ordered manifests, completeness gate, idempotent puts, + capped budget, prefix index). `exact_state_identity` implements the + numerical half of the numerical-vs-placement identity split. +- **Radix L2/L3 integration** — `skippy-cache::tier::L3Tier` + + `skippy-server::kv_integration`: exact-state records write through to the + durable tier and radix misses fill back from it (re-warming RAM off the + request path). Enabled in serving via `SKIPPY_L3_DIR` + (+ `SKIPPY_L3_BUDGET_BYTES`); identity-guarded, best-effort on disk + failure. +- **`skippy-kv/1` peer fetch** — `skippy-cache::l3_remote` + the ALPN in + `skippy-protocol`: serve a store to peers, pull manifests and segments by + digest with per-segment verification; idempotent re-fetch moves zero + bytes. Harness roles `serve`/`fetch` demonstrate cross-node prefix reuse + (fetch a peer's prefilled state, decode it, byte-exact). +- **Streaming handoff with two-phase commit** — `remote-handoff + --streaming`: KV pages export per prefill chunk and stream while later + chunks compute; the receiver stages pages into a session as they arrive + but cannot generate until the commit record validates tiling, counts, and + the running digest; uncommitted state is dropped on any failure. The + recurrent snapshot is the serialized tail, as the byte math predicts. +- **Cost-based phase placement** — `skippy-topology::phase_placement`: + role assignment from gossiped compute/bandwidth signals and the + `HandoffCostModel` break-even gate (per-token KV slope, fixed + recurrent floor, link throughput, overlap fraction). +- **Harness** — `skippy-correctness remote-handoff` roles + send/recv/restore/serve/fetch with EXPERIMENTS.md counters, TTFT + baseline, per-connection reports, and a sweep script. See + REMOTE_HANDOFF_RUNBOOK.md. + +Remaining, gated on other workstreams by #1427's own sequencing: +multi-request serving waits on the #1416 iteration-level scheduler +cutover; split-prefill → collapsed decode is the step-6 generalization; +routing the openai ingress through `phase_placement` lands with #1416. + +## Goal + +Serve a request with prefill on a compute-strong node and decode on a +bandwidth-strong node, with continuation state streamed between them while +prefill is still running, as a per-request routing decision — not a static +fleet mode. Short prompts keep prefilling in place; long prompts hand off. +Unlike layer split there is no per-token network dependency: the decoder +holds the full model and generates locally. + +Lab target (per #1427): M3 Ultra prefill → M1 Ultra decode, +Nemotron 3 Super Q4 + MTPv2 — with the reverse arm (M1 prefill → M3 decode) +in the perf matrix so role assignment is settled by measurement, not +assumption. + +Prior art: EXO 1.0 shipped exactly this shape (DGX Spark prefill → M3 Ultra +decode, layer-by-layer KV streaming) in Oct 2025, and every datacenter stack +(vLLM, SGLang, TRT-LLM, Dynamo) has a PD mode. The concept needs no proving. +What is open, and what this plan targets: PD disaggregation in the +llama.cpp/GGUF world, composed with continuous batching, chunked prefill, +admission control, and quantized KV caches — as a byproduct of a general KV +mobility layer that also gives cross-node prefix-cache reuse, which EXO does +not have. + +## Design stance + +**Disaggregation is remote prefix restore.** We already have the exact shape +of the decode-side mechanism: `ProbePrefill` / `TryRestorePrefillDecode` +control frames let a driver attach a sequence whose KV was produced earlier, +with each stage restoring from its own `UnifiedRadixCache`. The only thing +missing is that today the cache a stage restores from must be local. If the +radix cache gains an L3 tier whose pages can live on disk **or on a peer**, +then "prefill over there, decode here" is just "restore a prefix whose pages +happen to be remote." One mechanism, three features: restart-surviving prefix +cache, cross-node prefix reuse, PD disaggregation. + +This is deliberately **not** a revival of the wire +`StateExport`/`StateImport` path (`binary_messaging/connection.rs:285` rejects +it, and should keep rejecting it). State moves over a dedicated backpressured +QUIC handoff stream, outside the stage activation lanes, addressed by content +digest. Two consequences of the existing plumbing must be designed out: +`MAX_STAGE_STATE_IMPORT_BYTES` (512 MiB) caps useful contexts, so the handoff +needs chunked streaming export/import APIs (none exist today — +`kv_pages.rs` is whole-buffer); and the transfer must be **two-phase +committed** on the decode side, so partially imported state can never +generate. + +**The handoff is hybrid state, not just KV.** The lab target makes this +unavoidable: Nemotron 3 Super is 8 attention + 40 Mamba + 40 MoE layers, so +an exact handoff carries KV pages + recurrent/SSM + conv state + position +metadata + MTP context/bookkeeping + committed token history for the suffix +proposer (#1037's request-local suffix history must be rebuildable on the +decoder). The per-layer-range export/import primitives for all of this exist +in `skippy-runtime/src/kv_pages.rs` (including `_for_token_count` variants); +what's missing is streaming, transport, and commit semantics. + +**MVP topology: full replicas, same backend.** One prefill node and one +decode node, each holding the whole model (single-stage), Metal↔Metal. +This sidesteps paired multi-stage pipelines and cross-backend portability. +But the identity guard does **not** come for free even then: +`PageIdentity` (`skippy-cache/src/identity.rs`) pins topology/split, so a +prefill replica and a decode replica will not match without splitting +identity into **numerical fields** (backend numerics, arch, cache types, +layer content — must match) and **placement fields** (topology, split, +node — allowed to differ). That split is the subtle correctness work and it +lands in Phase 1, not later. + +## Byte math (why per-request policy, not a mode) + +For the Nemotron 3 Super target: ~8 KiB attention-KV per prompt token → +256 MiB at 32K, plus ~160 MiB *fixed* recurrent/SSM state; roughly 350 ms at +10 Gb/s, 140 ms at 25 Gb/s line rate. Streaming KV pages behind later prompt +chunks hides most of the bulk; the recurrent snapshot + MTP context + commit +are the uncovered tail, because recurrent state is only final once prefill +finishes. The fixed 160 MiB floor also means short prompts are strictly +worse to disaggregate — break-even prompt length is a measurable function of +(state bytes, link throughput, prefill speedup ratio), and placement must be +cost-gated: reject short prompts and poor links. + +## Phases + +### Phase 0 — Falsifier benchmark (days, no product code) + +Measure the handoff cost before building anything, per the accounting already +specified in `docs/skippy/EXPERIMENTS.md:323`: +`state_export_bytes/seconds`, `state_import_bytes/seconds`, +`kv_attach_seconds`, TTFT, TPOT. + +- Use the existing offline harness (`skippy-correctness/src/runner/state_handoff.rs`) + driving `export_full_state`/`import_full_state` + (`skippy-runtime/src/kv_pages.rs`) between processes on the lab pair, with + a plain file/socket copy over the Thunderbolt bridge standing in for the + transport. +- Sweep prompt length {512, 2k, 8k, 32k} × cache type {f16, Q4_0} on + Nemotron 3 Super (hybrid: measures the fixed recurrent floor) and one + dense attention-only model (isolates the KV curve). +- Compare projected disaggregated TTFT (prefill-on-fast + transfer + attach) + vs. measured prefill-in-place on the decode node — **both role + assignments** (M3 prefill → M1 decode and the reverse), since the #1427 + perf matrix currently has no reverse arm. + +**Gate:** disaggregation must project a TTFT win at some realistic prompt +length with transfer *not* overlapped (overlap only improves it). If flat +transfer never wins below 32k tokens, stop here and write up why. + +### Phase 1 — KV page mobility in `skippy-cache` (the L3 substrate) + +This phase *is* the radix L1/L2/L3 work — the gating dependency, owned by +that workstream; #1427 consumes its L3 stream contract (segment +identity/ordering, completeness, backpressure, idempotency) with disk and +network as interchangeable backends. Note #1399 deliberately **removed** +durable disk persistence from main (deleted `disk_tier.rs`, `miss_reason.rs`, +most of `exact_state.rs`), so this is a redesign informed by that removal, +not a small delta. Mic's in-channel acceptance criteria apply: solo-mode +proof across families, measurable 19K-prefix warmup reuse, concurrent +non-duplicating loads, capped write-behind disk budget. + +- Page codec: serialize KV pages via `export_kv_page` into content-addressed + blobs in the existing BLAKE3 `blob_store`, described by `RuntimeKvPageDesc` + (layer range, token range, k/v types, row bytes, codec). +- **`PageIdentity` split into numerical vs placement fields** — numerics + (backend numerics, arch, cache types) must match across the handoff; + placement (topology, split, node) must not be pinned. The correctness- + sensitive change; do it here where the harness can gate it. +- Chunked streaming export/import APIs (today `kv_pages.rs` is whole-buffer + only, and `MAX_STAGE_STATE_IMPORT_BYTES` = 512 MiB caps a full-context + import). +- L3 disk backend: spill/fill under the radix index, eviction ladder + alongside the existing `SparseCheckpointPolicy`. + +Standalone deliverable: prefix cache that survives process restart. Ship and +validate this before any networking. + +### Phase 2 — Peer page fetch (`skippy-kv/1`) + +- New iroh subprotocol (pattern: `STAGE_ALPN_V2` registration in + `skippy-protocol/src/validation.rs` + the tunnel bridging in + `mesh-llm-host-runtime/src/network/tunnel.rs`): request pages by digest, + stream blobs, mesh-membership auth, bandwidth accounting. +- Extend the probe step: a `ProbePrefill` miss on local L1/L2/L3 may consult + peer inventory. `CacheAffinityAdvertisement` (salted prefix digests, + already gossiped — `mesh-llm-routing/src/cache_inventory.rs`) tells us + which peer to ask without leaking tokens. + +Standalone deliverable: cross-node prefix reuse — a prompt prefilled anywhere +in the mesh warms every node. This is the feature EXO doesn't have, and it is +also the entire decode-side machinery for Phase 3. + +### Phase 3 — Disaggregated serving, single-request prototype + +Single-request correctness first (#1427 is explicit); multi-request +production waits for the #1416 iteration-level scheduler cutover in Phase 4. + +- Step 1: **full-prefill → full-decode** — prefill completes, the entire + hybrid state snapshot transfers, decode attaches and continues. Behind the + `state-handoff` correctness gates before any streaming. +- Step 2: **streaming handoff with two-phase commit** — KV pages push per + completed prefill chunk over the backpressured QUIC stream, overlapping + later chunks; recurrent snapshot + MTP context + commit record are the + final segment; the decoder attaches nothing until the commit record + validates completeness. Partial state can never generate. +- Config-pinned roles for the lab: `--phase-role prefill|decode|auto` + (planner auto-placement deferred to Phase 4). +- First token samples on the prefill side (`prefill_final_frame_sampled` + already exists); decode continues from token 2. +- Failure = fallback, not resume: prefill peer dies mid-handoff → decode + node prefills locally from scratch; uncommitted segments are discarded. +- **Correctness gates** (the harness arms, in bisection order): + 1. dense attention-only model — isolates transport/commit bugs from + state-family bugs; + 2. Nemotron 3 Super without MTP — adds recurrent/SSM + conv state; + 3. Nemotron 3 Super + MTPv2 + suffix N-gram — post-handoff output must + deterministically match local continuation, which requires rebuilding + request-local suffix history (#1037) on the decoder. Watch the #1385 + failure mode: MTP weights shipped but `speculative_decoding` omitted + from the package manifest → MTP silently never selected (depends on + the open #1425 Nemotron package). +- Family gate: add a `phase_disaggregation` capability to + `reviewed-family-capabilities.json`, granted per family as it passes the + harness ladder. + +**Gate:** llama-benchy A/B on the lab — disaggregated vs decode-node-solo vs +2-stage layer split, both role assignments, reporting TTFT/TPOT and the +Phase 0 counters. Win condition: TTFT improves at long prompts with TPOT no +worse than solo decode. + +### Phase 4 — Cost-based placement and multi-request production + +- Feed the capability signals nodes already gossip but the planner ignores + (`compute_tflops_fp16`, `mem_bandwidth_gbps` in `PeerAnnouncement`) into + phase placement scoring in `skippy-topology/src/planning.rs` — there is + no role concept in the planner today. +- Activate the peer-to-peer RTT/bandwidth matrix (`edge_order.rs` is + currently dead code in production — known gap in + `docs/skippy/TOPOLOGY_PLANNER.md`); handoff needs link throughput between + the *pair*, not coordinator RTT. +- Cost-gated routing: break-even prompt length computed from measured link + throughput, state bytes/token, and the fixed recurrent floor; reject + short prompts and poor links. Decode-side admission accounts imported- + state budget in the existing `MemoryComponent` capacity model. +- Multi-request serving via the #1416 iteration-level scheduler cutover. + +### Phase 5 — Extensions (post-MVP, in pitch-value order) + +1. **Quantized-KV transfer** (Q4_0 pages over the wire) — 4× fewer bytes + exactly where the network is the constraint; nobody has shipped this. +2. **Split-prefill → collapsed decode** — multiple prefill workers each + prefilling a slice, converging on one decoder (the #1427 generalization). +3. **Cross-backend canonical page layout** — CUDA/DGX-class prefill for a + Mac decode fleet; requires an interchange codec and extending the + numerical/placement identity split with an explicit portability mode. +4. **Disagg × pipeline composition** — prefill *pipeline* feeding a decode + *pipeline* for models that fit on neither pool's single node, stage-i to + stage-i page streaming. + +## Non-goals + +- Re-enabling wire `StateExport`/`StateImport` (rejected by the server as + "not executable"; #1427 says do not re-enable — superseded by the + dedicated backpressured QUIC handoff stream). +- In-flight sequence migration or resume after node failure. +- WAN disaggregation — byte math rules it out; LAN/Thunderbolt only. +- Recurrent-state mobility *during decode* (sticky-owner stands for layer- + split serving; the one-time prefill→decode snapshot is in scope, per-token + mobility is not). + +## Open questions for review + +1. Does the L3 page granularity match the restore ladder (`SparseCheckpointPolicy` + checkpoints) or do we need a finer page size for streaming overlap? +2. Push vs pull for the Phase 3 hot path: push-on-chunk-complete is simplest; + is pull-by-digest with prefetch hints worth the extra round trips to keep + one code path with Phase 2? +3. Where does the disaggregation routing decision live — host ingress + (`mesh-llm-host-runtime/src/api`) or driver? Ingress sees the prompt + before tokenization; the driver knows scheduler state. +4. Where exactly does the numerical/placement line fall in `PageIdentity` — + in particular `ctx_size`: operationally pin it equal across the replica + pair, or classify it as placement and prove numerics are ctx-independent? +5. Role assignment prior: M1 Ultra (~800 GB/s) and M3 Ultra (~819 GB/s) are + near-equal on bandwidth but ~2× apart on compute, which argues M3=prefill + as #1427 pins it — but decode also carries MTP draft/verify compute, so + the reverse arm in the Phase 0/3 matrices settles it by measurement. diff --git a/docs/skippy/REMOTE_HANDOFF_RUNBOOK.md b/docs/skippy/REMOTE_HANDOFF_RUNBOOK.md new file mode 100644 index 0000000000..90d27df46e --- /dev/null +++ b/docs/skippy/REMOTE_HANDOFF_RUNBOOK.md @@ -0,0 +1,162 @@ +# Remote Handoff Runbook (PD disaggregation, Phase 3 step 1) + +`skippy-correctness remote-handoff` runs the full-prefill → full-decode +handoff between two machines: the sender prefills a prompt, exports the +continuation state, streams it to the receiver in digest-verified segments, +and the receiver imports nothing until the commit record validates +completeness — then both sides greedy-decode the same continuation and the +tokens are compared one-for-one. With `--baseline` the receiver also +measures prefill-in-place, so one run yields the disaggregated-vs-local TTFT +comparison with the counters from `EXPERIMENTS.md` (export/transfer/import +bytes and seconds, attach, first-decode). + +## Requirements + +- Same model file, `--ctx-size`, `--layer-end`, lane count, and payload kind + on both sides (validated in the handshake; mismatches are rejected). +- A fresh native ABI build: the full-state header parsing was fixed in patch + 0023 (2026-08-27), so stale `.deps/llama-build` archives fail with + "full-state import restored native position 0". Run `just llama-prepare && + just llama-build` if in doubt, then `cargo build -p skippy-correctness`. +- Payload kinds: `full-state` (default; dense attention models) or + `kv-recurrent` (hybrid models — untested until a hybrid package is + available, see #1425). + +## L3 store integration + +Pass `--store-dir ` on either side to route the handoff through the +L3 segment store (`skippy-cache::l3`): the sender spills its exported state +(segments + manifest, off the transfer critical path), and the receiver +write-behinds incoming segments to disk and imports via the store's +`assemble` path — every segment digest, the tiling, and the whole-payload +digest re-verified. `--store-budget-bytes` caps the on-disk footprint +(oldest manifests evict first). The manifest records the sender's reference +continuation, so state in the store is self-verifying. + +Restart survival: reattach from a store with no network and no exporter — + +```bash +target/release/skippy-correctness remote-handoff --role restore \ + --store-dir --model --layer-end \ + --ctx-size 8192 --n-gpu-layers 99 --decode-tokens 32 \ + --report-out restore-report.json +``` + +`--manifest ` selects a specific manifest (default: +newest). The restore refuses manifests whose `exact_state_identity` (the +numerical identity: weights, cache dtypes, flash-attn, backend, layer +range, context shape — never stage/topology placement) does not match the +local configuration. + +## Streaming handoff (overlap + two-phase commit) + +Pass `--streaming` on **both** sides: the sender exports the KV page for +each prefill chunk (`--stream-chunk-tokens`, default 512) and streams it +while later chunks compute, so transfer and the receiver's import hide +inside the prefill wall. The recurrent snapshot (hybrid families) is the +serialized tail. The receiver stages pages into a session as they arrive +but **cannot generate until the commit record validates** page tiling, +counts, and the running payload digest — any failure drops the staged +session. The report's `overlap_wall_ms` plus the receiver's +`attach_residual_ms` and `first_decode_ms` compose the streaming TTFT. +Restore from a page-stream manifest re-imports page by page (pass +`--streaming` to restore too so identity matches). + +## Peer fetch (`skippy-kv/1`) — cross-node prefix reuse + +The direct TCP listener is an unauthenticated lab transport. It defaults to +loopback; bind it to a non-loopback address only on a trusted private network +with host firewall rules limiting both peers. Production mesh exposure must use +the registered `skippy-kv/1` iroh ALPN so mesh membership authenticates the +peer. Each TCP connection has bounded read and write deadlines, but those +deadlines do not provide authentication or confidentiality. + +Any node can serve its store and any node can pull by digest: + +```bash +# node A: serve the store (no model load) +target/release/skippy-correctness remote-handoff --role serve \ + --listen 0.0.0.0:19092 --store-dir --model --layer-end + +# node B: pull the newest manifest, then restore + decode from it +target/release/skippy-correctness remote-handoff --role fetch \ + --peer :19092 --store-dir \ + --model --layer-end --ctx-size 8192 --n-gpu-layers 99 \ + --decode-tokens 32 --report-out fetch-report.json +``` + +Fetches are idempotent (content-addressed: held segments transfer zero +bytes) and every segment is digest-verified before it lands locally. + +## Serving-path L3 tier + +The same store backs real serving: set `SKIPPY_L3_DIR=` (and +optionally `SKIPPY_L3_BUDGET_BYTES`) on a stage with the exact-state prefix +cache enabled, and recorded exact-state entries write through to disk while +radix misses fill back from it — prefix reuse that survives restarts and +RAM eviction. The tier identity is the radix namespace hash, so a +configuration change refuses stale state. + +## Two-machine run + +Use a **release** build for measurements (`cargo build --release -p +skippy-correctness`) — debug-build byte handling distorts transfer and +export timings. + +Receiver (decode node) first — it loads the model, then listens: + +```bash +target/release/skippy-correctness remote-handoff --role recv \ + --listen 0.0.0.0:19081 \ + --model --layer-end --ctx-size 8192 \ + --n-gpu-layers 99 --prefix-token-count 4096 --decode-tokens 32 \ + --report-out recv-report.json +``` + +Sender (prefill node), once the receiver prints `ready`: + +```bash +target/release/skippy-correctness remote-handoff --role send \ + --peer :19081 \ + --model --layer-end --ctx-size 8192 \ + --n-gpu-layers 99 --prefix-token-count 4096 --decode-tokens 32 \ + --baseline --report-out send-report.json +``` + +The sender's report is the primary artifact: `ttft_disaggregated_ms` +(prefill + export + transfer + attach + first decode) vs `ttft_local_ms` +(receiver's prefill-in-place + first decode), `ttft_speedup`, +`transfer_gbps`, and `matches` (exact token agreement, the correctness +gate). Non-zero exit on mismatch unless `--allow-mismatch`. + +## Sweep for the perf matrix + +Start the receiver once with `--accept-count --allow-mismatch` (it +serves n handoffs, writing `report-1.json … report-n.json`), then drive the +sender side with `scripts/remote-handoff-sweep.sh`: + +```bash +scripts/remote-handoff-sweep.sh :19081 \ + out/ 512 2048 4096 8192 +``` + +It prints a summary table (state MiB, link Gbps, per-phase ms, TTFT +disaggregated vs local, speedup, match). Run both role assignments (fast box +sends, then fast box receives). Keep `--ctx-size` at least prefix + decode +tokens on both sides. `transfer_gbps` on the Thunderbolt bridge tells you +whether the link, not the runtime, bounds the handoff. + +Timing caveat: Metal execution is asynchronous, so `source_prefill_ms` can +under-report with the balance absorbed into `state_export_ms` (export +synchronizes). The TTFT aggregates are correct; per-phase attribution +between those two columns is approximate. + +## Interpreting + +- `matches: true` — exact-state handoff is deterministic; the correctness + half of the #1427 step-1 gate. +- `ttft_speedup > 1` at some prefix length — disaggregation pays on this + pair; the break-even length feeds the Phase 4 cost gate. +- This prototype transfers after prefill completes (no chunk streaming) and + runs one request; it is the measurement harness for the EXPERIMENTS.md + falsifier, not the serving integration. diff --git a/docs/skippy/cachegen-lmcache-qwen3-0.6b-19k-summary.json b/docs/skippy/cachegen-lmcache-qwen3-0.6b-19k-summary.json new file mode 100644 index 0000000000..f0e136ad92 --- /dev/null +++ b/docs/skippy/cachegen-lmcache-qwen3-0.6b-19k-summary.json @@ -0,0 +1,76 @@ +{ + "schema_version": 1, + "decision": "quality-pass-latency-stop", + "date": "2026-09-11", + "commit": "4bf0865be6e7a13dbf7159b0fb6e6c41f29db72c", + "reference": { + "implementation": "LMCache CacheGen", + "revision": "b5d109ea99a89b4d8a670ee4fc2e8cb76411ee5c", + "license": "Apache-2.0" + }, + "hardware": { + "model": "Mac Studio", + "chip": "Apple M1 Ultra", + "memory_gib": 128, + "runtime_backend": "Metal", + "codec_backend": "scalar Rust CPU" + }, + "model": { + "id": "Qwen/Qwen3-0.6B-GGUF:Q8_0", + "sha256": "12fae8b8f78f0360b498d04c8db7d33aff29ab7d8080231f93a17c18119e6735", + "layer_end": 28 + }, + "workload": { + "prefix_tokens": 19000, + "continuation_steps": 64, + "ctx_size": 19200, + "cache_types": { + "k": "f16", + "v": "f16" + }, + "chunk_rows": 256 + }, + "thresholds": { + "minimum_token_agreement": 0.95, + "maximum_p99_decode_regression": 0.05, + "restore_to_first_token_must_beat_native": true + }, + "metrics": { + "native_storage_bytes": 2179072000, + "cachegen_storage_bytes": 446903003, + "compression_ratio": 0.20508868132856556, + "tile_count": 4200, + "encode_ms": 20732.857249999997, + "decode_ms": 41285.729709, + "native_write_ms": 2402.914834, + "cachegen_write_ms": 693.576666, + "native_persist_ms": 2402.914834, + "cachegen_persist_ms": 21426.433915999998, + "native_read_ms": 281.062542, + "cachegen_read_ms": 58.829709, + "native_import_ms": 60.07575, + "cachegen_import_ms": 57.972333, + "native_ttft_ms": 374.379459, + "cachegen_ttft_ms": 41411.126585, + "native_decode_tokens_per_second": 119.84529650055929, + "cachegen_decode_tokens_per_second": 124.166031091795, + "native_p99_decode_ms": 33.241167000000004, + "cachegen_p99_decode_ms": 8.594834, + "p99_decode_regression": -0.7414400643635647, + "matching_tokens": 64, + "token_agreement": 1.0, + "mean_entropy_abs_drift": 0.34619530792406294, + "max_entropy_abs_drift": 2.1676456928253174, + "mean_top_logprob_abs_drift": 0.06626987643903703, + "max_top_logprob_abs_drift": 0.4230744540691376, + "estimated_peak_codec_working_bytes": 2627965638 + }, + "failure_reasons": [ + "restore-to-first-token did not beat native (41411.127 ms >= 374.379 ms)" + ], + "notes": [ + "The 64/64 agreement result clears the predeclared 95% quality threshold.", + "The scalar CPU codec is retained as a deterministic oracle; request-path promotion requires a parallel device implementation to beat native restore-to-first-token.", + "Filesystem writes used sync_all; same-run reads may be page-cache warm." + ] +} diff --git a/docs/skippy/cachegen-metal-device-qwen3-0.6b-19k-summary.json b/docs/skippy/cachegen-metal-device-qwen3-0.6b-19k-summary.json new file mode 100644 index 0000000000..6f18445d42 --- /dev/null +++ b/docs/skippy/cachegen-metal-device-qwen3-0.6b-19k-summary.json @@ -0,0 +1,80 @@ +{ + "schema_version": 1, + "decision": "quality-pass-local-latency-stop", + "date": "2026-09-11", + "commit": "c86cf848b0fa183f2c2e594fe1e51bf1024d7185", + "reference": { + "implementation": "LMCache CacheGen", + "revision": "b5d109ea99a89b4d8a670ee4fc2e8cb76411ee5c", + "license": "Apache-2.0" + }, + "hardware": { + "model": "Mac Studio", + "chip": "Apple M1 Ultra", + "memory_gib": 128, + "runtime_backend": "Metal", + "codec_backend": "native MSL direct-to-resident-KV" + }, + "model": { + "id": "Qwen/Qwen3-0.6B-GGUF:Q8_0", + "sha256": "12fae8b8f78f0360b498d04c8db7d33aff29ab7d8080231f93a17c18119e6735", + "layer_end": 28 + }, + "workload": { + "prefix_tokens": 19000, + "continuation_steps": 64, + "ctx_size": 19200, + "cache_types": { + "k": "f16", + "v": "f16" + }, + "chunk_rows": 256, + "continuation_sessions": "isolated" + }, + "thresholds": { + "minimum_token_agreement": 0.95, + "maximum_p99_decode_regression": 0.05, + "restore_to_first_token_must_beat_native": true + }, + "metrics": { + "restore_path": "native-device", + "native_storage_bytes": 2179072000, + "cachegen_storage_bytes": 446903003, + "compression_ratio": 0.20508868132856556, + "tile_count": 4200, + "encode_ms": 21133.011542, + "scalar_oracle_decode_ms": 41436.069375, + "native_write_ms": 424.18879100000004, + "cachegen_write_ms": 199.407333, + "native_persist_ms": 424.18879100000004, + "cachegen_persist_ms": 21332.418875, + "native_read_ms": 296.2475, + "cachegen_read_ms": 54.211334, + "native_import_ms": 56.690583000000004, + "cachegen_import_ms": 540.7443750000001, + "native_ttft_ms": 385.41308300000003, + "cachegen_ttft_ms": 604.8354170000001, + "native_decode_tokens_per_second": 122.33699996373663, + "cachegen_decode_tokens_per_second": 128.07457193180554, + "native_p99_decode_ms": 32.474999999999994, + "cachegen_p99_decode_ms": 9.879707999999999, + "p99_decode_regression": -0.6957749653579676, + "matching_tokens": 64, + "token_agreement": 1.0, + "mean_entropy_abs_drift": 0.3461667150259018, + "max_entropy_abs_drift": 2.166335344314575, + "mean_top_logprob_abs_drift": 0.06625335784883646, + "max_top_logprob_abs_drift": 0.4222277104854584, + "estimated_peak_codec_working_bytes": 2627965638 + }, + "failure_reasons": [ + "restore-to-first-token did not beat native (604.835 ms >= 385.413 ms)" + ], + "notes": [ + "The direct Metal decoder consumes the persisted archive and writes resident K/V cells without materializing a decoded host page.", + "The scalar CPU decoder runs as an independently timed oracle and is excluded from direct-device TTFT.", + "The 64/64 agreement result clears the predeclared 95% quality threshold.", + "Filesystem writes used sync_all; same-run reads may be page-cache warm.", + "This local-tier result does not qualify or reject a transport-bound remote tier." + ] +} diff --git a/docs/skippy/cachegen-metal-typed-qwen3-0.6b-19k-summary.json b/docs/skippy/cachegen-metal-typed-qwen3-0.6b-19k-summary.json new file mode 100644 index 0000000000..83b7499268 --- /dev/null +++ b/docs/skippy/cachegen-metal-typed-qwen3-0.6b-19k-summary.json @@ -0,0 +1,228 @@ +{ + "schema_version": 1, + "decision": "f32-and-f32-f16-local-pass-lower-width-types-stop", + "date": "2026-09-12", + "commits": { + "quantized_and_mixed": "2e26d46ea87e1b8ee783460998e703f669513f91", + "f32": "74b4f60719d09dee7d9579c1365ac6f95d14c20c", + "f32_mixed_crossover": "ddddf34aa5e64262c9e113d07f9dcb506e5f7aab", + "metal_direct_staging": "e2bdc935e24cbed3bff767762547098db6028459" + }, + "reference": { + "implementation": "LMCache CacheGen", + "revision": "b5d109ea99a89b4d8a670ee4fc2e8cb76411ee5c", + "license": "Apache-2.0" + }, + "hardware": { + "model": "Mac Studio", + "chip": "Apple M1 Ultra", + "memory_gib": 128, + "runtime_backend": "Metal", + "codec_backend": "native MSL direct-to-resident-KV" + }, + "model": { + "id": "Qwen/Qwen3-0.6B-GGUF:Q8_0", + "sha256": "12fae8b8f78f0360b498d04c8db7d33aff29ab7d8080231f93a17c18119e6735", + "layer_end": 28 + }, + "workload": { + "prefix_tokens": 19000, + "continuation_steps": 64, + "ctx_size": 19200, + "chunk_rows": 256, + "continuation_sessions": "isolated" + }, + "thresholds": { + "minimum_token_agreement": 0.95, + "maximum_p99_decode_regression": 0.05, + "encoded_payload_must_be_smaller_than_native": true, + "restore_to_first_token_must_beat_native": true + }, + "metal_direct_staging": { + "change": "write validated jobs directly into final shared MTLBuffers", + "q8_0_f32": { + "baseline_cachegen_import_ms": [586.035708, 574.8656669999999], + "optimized_cachegen_import_ms": [497.411834, 486.434625], + "mean_import_reduction": 0.1525, + "optimized_cachegen_ttft_ms": [713.511958, 701.9677919999999], + "optimized_native_ttft_ms": [571.8813749999999, 547.7787920000001], + "decision": "latency-stop" + }, + "f16_f16": { + "baseline_cachegen_import_ms": 540.74, + "optimized_cachegen_import_ms": 451.817375, + "import_reduction": 0.1644, + "optimized_cachegen_ttft_ms": 514.338417, + "optimized_native_ttft_ms": 343.15962500000006, + "decision": "latency-stop" + }, + "chunk_128_diagnostic": { + "tile_count": 8344, + "cachegen_storage_bytes": 850615819, + "compression_ratio": 0.3084301534362527, + "cachegen_import_ms": 727.739791, + "decision": "reject-smaller-chunks" + } + }, + "cases": [ + { + "cache_type_k": "f32", + "cache_type_v": "f32", + "passed": true, + "native_storage_bytes": 4358144000, + "cachegen_storage_bytes": 446903003, + "compression_ratio": 0.10254434066428278, + "native_import_ms": 494.848083, + "cachegen_import_ms": 521.351959, + "native_ttft_ms": 1199.1292079999998, + "cachegen_ttft_ms": 766.1168339999999, + "p99_decode_regression": -0.047207090345013546, + "matching_tokens": 64, + "token_agreement": 1.0, + "first_token_mismatch_step": null, + "failure_reasons": [] + }, + { + "cache_type_k": "q8_0", + "cache_type_v": "f32", + "passed": false, + "stable_result": "one-pass-one-fail", + "native_storage_bytes": 2757888000, + "cachegen_storage_bytes": 565441003, + "compression_ratio": 0.20502681871054954, + "runs": [ + { + "native_ttft_ms": 800.57175, + "cachegen_ttft_ms": 796.943875, + "passed": true + }, + { + "native_ttft_ms": 727.173833, + "cachegen_ttft_ms": 792.1295419999999, + "passed": false + } + ], + "matching_tokens": 64, + "token_agreement": 1.0 + }, + { + "cache_type_k": "f32", + "cache_type_v": "f16", + "passed": true, + "native_storage_bytes": 3268608000, + "cachegen_storage_bytes": 446903003, + "compression_ratio": 0.13672578755237705, + "native_ttft_ms": 878.6078749999999, + "cachegen_ttft_ms": 686.1472490000001, + "p99_decode_regression": -0.19336615460457712, + "matching_tokens": 64, + "token_agreement": 1.0, + "first_token_mismatch_step": null, + "failure_reasons": [] + }, + { + "cache_type_k": "f16", + "cache_type_v": "f32", + "passed": true, + "native_storage_bytes": 3268608000, + "cachegen_storage_bytes": 446903003, + "compression_ratio": 0.13672578755237705, + "native_ttft_ms": 840.3155, + "cachegen_ttft_ms": 671.9314999999999, + "p99_decode_regression": -0.16524231813735718, + "matching_tokens": 64, + "token_agreement": 1.0, + "first_token_mismatch_step": null, + "failure_reasons": [] + }, + { + "cache_type_k": "q8_0", + "cache_type_v": "q8_0", + "passed": false, + "native_storage_bytes": 1157632000, + "cachegen_storage_bytes": 589803358, + "compression_ratio": 0.5094912355567227, + "native_import_ms": 129.030291, + "cachegen_import_ms": 571.4625410000001, + "native_ttft_ms": 333.361041, + "cachegen_ttft_ms": 650.6419990000002, + "p99_decode_regression": -0.6727379064515878, + "matching_tokens": 64, + "token_agreement": 1.0, + "first_token_mismatch_step": null, + "failure_reasons": [ + "restore-to-first-token did not beat native (650.642 ms >= 333.361 ms)" + ] + }, + { + "cache_type_k": "q4_0", + "cache_type_v": "q4_0", + "passed": false, + "native_storage_bytes": 612864000, + "cachegen_storage_bytes": 635037607, + "compression_ratio": 1.0361803059079993, + "native_import_ms": 80.699625, + "cachegen_import_ms": 619.267791, + "native_ttft_ms": 189.84379199999998, + "cachegen_ttft_ms": 699.35204, + "p99_decode_regression": -0.6919925490103622, + "matching_tokens": 61, + "token_agreement": 0.953125, + "first_token_mismatch_step": 15, + "failure_reasons": [ + "encoded payload is not smaller than native", + "restore-to-first-token did not beat native (699.352 ms >= 189.844 ms)" + ] + }, + { + "cache_type_k": "q8_0", + "cache_type_v": "f16", + "passed": false, + "native_storage_bytes": 1668352000, + "cachegen_storage_bytes": 565441003, + "compression_ratio": 0.3389218839909084, + "native_import_ms": 189.447709, + "cachegen_import_ms": 579.317334, + "native_ttft_ms": 558.277709, + "cachegen_ttft_ms": 792.0514169999999, + "p99_decode_regression": -0.16607278770613987, + "matching_tokens": 64, + "token_agreement": 1.0, + "first_token_mismatch_step": null, + "failure_reasons": [ + "restore-to-first-token did not beat native (792.051 ms >= 558.278 ms)" + ] + }, + { + "cache_type_k": "q4_0", + "cache_type_v": "f16", + "passed": false, + "native_storage_bytes": 1395968000, + "cachegen_storage_bytes": 610269862, + "compression_ratio": 0.43716608260361267, + "native_import_ms": 153.555959, + "cachegen_import_ms": 602.013291, + "native_ttft_ms": 483.978876, + "cachegen_ttft_ms": 807.0270820000001, + "p99_decode_regression": -0.17710234510665943, + "matching_tokens": 61, + "token_agreement": 0.953125, + "first_token_mismatch_step": 15, + "failure_reasons": [ + "restore-to-first-token did not beat native (807.027 ms >= 483.979 ms)" + ] + } + ], + "notes": [ + "The scalar CPU decoder ran as an independently timed oracle and is excluded from direct-device TTFT.", + "The F32/F32 case passed all declared gates after fixing native model configuration to preserve GGML_TYPE_F32 enum value zero instead of silently replacing it with F16.", + "Both F32/F16 directions passed with more than 160 ms of TTFT margin; Q8_0/F32 crossed the boundary between repeated runs and remains stopped.", + "F32/Q8_0 is not runtime-valid for this model because quantized V requires Flash Attention and this mixed pair has no compatible Flash Attention kernel.", + "Q8_0 cases preserved all 64 greedy continuation tokens; Q4_0 K cases first diverged at step 15 and preserved 61 of 64.", + "F32/F32 beat its matched native control on this local unified-memory tier; the four sampled quantized and mixed restores were slower.", + "Q4_0/Q4_0 expanded the native payload by 3.62 percent and therefore failed the storage gate independently of latency.", + "The matrix is representative rather than exhaustive; the remaining mixed pairings retain fixture coverage only.", + "Direct Metal staging reduced F16/F16 and Q8_0/F32 import by about 16 percent but did not make either case beat its matched native TTFT control.", + "A 128-row diagnostic doubled tile count, worsened compression, and increased import latency; further decoder parallelism needs shared calibration across multiple arithmetic substreams rather than smaller archive tiles." + ] +} diff --git a/docs/skippy/cachegen-quality-gate-qwen3-0.6b-19k-summary.json b/docs/skippy/cachegen-quality-gate-qwen3-0.6b-19k-summary.json new file mode 100644 index 0000000000..e24fcac351 --- /dev/null +++ b/docs/skippy/cachegen-quality-gate-qwen3-0.6b-19k-summary.json @@ -0,0 +1,71 @@ +{ + "schema_version": 1, + "decision": "stop", + "date": "2026-09-11", + "commit": "2677ad62e5295f6da2ac72ae7b8c978f87753d11", + "hardware": { + "model": "Mac Studio", + "chip": "Apple M1 Ultra", + "memory_gib": 128, + "backend": "Metal" + }, + "model": { + "id": "Qwen/Qwen3-0.6B-GGUF:Q8_0", + "sha256": "9465e63a22add5354d9bb4b99e90117043c7124007664907259bd16d043bb031", + "layer_end": 28 + }, + "workload": { + "prefix_tokens": 19000, + "continuation_steps": 64, + "ctx_size": 19200, + "cache_types": { + "k": "f16", + "v": "f16" + }, + "tile_rows": 4096 + }, + "thresholds": { + "minimum_token_agreement": 0.95, + "maximum_p99_decode_regression": 0.05, + "restore_to_first_token_must_beat_native": true + }, + "metrics": { + "native_storage_bytes": 2179072000, + "cachegen_storage_bytes": 202373739, + "compression_ratio": 0.09287152466738134, + "tile_count": 280, + "encode_ms": 15394.387625, + "decode_ms": 23521.694708, + "native_write_ms": 1369.8386249999999, + "cachegen_write_ms": 114.436834, + "native_persist_ms": 1369.8386249999999, + "cachegen_persist_ms": 15508.824459, + "native_read_ms": 239.139542, + "cachegen_read_ms": 22.285667, + "native_import_ms": 53.609541, + "cachegen_import_ms": 51.117250000000006, + "native_ttft_ms": 311.610791, + "cachegen_ttft_ms": 23604.395332999997, + "native_decode_tokens_per_second": 120.22645848205526, + "cachegen_decode_tokens_per_second": 121.2529795622552, + "native_p99_decode_ms": 18.861708, + "cachegen_p99_decode_ms": 9.297708, + "p99_decode_regression": -0.507059063792102, + "matching_tokens": 12, + "token_agreement": 0.1875, + "first_token_mismatch_step": 1, + "mean_entropy_abs_drift": 2.5719402572722174, + "max_entropy_abs_drift": 7.5579400062561035, + "mean_top_logprob_abs_drift": 0.9204300816782052, + "max_top_logprob_abs_drift": 3.5587148666381836, + "estimated_peak_codec_working_bytes": 2412371750 + }, + "failure_reasons": [ + "restore-to-first-token did not beat native (23604.395 ms >= 311.611 ms)", + "token agreement 0.1875 is below 0.9500" + ], + "notes": [ + "Filesystem writes used sync_all; reads were matched same-run filesystem reads and may be page-cache warm.", + "The codec decode cost alone exceeds the native full restore-to-first-token path by more than two orders of magnitude, so read-cache state cannot change the stop decision." + ] +} diff --git a/evals/README.md b/evals/README.md index 8a61eb5df4..23bc47cb14 100644 --- a/evals/README.md +++ b/evals/README.md @@ -125,6 +125,50 @@ across refs and passes. `--max-ttft-regression-pct` bounds candidate median TTFT relative to the first ref. Any configured failure is written into `run.json` and the Markdown report before the command exits non-zero. +### Disk L3 lifecycle certification + +Use `l3-plan` and `l3-run` for the persistent disk-cache release gate. Unlike +the comparative `run` command, this mode evaluates one release candidate while +preserving the same cache root across verified process restarts. It requires +named captured Buzz, OpenCode, and Goose sources plus c64/c128/c256 cohorts. + +```bash +python3 evals/agentic-replay.py l3-plan \ + --ref candidate= \ + --model '' \ + --trajectory-manifest /path/to/captured-l3-trajectories.json \ + --require-source-dataset buzz \ + --require-source-dataset opencode \ + --require-source-dataset goose + +python3 evals/agentic-replay.py l3-run \ + --ref candidate= \ + --model '' \ + --trajectory-manifest /path/to/captured-l3-trajectories.json \ + --require-source-dataset buzz \ + --require-source-dataset opencode \ + --require-source-dataset goose \ + --output /path/to/l3-artifact +``` + +The manifest must contain an `l3` cohort with each required source and `64`, +`128`, and `256` high-load cohorts with at least that many trajectories. The +runner performs disk-off cold restarts, an empty-root write, multi-turn growth, +same-process L1 reuse, first-request-after-restart L3 samples, 100-request +identical-prefix fill and record waves, prune and clear while traffic is in +flight, forced minimum-free fallback, and disk-off/disk-on high-load pairs. + +The command fails closed unless generated-output hashes match disk-off, every +qualifying prompt is in the configured token range, post-restart L3 p50 TTFT is +at most half of cold p50, fill and record waves each perform one physical +operation, payload writes stay within 1.2x, low space remains inference-safe, +and p99 decode-event latency at c64/c128/c256 regresses by no more than 5%. +Server/API status snapshots, raw requests, process commands and PIDs, logs, +binary/runtime hashes, a Markdown report, and an artifact hash inventory are +retained. Run the same artifact recipe on every supported backend/model family +and on each stage of the real two-machine split; the harness never treats a +single-node pass as full-chain evidence. + ### Compare Mesh with llama.cpp, vLLM, and SGLang Pass `--engine-config` to append external OpenAI-compatible server arms to the diff --git a/evals/agentic-replay.py b/evals/agentic-replay.py index 8e72ddc26b..5df1ccf463 100755 --- a/evals/agentic-replay.py +++ b/evals/agentic-replay.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -"""Agentic Replay: compare inference engines on ordered real-agent trajectories. +"""Agentic Replay: compare inference engines or certify disk-L3 lifecycle behavior. The runner creates detached worktrees, builds the release host and native runtime for every requested ref, replays a deterministic subset of the pinned @@ -7,11 +7,17 @@ tables, CSV, and dependency-free SVG charts. Mesh ref arms are deliberately launched without context-size, lane-count, -KV-budget, or backend-tuning arguments. External llama.cpp, vLLM, and SGLang -arms use an explicit engine configuration so their capacity and model choices -are visible in the artifact. Their identity is the SHA-256 of the engine's -reported version. Client concurrency is offered by this runner and is not a -server startup setting. +KV-budget, or backend-tuning arguments. The only serving argument is +``--model``; ``--log-format json`` is observational. External llama.cpp, vLLM, +and SGLang arms use an explicit engine configuration so their capacity and +model choices are visible in the artifact. Their identity is the SHA-256 of +the engine's reported version. Client concurrency is offered by this runner +and is not a server startup setting. + +The ``l3-*`` commands run captured trajectories through disk-off, empty-root, +same-process, post-restart, concurrent-fill, lifecycle-operation, and low-space +phases. They preserve the cache root across restarts and derive acceptance from +API counters plus generated-output hashes. """ from __future__ import annotations @@ -33,6 +39,7 @@ import subprocess import sys import tempfile +import threading import time from dataclasses import dataclass from datetime import datetime, timezone @@ -74,6 +81,8 @@ def verified_version_sha256_by_label( DEFAULT_ENDPOINT = urlsplit(DEFAULT_BASE_URL) DEFAULT_HOST = DEFAULT_ENDPOINT.hostname or "127.0.0.1" DEFAULT_PORT = DEFAULT_ENDPOINT.port or 80 +MANAGEMENT_HOST = "127.0.0.1" +MANAGEMENT_PORT = 3131 FORBIDDEN_STARTUP_OPTIONS = ( "--ctx-size", "--generation-concurrency", @@ -181,6 +190,16 @@ def parse_ref_specs(repo: Path, values: Sequence[str]) -> list[RefSpec]: return specs +def parse_single_ref_spec(repo: Path, value: str) -> RefSpec: + if "=" not in value: + raise ValueError(f"ref must use LABEL=GIT_REF syntax: {value}") + label, ref = value.split("=", 1) + label, ref = slug(label), ref.strip() + if not ref: + raise ValueError(f"empty git ref for label {label}") + return RefSpec(label=label, ref=ref, commit=git(repo, "rev-parse", f"{ref}^{{commit}}")) + + def external_config(args: argparse.Namespace) -> EngineConfig | None: path = getattr(args, "engine_config", None) if path is None: @@ -742,6 +761,7 @@ def stream_request( prompt_tokens = 0 cached_tokens = 0 content_events = 0 + decode_event_times: list[float] = [] content_parts: list[str] = [] reasoning_parts: list[str] = [] tool_call_parts: dict[int, dict[str, Any]] = {} @@ -820,8 +840,10 @@ def stream_request( reasoning_content = delta.get("reasoning_content") tool_calls = delta.get("tool_calls") if content or reasoning_content or tool_calls: + event_at = time.monotonic() if first_token_at is None: - first_token_at = time.monotonic() + first_token_at = event_at + decode_event_times.append(event_at) content_events += 1 if isinstance(content, str): content_parts.append(content) @@ -868,6 +890,10 @@ def stream_request( "ttft_seconds": first_token_at - started, "elapsed_seconds": completed - started, "generation_seconds": completed - first_token_at, + "decode_inter_token_seconds": [ + later - earlier + for earlier, later in zip(decode_event_times, decode_event_times[1:]) + ], "completion_tokens": completion_tokens, "prompt_tokens": prompt_tokens, "cached_tokens": cached_tokens, @@ -1031,6 +1057,11 @@ def summarize_requests( generation_seconds = sum( request["generation_seconds"] for request in successful ) + decode_inter_token = [ + sample + for request in successful + for sample in request.get("decode_inter_token_seconds", []) + ] if successful: workload_window = max(request["completed"] for request in successful) - min( request["started"] for request in successful @@ -1091,6 +1122,7 @@ def summarize_requests( if generation_seconds > 0 else None ), + "decode_inter_token_p99_seconds": percentile(decode_inter_token, 0.99), "mean_in_flight": mean_in_flight, "concurrency_utilization_pct": ( 100 * mean_in_flight / offered_concurrency @@ -1263,14 +1295,20 @@ def start_server( log_path: Path, hf_home: Optional[Path], ) -> tuple[subprocess.Popen[bytes], list[str]]: - if port_is_open(): + external = build.get("engine", "mesh") != "mesh" + management_port_open = not external and port_is_open(MANAGEMENT_HOST, MANAGEMENT_PORT) + if port_is_open() or management_port_open: raise RuntimeError( - f"TCP {DEFAULT_PORT} is already in use; stop the existing inference server" + ( + f"TCP {DEFAULT_PORT} or {MANAGEMENT_PORT} is already in use; " + "stop the existing Mesh instance" + if not external + else f"TCP {DEFAULT_PORT} is already in use; stop the existing inference server" + ) ) command = command_for_build(build, model) log_path.parent.mkdir(parents=True, exist_ok=True) log_handle = log_path.open("wb") - external = build.get("engine", "mesh") != "mesh" process = subprocess.Popen( command, cwd=Path(build["worktree"]), @@ -1289,7 +1327,9 @@ def start_server( return process, command -def stop_server(process: subprocess.Popen[bytes]) -> None: +def stop_server( + process: subprocess.Popen[bytes], *, include_management_port: bool = True +) -> None: if process.poll() is None: os.killpg(process.pid, signal.SIGINT) try: @@ -1302,12 +1342,240 @@ def stop_server(process: subprocess.Popen[bytes]) -> None: os.killpg(process.pid, signal.SIGKILL) process.wait(timeout=10) deadline = time.monotonic() + 10 - while port_is_open() and time.monotonic() < deadline: + while ( + port_is_open() + or (include_management_port and port_is_open(MANAGEMENT_HOST, MANAGEMENT_PORT)) + ) and time.monotonic() < deadline: time.sleep(0.2) - if port_is_open(): + occupied = [ + port + for host, port in ( + (DEFAULT_HOST, DEFAULT_PORT), + (MANAGEMENT_HOST, MANAGEMENT_PORT), + ) + if (port != MANAGEMENT_PORT or include_management_port) and port_is_open(host, port) + ] + if occupied: + raise RuntimeError(f"Mesh stopped but TCP ports {occupied} are still occupied") + if process.poll() is None: + raise RuntimeError(f"Mesh PID {process.pid} remained alive after shutdown") + + +def l3_server_command( + binary: Path, + model: str, + cache_root: Optional[Path], + budget: str, + minimum_free: str, +) -> list[str]: + command = server_command(binary, model) + if cache_root is not None: + command.extend( + ( + "--kv-cache-disk", + budget, + "--kv-cache-disk-dir", + str(cache_root), + "--kv-cache-min-free", + minimum_free, + ) + ) + return command + + +def start_l3_server( + build: dict[str, Any], + model: str, + state_dir: Path, + log_path: Path, + hf_home: Optional[Path], + cache_root: Optional[Path], + budget: str, + minimum_free: str, +) -> tuple[subprocess.Popen[bytes], list[str]]: + if port_is_open() or port_is_open(MANAGEMENT_HOST, MANAGEMENT_PORT): + raise RuntimeError( + f"TCP {DEFAULT_PORT} or {MANAGEMENT_PORT} is already in use; " + "stop the existing Mesh instance" + ) + command = l3_server_command( + Path(build["binary"]), model, cache_root, budget, minimum_free + ) + log_path.parent.mkdir(parents=True, exist_ok=True) + log_handle = log_path.open("wb") + process = subprocess.Popen( + command, + cwd=Path(build["worktree"]), + env=isolated_server_env(Path(build["runtime_root"]), state_dir, hf_home), + stdout=log_handle, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + log_handle.close() + return process, command + + +def management_json( + method: str, path: str, body: Optional[dict[str, Any]] = None +) -> dict[str, Any]: + connection = http.client.HTTPConnection( + MANAGEMENT_HOST, MANAGEMENT_PORT, timeout=30 + ) + encoded = json.dumps(body).encode() if body is not None else None + headers = {"Content-Type": "application/json"} if encoded is not None else {} + try: + connection.request(method, path, body=encoded, headers=headers) + response = connection.getresponse() + payload = response.read() + finally: + connection.close() + try: + document = json.loads(payload) + except json.JSONDecodeError as error: + raise RuntimeError( + f"management {method} {path} returned non-JSON HTTP {response.status}" + ) from error + if response.status < 200 or response.status >= 300: + raise RuntimeError( + f"management {method} {path} failed with HTTP {response.status}: {document}" + ) + if not isinstance(document, dict): + raise RuntimeError(f"management {method} {path} returned non-object JSON") + return document + + +def wait_for_l3_status( + process: subprocess.Popen[bytes], timeout: float +) -> dict[str, Any]: + deadline = time.monotonic() + timeout + last_error = "not ready" + while time.monotonic() < deadline: + if process.poll() is not None: + raise RuntimeError( + f"Mesh exited before L3 status readiness with status {process.returncode}" + ) + try: + status = management_json("GET", "/api/runtime/kv-cache") + if status.get("version") == 1: + return status + last_error = f"unexpected status version: {status.get('version')!r}" + except (OSError, RuntimeError) as error: + last_error = str(error) + time.sleep(0.5) + raise TimeoutError(f"L3 status did not become ready after {timeout}s: {last_error}") + + +def wait_for_committed_write( + process: subprocess.Popen[bytes], minimum_writes: int, timeout: float +) -> dict[str, Any]: + deadline = time.monotonic() + timeout + previous: Optional[tuple[int, int, int]] = None + while time.monotonic() < deadline: + status = wait_for_l3_status(process, min(timeout, 5)) + activity = status.get("activity") or {} + usage = status.get("usage") or {} + snapshot = ( + int(activity.get("writes") or 0), + int(usage.get("reserved_inflight_bytes") or 0), + int(usage.get("used_bytes") or 0), + ) + if ( + snapshot[0] >= minimum_writes + and snapshot[1] == 0 + and snapshot[2] > 0 + and snapshot == previous + ): + return status + previous = snapshot + time.sleep(0.5) + raise TimeoutError( + f"L3 write did not reach committed/stable state after {timeout}s" + ) + + +def clear_l3_until_empty(timeout: float) -> dict[str, Any]: + deadline = time.monotonic() + timeout + stable_empty = 0 + last: Optional[dict[str, Any]] = None + while time.monotonic() < deadline: + last = management_json("DELETE", "/api/runtime/kv-cache", {}) + manifests = ( + last.get("status", {}).get("usage", {}).get("manifests") + ) + if manifests == 0: + stable_empty += 1 + if stable_empty >= 4: + return last + else: + stable_empty = 0 + time.sleep(0.5) + raise TimeoutError(f"L3 clear did not remain empty after {timeout}s: {last}") + + +def activity_delta( + before: dict[str, Any], after: dict[str, Any] +) -> dict[str, int]: + before_activity = before.get("activity") or {} + after_activity = after.get("activity") or {} + fields = ( + "fills", + "hits", + "misses", + "writes", + "bytes_read", + "bytes_written", + "evictions", + "corrupt_entries", + ) + return { + field: int(after_activity.get(field) or 0) + - int(before_activity.get(field) or 0) + for field in fields + } + + +def final_checkpoint_request( + trajectory: dict[str, Any], + model_id: str, + max_output_tokens: int, + timeout: float, + phase: str, +) -> dict[str, Any]: + results = replay_trajectory( + trajectory, + model_id, + max_output_tokens, + timeout, + measured_assistant_turns={assistant_turn_count(trajectory) - 1}, + checkpoint_stage=phase, + ) + if len(results) != 1: raise RuntimeError( - f"inference server stopped but TCP {DEFAULT_PORT} is still occupied" + f"trajectory {trajectory['session_id']} produced {len(results)} final checkpoints" ) + return results[0] + + +def concurrent_final_checkpoints( + trajectory: dict[str, Any], + model_id: str, + max_output_tokens: int, + timeout: float, + phase: str, + requests: int, +) -> list[dict[str, Any]]: + barrier = threading.Barrier(requests) + + def run(index: int) -> dict[str, Any]: + barrier.wait(timeout=timeout) + result = final_checkpoint_request( + trajectory, model_id, max_output_tokens, timeout, phase + ) + result["request_id"] = f"{trajectory['session_id']}:{phase}:{index}" + return result + + with concurrent.futures.ThreadPoolExecutor(max_workers=requests) as pool: + return list(pool.map(run, range(requests))) def collect_runtime_logs(state_dir: Path, output_dir: Path) -> None: @@ -1377,7 +1645,10 @@ def run_arm_pass( finally: try: if process is not None: - stop_server(process) + stop_server( + process, + include_management_port=build.get("engine", "mesh") == "mesh", + ) finally: collect_runtime_logs(state_dir, pass_dir / "native-runtime") shutil.rmtree(state_dir, ignore_errors=True) @@ -2321,6 +2592,689 @@ def run_benchmark(args: argparse.Namespace) -> Path: return report +def l3_lifecycle_plan(args: argparse.Namespace, spec: RefSpec) -> dict[str, Any]: + prompt_min, prompt_max = parse_token_range(args.prompt_token_range) + return { + "schema_version": 1, + "kind": "disk-l3-lifecycle", + "repo": str(args.repo), + "ref": spec.__dict__, + "build_commands": [ + ["just", "release-host-build"], + ["just", "release-runtime-build", args.backend], + ], + "server_commands": { + "disk_off": [ + "", + "serve", + "--model", + args.model, + "--log-format", + "json", + ], + "disk_on": [ + "", + "serve", + "--model", + args.model, + "--log-format", + "json", + "--kv-cache-disk", + args.disk_budget, + "--kv-cache-disk-dir", + "", + "--kv-cache-min-free", + args.minimum_free, + ], + }, + "input": { + "trajectory_manifest": str(args.trajectory_manifest), + "lifecycle_cohort": args.lifecycle_cohort, + "required_source_datasets": args.require_source_dataset, + "high_load_cohorts": [str(value) for value in args.concurrency], + "prompt_token_range": [prompt_min, prompt_max], + "low_space_disk_budget": args.low_space_disk_budget, + "low_space_minimum_free": args.low_space_minimum_free, + }, + "phases": [ + "disk-off cold restarts", + "disk-on empty-root cold write", + "multi-turn growth across captured sessions and sources", + "same-process L1 repeat", + "post-restart first-request L3 samples", + "identical-prefix physical-fill wave", + "empty-root identical-prefix record wave", + "prune and clear during traffic", + "forced low-space cold fallback", + *( + f"captured high-load c{value} disk-off and disk-on" + for value in args.concurrency + ), + ], + "gates": { + "all_requests_succeed": True, + "greedy_seeded_output_sha256_matches_disk_off": True, + "post_restart_l3_ttft_p50_ratio_max": args.max_l3_ttft_ratio, + "identical_prefix_repeats": args.identical_repeats, + "physical_fill_delta": 1, + "physical_write_delta": 1, + "payload_write_amplification_max": args.max_payload_write_amplification, + "forced_low_space_state": "read_only_low_space", + "high_load_decode_inter_token_p99_regression_max_pct": ( + args.max_decode_p99_regression_pct + ), + }, + "outputs": [ + "raw request JSONL", + "phase status snapshots and counter deltas", + "server/build logs", + "versioned run JSON and Markdown report", + "SHA-256 inventory", + ], + } + + +def select_l3_trajectories( + cohorts: dict[str, list[dict[str, Any]]], args: argparse.Namespace +) -> list[dict[str, Any]]: + trajectories = cohorts[args.lifecycle_cohort] + if not args.require_source_dataset: + return [trajectories[0]] + selected: list[dict[str, Any]] = [] + for source in args.require_source_dataset: + match = next( + ( + trajectory + for trajectory in trajectories + if trajectory["source_dataset"] == source + ), + None, + ) + if match is None: + raise ValueError( + f"lifecycle cohort {args.lifecycle_cohort!r} is missing required " + f"source_dataset {source!r}" + ) + selected.append(match) + return selected + + +def l3_request_gate( + checks: list[dict[str, Any]], name: str, passed: bool, detail: str +) -> None: + checks.append({"name": name, "passed": passed, "detail": detail}) + + +def evaluate_l3_lifecycle_gates( + run: dict[str, Any], args: argparse.Namespace +) -> dict[str, Any]: + phases = run["phases"] + checks: list[dict[str, Any]] = [] + all_requests = [ + request + for phase in phases.values() + for request in phase.get("requests", []) + ] + failures = [request.get("request_id", "unknown") for request in all_requests if "error" in request] + l3_request_gate( + checks, + "all_requests_succeed", + not failures, + "all requests succeeded" if not failures else f"failed requests: {failures}", + ) + + baseline = phases["disk_off_cold"]["requests"] + baseline_hashes = { + (request["session_id"], request["assistant_turn"]): request.get( + "content_sha256" + ) + for request in baseline + } + mismatches = [ + request.get("request_id", "unknown") + for request in all_requests + if "error" not in request + and (request.get("session_id"), request.get("assistant_turn")) + in baseline_hashes + and request.get("content_sha256") + != baseline_hashes[ + (request["session_id"], request["assistant_turn"]) + ] + ] + l3_request_gate( + checks, + "output_identity", + not mismatches, + "all generated identities match disk-off" + if not mismatches + else f"mismatched requests: {mismatches}", + ) + + prompt_min, prompt_max = parse_token_range(args.prompt_token_range) + qualifying = [ + request + for request in baseline + phases["restart_l3"]["requests"] + if "error" not in request + ] + prompt_outliers = [ + request["request_id"] + for request in qualifying + if not prompt_min <= request["prompt_tokens"] <= prompt_max + ] + l3_request_gate( + checks, + "prompt_token_range", + not prompt_outliers, + f"all cold/restart prompts are within {prompt_min}:{prompt_max}" + if not prompt_outliers + else f"out-of-range requests: {prompt_outliers}", + ) + + cold_p50 = statistics.median( + request["ttft_seconds"] for request in baseline if "error" not in request + ) + restart_p50 = statistics.median( + request["ttft_seconds"] + for request in phases["restart_l3"]["requests"] + if "error" not in request + ) + ratio = restart_p50 / cold_p50 if cold_p50 else math.inf + l3_request_gate( + checks, + "post_restart_l3_ttft", + ratio <= args.max_l3_ttft_ratio, + f"restart/cold p50 ratio {ratio:.4f} (max {args.max_l3_ttft_ratio:.4f})", + ) + restart_deltas = phases["restart_l3"]["activity_deltas"] + l3_request_gate( + checks, + "every_restart_reads_l3", + bool(restart_deltas) + and all(delta["fills"] >= 1 and delta["bytes_read"] > 0 for delta in restart_deltas), + "each fresh process performed a physical L3 fill" + if restart_deltas + else "no restart samples recorded", + ) + + growth = phases["multi_turn_growth"] + growth_sources = {request["source_dataset"] for request in growth["requests"]} + l3_request_gate( + checks, + "multi_turn_growth", + growth["activity_delta"]["writes"] > 0 + and set(args.require_source_dataset).issubset(growth_sources), + f"writes={growth['activity_delta']['writes']} sources={sorted(growth_sources)}", + ) + + l1_delta = phases["same_process_l1"]["activity_delta"] + l3_request_gate( + checks, + "same_process_l1_avoids_disk", + l1_delta["fills"] == 0 and l1_delta["bytes_read"] == 0, + f"fills={l1_delta['fills']} bytes_read={l1_delta['bytes_read']}", + ) + fill_delta = phases["concurrent_fill"]["activity_delta"] + l3_request_gate( + checks, + "single_physical_fill", + fill_delta["fills"] == 1, + f"fills={fill_delta['fills']} across {args.identical_repeats} requests", + ) + record_delta = phases["concurrent_record"]["activity_delta"] + initial_delta = phases["disk_on_empty"]["activity_delta"] + allowed_write_bytes = math.ceil( + initial_delta["bytes_written"] * args.max_payload_write_amplification + ) + record_ok = ( + record_delta["writes"] == 1 + and record_delta["bytes_written"] <= allowed_write_bytes + ) + l3_request_gate( + checks, + "single_physical_write", + record_ok, + f"writes={record_delta['writes']} bytes={record_delta['bytes_written']} " + f"allowed={allowed_write_bytes}", + ) + low_space = phases["low_space"] + low_state = low_space["status_after"]["effective"]["state"] + low_success = all("error" not in request for request in low_space["requests"]) + l3_request_gate( + checks, + "low_space_falls_back_cold", + low_state == "read_only_low_space" + and low_success + and low_space["activity_delta"]["writes"] == 0, + f"effective_state={low_state} request_success={low_success} " + f"writes={low_space['activity_delta']['writes']}", + ) + lifecycle = phases["lifecycle_under_traffic"] + lifecycle_success = all("error" not in request for request in lifecycle["requests"]) + final_manifests = ( + lifecycle["final_clear"].get("status", {}).get("usage", {}).get("manifests") + ) + l3_request_gate( + checks, + "lifecycle_under_traffic", + lifecycle_success + and "prune" in lifecycle + and "clear" in lifecycle + and final_manifests == 0, + f"traffic_success={lifecycle_success} final_manifests={final_manifests}", + ) + + for concurrency in args.concurrency: + off = phases[f"high_load_off_c{concurrency}"]["summary"] + on = phases[f"high_load_on_c{concurrency}"]["summary"] + off_p99 = off.get("decode_inter_token_p99_seconds") + on_p99 = on.get("decode_inter_token_p99_seconds") + comparable = off_p99 not in (None, 0) and on_p99 is not None + regression = 100 * (on_p99 / off_p99 - 1) if comparable else math.inf + off_hashes = off.get("content_sha256_by_request", {}) + on_hashes = on.get("content_sha256_by_request", {}) + passed = ( + off["failed_requests"] == 0 + and on["failed_requests"] == 0 + and off_hashes == on_hashes + and regression <= args.max_decode_p99_regression_pct + ) + l3_request_gate( + checks, + f"high_load_c{concurrency}", + passed, + f"decode-gap p99 regression={regression:.3f}% output_match={off_hashes == on_hashes}", + ) + return { + "evaluated": True, + "passed": all(check["passed"] for check in checks), + "checks": checks, + "cold_ttft_p50_seconds": cold_p50, + "restart_l3_ttft_p50_seconds": restart_p50, + "restart_l3_ttft_ratio": ratio, + } + + +def write_l3_lifecycle_report(output: Path, run: dict[str, Any]) -> Path: + lines = [ + "# Disk L3 KV cache lifecycle certification", + "", + f"- Commit: `{run['build']['commit']}`", + f"- Model: `{run['config']['model']}`", + f"- Backend: `{run['config']['backend']}`", + f"- Result: **{'PASS' if run['gates']['passed'] else 'FAIL'}**", + f"- Cold TTFT p50: `{run['gates']['cold_ttft_p50_seconds']:.6f}s`", + f"- Restart L3 TTFT p50: `{run['gates']['restart_l3_ttft_p50_seconds']:.6f}s`", + f"- Restart/cold ratio: `{run['gates']['restart_l3_ttft_ratio']:.4f}`", + "", + "## Gates", + "", + ] + for check in run["gates"]["checks"]: + marker = "PASS" if check["passed"] else "FAIL" + lines.append(f"- **{marker}** `{check['name']}` — {check['detail']}") + report = output / "REPORT.md" + report.write_text("\n".join(lines) + "\n", encoding="utf-8") + inventory = [ + f"{sha256(path)} {path.relative_to(output).as_posix()}" + for path in sorted(item for item in output.rglob("*") if item.is_file()) + if path.name != "artifact-sha256.txt" + ] + (output / "artifact-sha256.txt").write_text( + "\n".join(inventory) + "\n", encoding="utf-8" + ) + return report + + +def run_l3_lifecycle(args: argparse.Namespace) -> Path: + args.repo = args.repo.resolve() + args.output = args.output.resolve() + args.trajectory_manifest = args.trajectory_manifest.resolve() + if args.hf_home is not None: + args.hf_home = args.hf_home.resolve() + spec = parse_single_ref_spec(args.repo, args.ref) + plan = l3_lifecycle_plan(args, spec) + args.output.mkdir(parents=True, exist_ok=False) + write_json(args.output / "plan.json", plan) + commands = CommandLog(args.output / "commands.jsonl") + expected_cohorts = [args.lifecycle_cohort, *(str(value) for value in args.concurrency)] + inputs = import_trajectory_manifest( + args.trajectory_manifest, args.output, expected_cohorts + ) + cohorts = load_trajectory_cohorts(Path(inputs["manifest"]), expected_cohorts) + lifecycle_trajectories = select_l3_trajectories(cohorts, args) + for concurrency in args.concurrency: + if len(cohorts[str(concurrency)]) < concurrency: + raise ValueError( + f"cohort {concurrency} has {len(cohorts[str(concurrency)])} " + f"trajectories; at least {concurrency} required" + ) + worktree_root = ( + args.worktree_root or (args.repo.parent / ".agentic-replay-worktrees") + ).resolve() + worktree = prepare_worktree(args.repo, worktree_root, spec) + build = build_ref( + spec, worktree, args.backend, args.output, commands, args.skip_build + ) + run: dict[str, Any] = { + "schema_version": 1, + "kind": "disk-l3-lifecycle", + "started_at": utc_now(), + "host": {"hostname": socket.gethostname(), "platform": sys.platform}, + "config": { + "model": args.model, + "backend": args.backend, + "disk_budget": args.disk_budget, + "minimum_free": args.minimum_free, + "low_space_disk_budget": args.low_space_disk_budget, + "prompt_token_range": args.prompt_token_range, + "cold_samples": args.cold_samples, + "restart_samples": args.restart_samples, + "identical_repeats": args.identical_repeats, + "concurrency": args.concurrency, + }, + "plan_sha256": stable_hash(plan), + "inputs": inputs, + "build": build, + "phases": {}, + } + run_path = args.output / "run.json" + scratch_root = Path(tempfile.mkdtemp(prefix="agentic-replay-l3-")) + cache_root = scratch_root / "cache" + cache_root.mkdir() + process: Optional[subprocess.Popen[bytes]] = None + server_index = 0 + + def serve( + cache: Optional[Path], + minimum_free: str = args.minimum_free, + budget: str = args.disk_budget, + ) -> tuple[str, Path]: + nonlocal process, server_index + server_index += 1 + state_dir = scratch_root / f"state-{server_index}" + state_dir.mkdir() + log_path = args.output / "logs" / f"server-{server_index}.log" + process, command = start_l3_server( + build, + args.model, + state_dir, + log_path, + args.hf_home, + cache, + budget, + minimum_free, + ) + model_id = wait_for_model(DEFAULT_BASE_URL, args.startup_timeout, process) + run.setdefault("servers", []).append( + { + "index": server_index, + "pid": process.pid, + "command": command, + "log": str(log_path), + } + ) + return model_id, state_dir + + def stop(state_dir: Path) -> None: + nonlocal process + if process is not None: + stop_server(process) + collect_runtime_logs( + state_dir, args.output / "native-runtime" / f"server-{server_index}" + ) + process = None + + def primary_request(model_id: str, phase: str) -> dict[str, Any]: + return final_checkpoint_request( + lifecycle_trajectories[0], + model_id, + args.max_output_tokens, + args.request_timeout, + phase, + ) + + try: + off_requests: list[dict[str, Any]] = [] + for sample in range(args.cold_samples): + model_id, state_dir = serve(None) + try: + for trajectory in lifecycle_trajectories: + request = final_checkpoint_request( + trajectory, + model_id, + args.max_output_tokens, + args.request_timeout, + f"disk_off_cold_{sample + 1}", + ) + request["request_id"] = ( + f"{trajectory['session_id']}:disk-off:{sample + 1}" + ) + off_requests.append(request) + finally: + stop(state_dir) + write_request_records( + args.output / "data/disk-off-cold.jsonl", off_requests, 1 + ) + run["phases"]["disk_off_cold"] = {"requests": off_requests} + write_json(run_path, run) + + model_id, state_dir = serve(cache_root) + growth_before = wait_for_l3_status(process, args.startup_timeout) + growth_requests = [ + request + for trajectory in lifecycle_trajectories + for request in replay_trajectory( + trajectory, + model_id, + args.max_output_tokens, + args.request_timeout, + checkpoint_stage="multi_turn_growth", + ) + ] + growth_after = wait_for_committed_write(process, 1, args.request_timeout) + run["phases"]["multi_turn_growth"] = { + "requests": growth_requests, + "status_before": growth_before, + "status_after": growth_after, + "activity_delta": activity_delta(growth_before, growth_after), + } + clear_l3_until_empty(args.request_timeout) + stop(state_dir) + + model_id, state_dir = serve(cache_root) + status_before = wait_for_l3_status(process, args.startup_timeout) + if (status_before.get("usage") or {}).get("manifests") != 0: + raise RuntimeError("disk-on empty-root phase started with cached manifests") + empty_request = primary_request(model_id, "disk_on_empty") + status_after = wait_for_committed_write(process, 1, args.request_timeout) + empty_phase = { + "requests": [empty_request], + "status_before": status_before, + "status_after": status_after, + "activity_delta": activity_delta(status_before, status_after), + } + run["phases"]["disk_on_empty"] = empty_phase + additional_requests = [ + final_checkpoint_request( + trajectory, + model_id, + args.max_output_tokens, + args.request_timeout, + "disk_on_additional_source", + ) + for trajectory in lifecycle_trajectories[1:] + ] + if additional_requests: + additional_after = wait_for_committed_write( + process, len(lifecycle_trajectories), args.request_timeout + ) + run["phases"]["disk_on_additional_sources"] = { + "requests": additional_requests, + "status_before": status_after, + "status_after": additional_after, + "activity_delta": activity_delta(status_after, additional_after), + } + status_after = additional_after + l1_before = status_after + l1_requests = [ + final_checkpoint_request( + trajectory, + model_id, + args.max_output_tokens, + args.request_timeout, + "same_process_l1", + ) + for trajectory in lifecycle_trajectories + ] + l1_after = wait_for_l3_status(process, args.request_timeout) + run["phases"]["same_process_l1"] = { + "requests": l1_requests, + "status_before": l1_before, + "status_after": l1_after, + "activity_delta": activity_delta(l1_before, l1_after), + } + stop(state_dir) + write_json(run_path, run) + + restart_requests: list[dict[str, Any]] = [] + restart_deltas: list[dict[str, int]] = [] + for sample in range(args.restart_samples): + model_id, state_dir = serve(cache_root) + before = wait_for_l3_status(process, args.startup_timeout) + try: + for trajectory in lifecycle_trajectories: + request = final_checkpoint_request( + trajectory, + model_id, + args.max_output_tokens, + args.request_timeout, + f"restart_l3_{sample + 1}", + ) + request["request_id"] = ( + f"{trajectory['session_id']}:restart:{sample + 1}" + ) + restart_requests.append(request) + after = wait_for_l3_status(process, args.request_timeout) + restart_deltas.append(activity_delta(before, after)) + finally: + stop(state_dir) + run["phases"]["restart_l3"] = { + "requests": restart_requests, + "activity_deltas": restart_deltas, + } + + model_id, state_dir = serve(cache_root) + before = wait_for_l3_status(process, args.startup_timeout) + fill_requests = concurrent_final_checkpoints( + lifecycle_trajectories[0], + model_id, + args.max_output_tokens, + args.request_timeout, + "concurrent_fill", + args.identical_repeats, + ) + after = wait_for_l3_status(process, args.request_timeout) + run["phases"]["concurrent_fill"] = { + "requests": fill_requests, + "status_before": before, + "status_after": after, + "activity_delta": activity_delta(before, after), + } + stop(state_dir) + + model_id, state_dir = serve(cache_root) + clear_l3_until_empty(args.request_timeout) + stop(state_dir) + model_id, state_dir = serve(cache_root) + before = wait_for_l3_status(process, args.startup_timeout) + record_requests = concurrent_final_checkpoints( + lifecycle_trajectories[0], + model_id, + args.max_output_tokens, + args.request_timeout, + "concurrent_record", + args.identical_repeats, + ) + after = wait_for_committed_write(process, 1, args.request_timeout) + run["phases"]["concurrent_record"] = { + "requests": record_requests, + "status_before": before, + "status_after": after, + "activity_delta": activity_delta(before, after), + } + stop(state_dir) + write_json(run_path, run) + + model_id, state_dir = serve(cache_root) + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + traffic = pool.submit(primary_request, model_id, "lifecycle_under_traffic") + time.sleep(0.05) + prune = management_json( + "POST", "/api/runtime/kv-cache/prune", {"target_bytes": 0} + ) + clear = management_json("DELETE", "/api/runtime/kv-cache", {}) + traffic_request = traffic.result() + final_clear = clear_l3_until_empty(args.request_timeout) + run["phases"]["lifecycle_under_traffic"] = { + "requests": [traffic_request], + "prune": prune, + "clear": clear, + "final_clear": final_clear, + } + stop(state_dir) + + model_id, state_dir = serve( + cache_root, args.low_space_minimum_free, args.low_space_disk_budget + ) + low_before = wait_for_l3_status(process, args.startup_timeout) + low_request = primary_request(model_id, "low_space") + low_after = wait_for_l3_status(process, args.request_timeout) + run["phases"]["low_space"] = { + "requests": [low_request], + "status_before": low_before, + "status_after": low_after, + "activity_delta": activity_delta(low_before, low_after), + } + stop(state_dir) + + for disk_enabled in (False, True): + model_id, state_dir = serve(cache_root if disk_enabled else None) + try: + for concurrency in args.concurrency: + phase = f"high_load_{'on' if disk_enabled else 'off'}_c{concurrency}" + raw_path = args.output / "data" / f"{phase}.jsonl" + summary = run_trajectory_cell( + trajectories=cohorts[str(concurrency)][:concurrency], + model_id=model_id, + concurrency=concurrency, + max_output_tokens=args.max_output_tokens, + timeout=args.request_timeout, + raw_path=raw_path, + replay_mode="final", + ) + requests = [ + json.loads(line) + for line in raw_path.read_text(encoding="utf-8").splitlines() + ] + run["phases"][phase] = { + "requests": requests, + "summary": summary, + } + finally: + stop(state_dir) + run["completed_at"] = utc_now() + run["gates"] = evaluate_l3_lifecycle_gates(run, args) + write_json(run_path, run) + report = write_l3_lifecycle_report(args.output, run) + if not run["gates"]["passed"]: + raise RuntimeError(f"disk L3 lifecycle gates failed; see {report}") + return report + finally: + if process is not None: + stop_server(process) + shutil.rmtree(scratch_root, ignore_errors=True) + + def add_common_arguments(parser: argparse.ArgumentParser) -> None: parser.add_argument("--repo", type=Path, default=REPO) parser.add_argument( @@ -2415,6 +3369,96 @@ def add_common_arguments(parser: argparse.ArgumentParser) -> None: ) +def add_l3_arguments(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--repo", type=Path, default=REPO) + parser.add_argument( + "--ref", + required=True, + help="single LABEL=GIT_REF release candidate", + ) + parser.add_argument("--model", required=True, help="model URI or local package path") + parser.add_argument("--backend", default="metal") + parser.add_argument("--trajectory-manifest", type=Path, required=True) + parser.add_argument("--lifecycle-cohort", default="l3") + parser.add_argument( + "--require-source-dataset", + action="append", + default=[], + help="repeat for required captured Buzz/OpenCode/Goose source names", + ) + parser.add_argument( + "--concurrency", + type=int, + action="append", + default=[], + help="captured high-load cohort; defaults to c64/c128/c256", + ) + parser.add_argument("--prompt-token-range", default="18000:24000") + parser.add_argument("--cold-samples", type=int, default=3) + parser.add_argument("--restart-samples", type=int, default=3) + parser.add_argument("--identical-repeats", type=int, default=100) + parser.add_argument("--max-output-tokens", type=int, default=2048) + parser.add_argument("--disk-budget", default="auto") + parser.add_argument("--low-space-disk-budget", default="32GiB") + parser.add_argument("--minimum-free", default="1GiB") + parser.add_argument( + "--low-space-minimum-free", + default="1TiB", + help="minimum-free threshold used to force read-only cold fallback", + ) + parser.add_argument("--max-l3-ttft-ratio", type=float, default=0.5) + parser.add_argument( + "--max-payload-write-amplification", type=float, default=1.2 + ) + parser.add_argument( + "--max-decode-p99-regression-pct", type=float, default=5.0 + ) + + +def validate_l3_args(args: argparse.Namespace, parser: argparse.ArgumentParser) -> None: + if not args.concurrency: + args.concurrency = [64, 128, 256] + positive = ( + args.cold_samples, + args.restart_samples, + args.identical_repeats, + args.max_output_tokens, + ) + if any(value <= 0 for value in positive): + parser.error("L3 sample, repeat, and output sizes must be positive") + if len(set(args.concurrency)) != len(args.concurrency) or any( + value <= 0 for value in args.concurrency + ): + parser.error("--concurrency values must be unique and positive") + if len(set(args.require_source_dataset)) != len(args.require_source_dataset): + parser.error("--require-source-dataset values must be unique") + if args.command == "l3-run" and len(args.require_source_dataset) < 3: + parser.error( + "l3-run requires three --require-source-dataset values for the " + "captured Buzz, OpenCode, and Goose workloads" + ) + try: + parse_token_range(args.prompt_token_range) + except ValueError as error: + parser.error(str(error)) + iec = re.compile(r"[1-9][0-9]*(?:KiB|MiB|GiB|TiB)\Z") + if args.disk_budget != "auto" and iec.fullmatch(args.disk_budget) is None: + parser.error("--disk-budget must be auto or a positive IEC size") + for option in ( + "minimum_free", + "low_space_minimum_free", + "low_space_disk_budget", + ): + if iec.fullmatch(getattr(args, option)) is None: + parser.error(f"--{option.replace('_', '-')} must be a positive IEC size") + if not 0 < args.max_l3_ttft_ratio <= 1: + parser.error("--max-l3-ttft-ratio must be greater than 0 and at most 1") + if args.max_payload_write_amplification < 1: + parser.error("--max-payload-write-amplification must be at least 1") + if args.max_decode_p99_regression_pct < 0: + parser.error("--max-decode-p99-regression-pct cannot be negative") + + def validate_args(args: argparse.Namespace, parser: argparse.ArgumentParser) -> None: if args.command == "run" and ( (args.dataset_file is None) == (args.trajectory_manifest is None) @@ -2503,9 +3547,29 @@ def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: run.add_argument("--resume", action="store_true") report = subparsers.add_parser("report", help="rerender tables and charts from run.json") report.add_argument("--artifact", type=Path, required=True) + l3_plan = subparsers.add_parser( + "l3-plan", help="print the side-effect-free disk-L3 certification plan" + ) + add_l3_arguments(l3_plan) + l3_run = subparsers.add_parser( + "l3-run", help="build and execute disk-L3 lifecycle certification" + ) + add_l3_arguments(l3_run) + l3_run.add_argument("--output", type=Path, required=True) + l3_run.add_argument("--worktree-root", type=Path) + l3_run.add_argument("--hf-home", type=Path) + l3_run.add_argument("--startup-timeout", type=float, default=1800) + l3_run.add_argument("--request-timeout", type=float, default=900) + l3_run.add_argument("--skip-build", action="store_true") + l3_report = subparsers.add_parser( + "l3-report", help="rerender a disk-L3 lifecycle report from run.json" + ) + l3_report.add_argument("--artifact", type=Path, required=True) args = parser.parse_args(argv) if args.command in {"plan", "run"}: validate_args(args, parser) + elif args.command in {"l3-plan", "l3-run"}: + validate_l3_args(args, parser) return args @@ -2525,9 +3589,21 @@ def main(argv: Optional[Sequence[str]] = None) -> int: if args.command == "run": print(run_benchmark(args)) return 0 + if args.command == "l3-plan": + args.repo = args.repo.resolve() + args.trajectory_manifest = args.trajectory_manifest.resolve() + spec = parse_single_ref_spec(args.repo, args.ref) + print(json.dumps(l3_lifecycle_plan(args, spec), indent=2, sort_keys=True)) + return 0 + if args.command == "l3-run": + print(run_l3_lifecycle(args)) + return 0 artifact = args.artifact.resolve() document = json.loads((artifact / "run.json").read_text(encoding="utf-8")) - print(write_report(artifact, document)) + if args.command == "l3-report": + print(write_l3_lifecycle_report(artifact, document)) + else: + print(write_report(artifact, document)) return 0 diff --git a/evals/kv-restart-replay.py b/evals/kv-restart-replay.py new file mode 100644 index 0000000000..21410078c4 --- /dev/null +++ b/evals/kv-restart-replay.py @@ -0,0 +1,682 @@ +#!/usr/bin/env python3 +"""KV restart replay: measure serving latency across a process restart. + +Issue #1647-A. Runs one frozen multi-turn conversation against a default-startup +``mesh-llm serve`` in three cohorts: + +- ``fill`` — cold server, conversation grows turn by turn (prefix reuse). +- ``restore`` — the server is stopped and restarted on the same state directory, + then the frozen full conversation is replayed. On a build with a + durable KV tier this measures first-request-after-restart + restoration; without one it is the cold-prefill reference. +- ``warm`` — repeat replays without restart (resident reuse reference). + +The runner never sets context size, lanes, KV budget, or backend tuning; the +only serving arguments are ``--model``, ``--log-format json`` and the explicit +``--serve-extra-args`` pass-through an operator asks for (for example a +``--kv-cache-disk`` mode under test). Everything measured is observational: +streaming TTFT from the first chunk, usage from ``stream_options.include_usage``, +cached tokens from ``prompt_tokens_details``. + +Artifacts: ``run.json`` (schema_version 1, provenance + config + cohort +summaries), ``requests.jsonl`` (one row per request), ``report.md`` (human +summary). The manifest itself is deterministic from ``--turns`` and +``--turn-target-tokens`` and is embedded (with its SHA-256) into ``run.json``. +""" + +from __future__ import annotations + +import argparse +import hashlib +import http.client +import json +import math +import os +import re +import signal +import socket +import statistics +import subprocess +import sys +import time +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Optional, Sequence + +REPO = Path(__file__).resolve().parents[1] +DEFAULT_BASE_URL = "http://127.0.0.1:9337/v1" +DEFAULT_HOST = "127.0.0.1" +DEFAULT_PORT = 9337 +SCHEMA_VERSION = 1 + +# Same discipline as evals/agentic-replay.py: a default-startup benchmark may +# not tune the server. Extra serving arguments must arrive explicitly via +# --serve-extra-args and are recorded verbatim in run.json. +FORBIDDEN_STARTUP_OPTIONS = ( + "--ctx-size", + "--generation-concurrency", + "--generation-queue-capacity", + "--host", + "--max-vram", + "--parallel", + "--port", +) + +# Deterministic manifest vocabulary. The conversation simulates a long-running +# coding-agent session: a stable scaffold, a growing project brief, and a +# per-turn request. Content is drawn from a fixed word list with a fixed PRNG +# seed so the same settings always produce the same conversation. +SEED = 20260909 +_VOCAB = ( + "cache prefix token restore restart segment manifest budget eviction " + "prefill decode latency throughput checkpoint durable radix tier admission " + "pipeline stream verify digest commit quarantine pin lease reserve node " + "mesh relay model runtime kernel attention matrix layer head batch queue " + "trace replay harness baseline cohort percentile regression gate promote" +).split() + + +def utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def stable_hash(value: Any) -> str: + return hashlib.sha256( + json.dumps(value, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() + + +# --------------------------------------------------------------------------- +# Manifest +# --------------------------------------------------------------------------- + + +class DeterministicRandom: + """Small LCG so manifests do not depend on the host Python random module.""" + + def __init__(self, seed: int) -> None: + self.state = seed & 0xFFFFFFFFFFFF + + def next(self) -> int: + self.state = (self.state * 25214903917 + 11) & 0xFFFFFFFFFFFF + return self.state >> 16 + + def below(self, bound: int) -> int: + return self.next() % bound + + def words(self, count: int) -> list[str]: + return [_VOCAB[self.below(len(_VOCAB))] for _ in range(count)] + + +def build_manifest(turns: int, turn_target_tokens: int, system_tokens: int) -> dict[str, Any]: + """Build the frozen conversation. Turn sizes are approximate (words * 4/3); + the authoritative prompt token counts come from server usage at run time.""" + + rng = DeterministicRandom(SEED) + + def block(target_tokens: int, topic: str) -> str: + words = max(1, int(target_tokens * 3 / 4)) + chunks = [] + while len(chunks) * 8 < words: + chunks.append(" ".join(rng.words(8))) + return f"[{topic}] " + " ".join(chunks) + + scaffold = block(system_tokens, "scaffold") + turn_specs = [] + for index in range(turns): + body = block(turn_target_tokens, f"turn-{index + 1}-context") + request = ( + f"Turn {index + 1}: given the project brief above, summarize the " + f"{' '.join(rng.words(6))} constraint in one sentence and list the " + f"{' '.join(rng.words(4))} next step." + ) + response = ( + f"Turn {index + 1} answer: preserve the {' '.join(rng.words(5))} " + f"constraint. Next step: verify {' '.join(rng.words(4))}." + ) + turn_specs.append({"context": body, "request": request, "response": response}) + + return { + "schema_version": SCHEMA_VERSION, + "kind": "kv-restart-replay/manifest", + "seed": SEED, + "settings": { + "turns": turns, + "turn_target_tokens": turn_target_tokens, + "system_tokens": system_tokens, + "approx_total_prompt_tokens": system_tokens + turns * turn_target_tokens, + }, + "system": scaffold, + "turns": turn_specs, + } + + +# --------------------------------------------------------------------------- +# Server lifecycle (mirrors evals/agentic-replay.py) +# --------------------------------------------------------------------------- + + +def server_command(binary: Path, model: str, extra_args: Sequence[str]) -> list[str]: + command = [str(binary), "serve", "--model", model, "--log-format", "json"] + command.extend(extra_args) + for argument in command: + for option in FORBIDDEN_STARTUP_OPTIONS: + if argument == option or argument.startswith(f"{option}="): + raise AssertionError(f"default-startup benchmark cannot use {option}") + return command + + +def port_is_open(host: str = DEFAULT_HOST, port: int = DEFAULT_PORT) -> bool: + with socket.socket() as connection: + connection.settimeout(0.2) + return connection.connect_ex((host, port)) == 0 + + +def wait_for_model(timeout: float, process: subprocess.Popen[bytes]) -> str: + deadline = time.monotonic() + timeout + last_error = "not ready" + while time.monotonic() < deadline: + if process.poll() is not None: + raise RuntimeError(f"Mesh exited before readiness with status {process.returncode}") + connection = http.client.HTTPConnection(DEFAULT_HOST, DEFAULT_PORT, timeout=5) + try: + connection.request("GET", "/v1/models") + response = connection.getresponse() + body = response.read() + if response.status == 200: + document = json.loads(body) + models = document.get("data") or [] + if models: + return models[0]["id"] + last_error = f"HTTP {response.status}: {body[:300]!r}" + except (OSError, json.JSONDecodeError) as error: + last_error = str(error) + finally: + connection.close() + time.sleep(1) + raise TimeoutError(f"Mesh did not become ready after {timeout}s: {last_error}") + + +def stop_server(process: subprocess.Popen[bytes]) -> None: + if process.poll() is None: + os.killpg(process.pid, signal.SIGINT) + try: + process.wait(timeout=30) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGTERM) + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGKILL) + process.wait(timeout=10) + deadline = time.monotonic() + 10 + while port_is_open() and time.monotonic() < deadline: + time.sleep(0.2) + if port_is_open(): + raise RuntimeError("Mesh stopped but the serving port is still occupied") + + +def isolated_server_env(state_dir: Path) -> dict[str, str]: + env = os.environ.copy() + home = state_dir / "home" + home.mkdir(parents=True, exist_ok=True) + env.update( + { + "HOME": str(home), + "XDG_CACHE_HOME": str(state_dir / "xdg-cache"), + "XDG_CONFIG_HOME": str(state_dir / "xdg-config"), + "MESH_LLM_RUNTIME_ROOT": str(state_dir / "runtime"), + } + ) + if "HF_HOME" not in env: + env["HF_HOME"] = str(Path.home() / ".cache/huggingface") + return env + + +def start_server( + binary: Path, + model: str, + extra_args: Sequence[str], + state_dir: Path, + log_path: Path, +) -> tuple[subprocess.Popen[bytes], list[str]]: + if port_is_open(): + raise RuntimeError(f"TCP {DEFAULT_PORT} is already in use; stop the existing Mesh instance") + command = server_command(binary, model, extra_args) + log_path.parent.mkdir(parents=True, exist_ok=True) + log_handle = log_path.open("wb") + process = subprocess.Popen( + command, + cwd=str(REPO), + env=isolated_server_env(state_dir), + stdout=log_handle, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + log_handle.close() + return process, command + + +# --------------------------------------------------------------------------- +# Requests (mirrors evals/agentic-replay.py stream_request) +# --------------------------------------------------------------------------- + + +def stream_request( + request_id: str, + messages: Sequence[dict[str, Any]], + model_id: str, + max_output_tokens: int, + timeout: float, +) -> dict[str, Any]: + started = time.monotonic() + first_token_at: Optional[float] = None + completion_tokens = 0 + prompt_tokens = 0 + cached_tokens = 0 + saw_prompt_tokens = False + saw_cached_tokens = False + saw_done = False + connection = http.client.HTTPConnection(DEFAULT_HOST, DEFAULT_PORT, timeout=timeout) + payload = { + "model": model_id, + "messages": list(messages), + "max_tokens": max_output_tokens, + "temperature": 0, + "seed": 42, + "stream": True, + "stream_options": {"include_usage": True}, + } + try: + connection.request( + "POST", + "/v1/chat/completions", + json.dumps(payload), + {"Content-Type": "application/json", "Authorization": "Bearer EMPTY"}, + ) + response = connection.getresponse() + if response.status != 200: + body = response.read(4096).decode("utf-8", errors="replace") + return {"request_id": request_id, "error": f"HTTP {response.status}: {body}"} + for raw_line in response: + line = raw_line.strip() + if not line.startswith(b"data: "): + continue + event_bytes = line[6:] + if event_bytes == b"[DONE]": + saw_done = True + break + try: + event = json.loads(event_bytes) + except json.JSONDecodeError: + continue + server_error = event.get("error") + if server_error is not None: + return { + "request_id": request_id, + "error": f"stream failed with server error: {server_error}", + } + usage = event.get("usage") + if isinstance(usage, dict): + completion_tokens = int(usage.get("completion_tokens") or completion_tokens) + if "prompt_tokens" in usage and usage["prompt_tokens"] is not None: + prompt_tokens = int(usage["prompt_tokens"]) + saw_prompt_tokens = True + details = usage.get("prompt_tokens_details") + if ( + isinstance(details, dict) + and "cached_tokens" in details + and details["cached_tokens"] is not None + ): + cached_tokens = int(details["cached_tokens"]) + saw_cached_tokens = True + choices = event.get("choices") + if not isinstance(choices, list) or not choices: + continue + delta = choices[0].get("delta") if isinstance(choices[0], dict) else None + if isinstance(delta, dict) and delta.get("content") and first_token_at is None: + first_token_at = time.monotonic() + if first_token_at is None: + return {"request_id": request_id, "error": "stream completed without content tokens"} + if not saw_done: + return {"request_id": request_id, "error": "stream ended without terminal [DONE] marker"} + if not saw_prompt_tokens: + return {"request_id": request_id, "error": "stream completed without prompt token usage"} + if not saw_cached_tokens: + return {"request_id": request_id, "error": "stream completed without cached token usage"} + ended = time.monotonic() + return { + "request_id": request_id, + "ttft_seconds": first_token_at - started, + "total_seconds": ended - started, + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "cached_tokens": cached_tokens, + "decode_tokens_per_second": ( + completion_tokens / (ended - first_token_at) if ended > first_token_at else None + ), + } + except (OSError, TimeoutError, http.client.HTTPException) as error: + return {"request_id": request_id, "error": str(error)} + finally: + connection.close() + + +# --------------------------------------------------------------------------- +# Provenance +# --------------------------------------------------------------------------- + + +def hardware_fingerprint() -> dict[str, Any]: + fingerprint: dict[str, Any] = { + "platform": sys.platform, + "python": sys.version.split()[0], + "hostname": socket.gethostname(), + } + if sys.platform == "darwin": + try: + fingerprint["chip"] = ( + subprocess.run( + ["sysctl", "-n", "machdep.cpu.brand_string"], + capture_output=True, + text=True, + check=True, + ).stdout.strip() + ) + fingerprint["machine_model"] = ( + subprocess.run( + ["sysctl", "-n", "hw.model"], capture_output=True, text=True, check=True + ).stdout.strip() + ) + except (OSError, subprocess.CalledProcessError): + pass + try: + memory = subprocess.run( + ["sysctl", "-n", "hw.memsize"], + capture_output=True, + text=True, + check=True, + ) + fingerprint["physical_memory_bytes"] = int(memory.stdout.strip()) + except (OSError, subprocess.CalledProcessError, ValueError): + fingerprint["physical_memory_bytes"] = None + fingerprint["cpu_core_count"] = os.cpu_count() + else: + fingerprint["cpu_core_count"] = os.cpu_count() + try: + meminfo = Path("/proc/meminfo").read_text(encoding="utf-8") + match = re.search(r"MemTotal:\s+(\d+)\s+kB", meminfo) + if match: + fingerprint["physical_memory_bytes"] = int(match.group(1)) * 1024 + except OSError: + pass + return fingerprint + + +def binary_provenance(binary: Path) -> dict[str, Any]: + provenance: dict[str, Any] = {"binary": str(binary), "binary_sha256": sha256_file(binary)} + try: + described = subprocess.run( + ["git", "describe", "--always", "--dirty", "--tags"], + cwd=str(REPO), + capture_output=True, + text=True, + check=True, + ) + provenance["git_describe"] = described.stdout.strip() + commit = subprocess.run( + ["git", "rev-parse", "HEAD"], cwd=str(REPO), capture_output=True, text=True, check=True + ) + provenance["source_sha"] = commit.stdout.strip() + except (subprocess.CalledProcessError, OSError): + provenance["git_describe"] = "unknown" + provenance["source_sha"] = "unknown" + return provenance + + +# --------------------------------------------------------------------------- +# Cohorts and run +# --------------------------------------------------------------------------- + + +def summarize_cohort(name: str, rows: Sequence[dict[str, Any]]) -> dict[str, Any]: + successful = [row for row in rows if "error" not in row] + failed = [row for row in rows if "error" in row] + ttft = [row["ttft_seconds"] for row in successful] + + def percentile(values: Sequence[float], fraction: float) -> Optional[float]: + if not values: + return None + ordered = sorted(values) + index = min(math.ceil(len(ordered) * fraction) - 1, len(ordered) - 1) + return ordered[max(index, 0)] + + prompt_tokens = sum(row.get("prompt_tokens", 0) for row in successful) + cached_tokens = sum(row.get("cached_tokens", 0) for row in successful) + decode = [row["decode_tokens_per_second"] for row in successful if row.get("decode_tokens_per_second")] + return { + "cohort": name, + "requests": len(rows), + "failed": len(failed), + "ttft_p50_seconds": percentile(ttft, 0.50), + "ttft_p95_seconds": percentile(ttft, 0.95) if len(ttft) > 1 else None, + "total_seconds_mean": statistics.fmean(row["total_seconds"] for row in successful) if successful else None, + "prompt_tokens": prompt_tokens, + "cached_tokens": cached_tokens, + "cache_pct": (100 * cached_tokens / prompt_tokens) if prompt_tokens else None, + "decode_tokens_per_second_mean": statistics.fmean(decode) if decode else None, + } + + +def messages_through(messages: list[dict[str, Any]], turn_index: int) -> list[dict[str, Any]]: + """Conversation prefix ending on the user message of ``turn_index`` (0-based). + + Fill requests must look like the requests a real session produces: the last + message is always the user turn being answered, never the canned assistant + reply that follows it in the canonical conversation. + """ + return messages[: 2 * (turn_index + 1)] + + +def run_arm(args: argparse.Namespace, output: Path) -> dict[str, Any]: + binary = Path(args.binary).resolve() + model_path = Path(args.model).resolve() + if not binary.exists(): + raise FileNotFoundError(f"binary not found: {binary}") + if not model_path.exists(): + raise FileNotFoundError(f"model not found: {model_path}") + if args.turns < 1: + raise ValueError("turns must be at least 1") + if args.restore_repeats < 1: + raise ValueError("restore-repeats must be at least 1") + + manifest = build_manifest(args.turns, args.turn_target_tokens, args.system_tokens) + manifest_sha = stable_hash(manifest) + conversation: list[dict[str, Any]] = [{"role": "system", "content": manifest["system"]}] + for spec in manifest["turns"]: + conversation.append({"role": "user", "content": f"{spec['context']}\n\n{spec['request']}"}) + conversation.append({"role": "assistant", "content": spec["response"]}) + frozen_prompt = messages_through(conversation, args.turns - 1) + + state_dir = (output / "server-state").resolve() + state_dir.mkdir(parents=True, exist_ok=True) + requests_path = output / "requests.jsonl" + rows: list[dict[str, Any]] = [] + + def record(cohort: str, index: int, result: dict[str, Any]) -> None: + row = {"cohort": cohort, "request_index": index, **result} + rows.append(row) + with requests_path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(row, sort_keys=True) + "\n") + + def replay_frozen(cohort: str, model_id: str, repeats: Optional[int] = None) -> None: + for repeat in range(args.restore_repeats if repeats is None else repeats): + result = stream_request( + f"{cohort}-{repeat + 1}", + frozen_prompt, + model_id, + args.max_output_tokens, + args.request_timeout, + ) + record(cohort, repeat, result) + + provenance = { + "schema_version": SCHEMA_VERSION, + "kind": "kv-restart-replay/run", + "started_at": utc_now(), + "binary": binary_provenance(binary), + "model": { + "path": str(model_path), + "sha256": sha256_file(model_path), + "size_bytes": model_path.stat().st_size, + }, + "hardware": hardware_fingerprint(), + "config": { + "base_url": DEFAULT_BASE_URL, + "turns": args.turns, + "turn_target_tokens": args.turn_target_tokens, + "system_tokens": args.system_tokens, + "restore_repeats": args.restore_repeats, + "max_output_tokens": args.max_output_tokens, + "request_timeout": args.request_timeout, + "serve_extra_args": list(args.serve_extra_args), + }, + "manifest_sha256": manifest_sha, + "manifest": manifest, + } + + process = None + try: + # Cohort: fill — cold server, conversation grows turn by turn. + process, command = start_server( + binary, str(model_path), args.serve_extra_args, state_dir, output / "logs" / "fill.log" + ) + provenance["serve_command"] = command + model_id = wait_for_model(args.ready_timeout, process) + for index in range(args.turns): + result = stream_request( + f"fill-{index + 1}", + messages_through(conversation, index), + model_id, + args.max_output_tokens, + args.request_timeout, + ) + record("fill", index, result) + + # Cohort: restore — full process restart on the same state directory. + stop_server(process) + process = None + stopped_at = time.monotonic() + process, _ = start_server( + binary, str(model_path), args.serve_extra_args, state_dir, output / "logs" / "restore.log" + ) + model_id = wait_for_model(args.ready_timeout, process) + restart_gap_seconds = time.monotonic() - stopped_at + provenance["restart"] = { + "method": "SIGINT to the serving process group, then fresh start on the same state directory", + "restart_to_ready_seconds": restart_gap_seconds, + } + # Cohort: restore — the FIRST post-restart replay alone is the + # first-request-after-restart measurement; later replays warm from the + # resident cache and are recorded under the warm cohort instead. + replay_frozen("restore", model_id, repeats=1) + replay_frozen("warm", model_id, repeats=max(args.restore_repeats - 1, 0)) + finally: + if process is not None: + stop_server(process) + + provenance["cohorts"] = [ + summarize_cohort("fill", [row for row in rows if row["cohort"] == "fill"]), + summarize_cohort("restore", [row for row in rows if row["cohort"] == "restore"]), + summarize_cohort("warm", [row for row in rows if row["cohort"] == "warm"]), + ] + provenance["completed_at"] = utc_now() + return provenance + + +def write_report(run: dict[str, Any], path: Path) -> None: + lines = [ + "# KV restart replay", + "", + f"- source: `{run['binary'].get('source_sha', 'unknown')}` (`{run['binary'].get('git_describe', '?')}`)", + f"- model sha256: `{run['model']['sha256'][:16]}…`", + f"- manifest sha256: `{run['manifest_sha256'][:16]}…` ({run['config']['turns']} turns, " + f"≈{run['manifest']['settings']['approx_total_prompt_tokens']} prompt tokens)", + f"- serve extra args: `{run['config']['serve_extra_args'] or 'none'}`", + f"- restart-to-ready: {run.get('restart', {}).get('restart_to_ready_seconds', float('nan')):.1f}s", + "", + "| cohort | requests | failed | TTFT p50 (s) | TTFT p95 (s) | cached % | decode tok/s |", + "|---|---:|---:|---:|---:|---:|---:|", + ] + for cohort in run["cohorts"]: + fmt = lambda value: "—" if value is None else f"{value:.3f}" if isinstance(value, float) else value + lines.append( + f"| {cohort['cohort']} | {cohort['requests']} | {cohort['failed']} " + f"| {fmt(cohort['ttft_p50_seconds'])} | {fmt(cohort['ttft_p95_seconds'])} " + f"| {fmt(cohort['cache_pct'])} | {fmt(cohort['decode_tokens_per_second_mean'])} |" + ) + lines += [ + "", + "`restore` is the first-request-after-restart cohort; on a build without a", + "durable KV tier it is the cold-prefill reference. `warm` repeats the same", + "replay without restart (resident reuse reference).", + "", + ] + path.write_text("\n".join(lines), encoding="utf-8") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--binary", default=str(REPO / "target/release/mesh-llm")) + parser.add_argument("--model", required=True, help="path to a GGUF model file") + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--turns", type=int, default=4) + parser.add_argument("--turn-target-tokens", type=int, default=4750) + parser.add_argument("--system-tokens", type=int, default=500) + parser.add_argument( + "--restore-repeats", + type=int, + default=3, + help="total post-restart requests: one restore followed by warm repeats", + ) + parser.add_argument("--max-output-tokens", type=int, default=256) + parser.add_argument("--request-timeout", type=float, default=900.0) + parser.add_argument("--ready-timeout", type=float, default=900.0) + parser.add_argument( + "--serve-extra-args", + nargs=argparse.REMAINDER, + default=[], + help="explicit extra serving arguments (recorded verbatim in run.json)", + ) + args = parser.parse_args() + + output = args.output.resolve() + if output.exists() and any(output.iterdir()): + raise SystemExit(f"output directory is not empty: {output}") + output.mkdir(parents=True, exist_ok=True) + + run = run_arm(args, output) + write_json = output / "run.json" + write_json.write_text(json.dumps(run, indent=2, sort_keys=True) + "\n", encoding="utf-8") + write_report(run, output / "report.md") + print(f"wrote {write_json}") + for cohort in run["cohorts"]: + p50 = cohort["ttft_p50_seconds"] + print( + f" {cohort['cohort']:8s} p50={p50 if p50 is None else round(p50, 3)}s " + f"cache={cohort['cache_pct'] if cohort['cache_pct'] is None else round(cohort['cache_pct'], 1)}% " + f"failed={cohort['failed']}" + ) + return 0 if all(cohort["failed"] == 0 for cohort in run["cohorts"]) else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/evals/test_agentic_replay_l3.py b/evals/test_agentic_replay_l3.py new file mode 100644 index 0000000000..f283e7b6bf --- /dev/null +++ b/evals/test_agentic_replay_l3.py @@ -0,0 +1,129 @@ +import copy +import importlib.util +import sys +import unittest +from pathlib import Path +from types import SimpleNamespace + + +SCRIPT = Path(__file__).with_name("agentic-replay.py") +SPEC = importlib.util.spec_from_file_location("agentic_replay", SCRIPT) +AGENTIC_REPLAY = importlib.util.module_from_spec(SPEC) +assert SPEC.loader is not None +sys.modules[SPEC.name] = AGENTIC_REPLAY +SPEC.loader.exec_module(AGENTIC_REPLAY) + + +def activity(**changes): + snapshot = { + "fills": 0, + "hits": 0, + "misses": 0, + "writes": 0, + "bytes_read": 0, + "bytes_written": 0, + "evictions": 0, + "corrupt_entries": 0, + } + snapshot.update(changes) + return snapshot + + +def request(request_id="s:0", ttft=2.0): + return { + "request_id": request_id, + "session_id": "s", + "source_dataset": "buzz", + "assistant_turn": 0, + "prompt_tokens": 19_000, + "ttft_seconds": ttft, + "content_sha256": "same", + } + + +def passing_run(): + baseline = request() + high_load = { + "failed_requests": 0, + "decode_inter_token_p99_seconds": 0.1, + "content_sha256_by_request": {"s:0": "same"}, + } + return { + "phases": { + "disk_off_cold": {"requests": [baseline]}, + "disk_on_empty": { + "requests": [baseline], + "activity_delta": activity(writes=1, bytes_written=100), + }, + "multi_turn_growth": { + "requests": [baseline], + "activity_delta": activity(writes=1), + }, + "same_process_l1": { + "requests": [baseline], + "activity_delta": activity(), + }, + "restart_l3": { + "requests": [request("s:restart", 0.5)], + "activity_deltas": [activity(fills=1, bytes_read=100)], + }, + "concurrent_fill": { + "requests": [baseline], + "activity_delta": activity(fills=1, bytes_read=100), + }, + "concurrent_record": { + "requests": [baseline], + "activity_delta": activity(writes=1, bytes_written=100), + }, + "low_space": { + "requests": [baseline], + "status_after": {"effective": {"state": "read_only_low_space"}}, + "activity_delta": activity(), + }, + "lifecycle_under_traffic": { + "requests": [baseline], + "prune": {}, + "clear": {}, + "final_clear": {"status": {"usage": {"manifests": 0}}}, + }, + "high_load_off_c1": { + "requests": [baseline], + "summary": high_load, + }, + "high_load_on_c1": { + "requests": [baseline], + "summary": dict(high_load, decode_inter_token_p99_seconds=0.104), + }, + } + } + + +ARGS = SimpleNamespace( + prompt_token_range="18000:24000", + max_l3_ttft_ratio=0.5, + identical_repeats=100, + max_payload_write_amplification=1.2, + require_source_dataset=["buzz"], + concurrency=[1], + max_decode_p99_regression_pct=5.0, +) + + +class DiskL3LifecycleGateTests(unittest.TestCase): + def test_complete_evidence_passes(self): + gates = AGENTIC_REPLAY.evaluate_l3_lifecycle_gates(passing_run(), ARGS) + self.assertTrue(gates["passed"], gates) + + def test_duplicate_physical_fill_fails_closed(self): + run = copy.deepcopy(passing_run()) + run["phases"]["concurrent_fill"]["activity_delta"]["fills"] = 2 + gates = AGENTIC_REPLAY.evaluate_l3_lifecycle_gates(run, ARGS) + self.assertFalse(gates["passed"]) + check = next( + item for item in gates["checks"] if item["name"] == "single_physical_fill" + ) + self.assertFalse(check["passed"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/build-llama.sh b/scripts/build-llama.sh index 8075fd9f7b..7612dc445a 100755 --- a/scripts/build-llama.sh +++ b/scripts/build-llama.sh @@ -400,6 +400,9 @@ if [[ "$LLAMA_STAGE_FULL_REPLAY" == "ON" ]]; then test-skippy-recurrent-state-roundtrip test-skippy-verify-checkpoint-retirement ) + if [[ "$LLAMA_BACKEND" == "metal" ]]; then + BUILD_TARGETS+=(test-skippy-cachegen-metal) + fi fi if [[ "$LLAMA_STAGE_UPSTREAM_TESTS" == "ON" ]]; then diff --git a/scripts/ci-product-integration-smoke.sh b/scripts/ci-product-integration-smoke.sh index 5f6c246593..92c6497671 100755 --- a/scripts/ci-product-integration-smoke.sh +++ b/scripts/ci-product-integration-smoke.sh @@ -23,19 +23,28 @@ PHASE_MANIFEST="$PHASE_ROOT/phase-results.json" PHASE_RECORDS="$PHASE_ROOT/.phase-results.jsonl" SUITE_STARTED_AT_UNIX_NS="" SUITE_FINALIZED=0 +DURABLE_ONLY="${MESH_PRODUCT_INTEGRATION_DURABLE_ONLY:-0}" # This suite qualifies the runtime bundled in the product assembled by the # current workflow. A same-version published catalog may describe an older # artifact, so force the resolver onto its documented bundle fallback path. export MESH_LLM_NATIVE_RUNTIME_MANIFEST_URL="http://127.0.0.1:9/native-runtimes.json" -readonly -a REQUIRED_PHASES=( - dense-standalone - dense-openai-sdk - dense-constrained-tokio-restart - dense-split-kv - recurrent-split-kv -) +if [[ "$DURABLE_ONLY" == "1" ]]; then + readonly -a REQUIRED_PHASES=(durable-l3) +elif [[ "$DURABLE_ONLY" == "0" ]]; then + readonly -a REQUIRED_PHASES=( + dense-standalone + dense-openai-sdk + dense-constrained-tokio-restart + dense-split-kv + recurrent-split-kv + durable-l3 + ) +else + echo "MESH_PRODUCT_INTEGRATION_DURABLE_ONLY must be 0 or 1" >&2 + exit 2 +fi case "${PLATFORM}/${BACKEND}" in linux/cpu) DEVICE=CPU ;; @@ -43,6 +52,7 @@ case "${PLATFORM}/${BACKEND}" in linux/vulkan) DEVICE=Vulkan0 ;; linux/rocm) DEVICE=ROCm0 ;; macos/metal) DEVICE=MTL0 ;; + windows/cpu) DEVICE=CPU ;; *) echo "unsupported typed product suite combination: ${PLATFORM}/${BACKEND}" >&2 exit 2 @@ -241,13 +251,14 @@ append_phase_record() { local workdir="$4" local log_paths_json="$5" local split_evidence_json="$6" - local started_at_unix_ns="$7" - local ended_at_unix_ns="$8" - local exit_code="$9" + local durable_l3_evidence_json="$7" + local started_at_unix_ns="$8" + local ended_at_unix_ns="$9" + local exit_code="${10}" python3 - "$PHASE_RECORDS" "$phase" "$status" "$model_json" \ - "$workdir" "$log_paths_json" "$split_evidence_json" "$started_at_unix_ns" \ - "$ended_at_unix_ns" "$exit_code" <<'PY' + "$workdir" "$log_paths_json" "$split_evidence_json" "$durable_l3_evidence_json" \ + "$started_at_unix_ns" "$ended_at_unix_ns" "$exit_code" <<'PY' import json import sys @@ -259,6 +270,7 @@ import sys workdir, log_paths_json, split_evidence_json, + durable_l3_evidence_json, started_at_unix_ns, ended_at_unix_ns, exit_code, @@ -271,6 +283,7 @@ record = { "workdir": workdir, "log_paths": json.loads(log_paths_json), "split_evidence": json.loads(split_evidence_json), + "durable_l3_evidence": json.loads(durable_l3_evidence_json), "started_at_unix_ns": int(started_at_unix_ns), "ended_at_unix_ns": int(ended_at_unix_ns), "exit_code": int(exit_code), @@ -337,6 +350,7 @@ for record in records: seen.add(phase) model = record.get("model") split_evidence = record.get("split_evidence") + durable_l3_evidence = record.get("durable_l3_evidence") if ( record.get("status") not in {"passed", "failed"} or not isinstance(model, dict) @@ -392,6 +406,41 @@ for record in records: or evidence_payload.get("model_label") != expected_model_label ): errors.append(f"invalid split evidence payload for phase: {phase!r}") + if phase != "durable-l3": + if durable_l3_evidence is not None: + errors.append(f"unexpected durable L3 evidence for phase: {phase!r}") + elif record.get("status") == "passed": + if ( + not isinstance(durable_l3_evidence, dict) + or not isinstance(durable_l3_evidence.get("path"), str) + or not durable_l3_evidence["path"] + or not isinstance(durable_l3_evidence.get("sha256"), str) + or not re.fullmatch(r"[0-9a-f]{64}", durable_l3_evidence["sha256"]) + ): + errors.append("missing durable L3 evidence for passed phase: 'durable-l3'") + else: + evidence_path = durable_l3_evidence["path"] + try: + with open(evidence_path, "rb") as evidence_file: + raw_evidence = evidence_file.read() + actual_digest = hashlib.sha256(raw_evidence).hexdigest() + evidence_payload = json.loads(raw_evidence) + except (OSError, UnicodeDecodeError, json.JSONDecodeError) as error: + errors.append(f"unreadable durable L3 evidence: {error}") + else: + labels = [ + model.get("model", {}).get("label") + for model in evidence_payload.get("models", []) + if isinstance(model, dict) + ] + if actual_digest != durable_l3_evidence["sha256"]: + errors.append("durable L3 evidence digest mismatch") + if ( + evidence_payload.get("kind") != "mesh-llm-durable-l3-restart" + or evidence_payload.get("status") != "passed" + or labels != ["dense", "recurrent"] + ): + errors.append("invalid durable L3 evidence payload") missing = [phase for phase in required_phases if phase not in seen] if finalize: @@ -499,6 +548,58 @@ print(json.dumps({"path": path, "sha256": sha256}, sort_keys=True, separators=(" PY } +verify_durable_l3_phase_evidence() { + local phase="$1" + local phase_dir="$2" + local evidence_path="${phase_dir}/durable-l3-evidence.json" + local evidence_sha256 + + if [[ "$phase" != "durable-l3" ]]; then + printf 'null\n' + return 0 + fi + if ! python3 - "$evidence_path" <<'PY' >&2; then +import json +import re +import sys + +with open(sys.argv[1], encoding="utf-8") as handle: + evidence = json.load(handle) +if evidence.get("kind") != "mesh-llm-durable-l3-restart" or evidence.get("status") != "passed": + raise SystemExit("durable L3 evidence did not pass") +models = evidence.get("models") +if not isinstance(models, list) or [row.get("model", {}).get("label") for row in models] != ["dense", "recurrent"]: + raise SystemExit("durable L3 evidence must contain dense and recurrent rows") +for row in models: + if ( + row.get("process_boundary") is not True + or row.get("cache_root_lifecycle") != "preserved" + or row.get("exact_output_match") is not True + or not isinstance(row.get("restored_cached_tokens"), int) + or row["restored_cached_tokens"] <= 0 + or not isinstance(row.get("l3_fill_count"), int) + or row["l3_fill_count"] <= 0 + or not row.get("payload_kinds") + or row.get("configuration", {}).get("source") != "cli" + ): + raise SystemExit(f"incomplete durable L3 model evidence: {row!r}") + model_sha256 = row.get("model", {}).get("sha256") + if not isinstance(model_sha256, str) or not re.fullmatch(r"[0-9a-f]{64}", model_sha256): + raise SystemExit("durable L3 model evidence has invalid fixture digest") +PY + echo "${phase} did not produce independently verifiable durable L3 evidence" >&2 + return 1 + fi + evidence_sha256="$(fixture_sha256 "$evidence_path")" + python3 - "$evidence_path" "$evidence_sha256" <<'PY' +import json +import sys + +path, sha256 = sys.argv[1:] +print(json.dumps({"path": path, "sha256": sha256}, sort_keys=True, separators=(",", ":"))) +PY +} + finalize_interrupted_suite() { local exit_code="$?" trap - EXIT @@ -531,6 +632,7 @@ run_phase() { local phase_exit_code local phase_status local split_evidence_json=null + local durable_l3_evidence_json=null if ! ensure_phase_is_planned_once "$phase"; then write_phase_manifest failed "$phase" 1 || true @@ -555,10 +657,17 @@ run_phase() { split_evidence_json=null fi fi + if [[ "$phase_exit_code" -eq 0 ]]; then + if ! durable_l3_evidence_json="$(verify_durable_l3_phase_evidence "$phase" "$phase_dir")"; then + phase_exit_code=72 + phase_status=failed + durable_l3_evidence_json=null + fi + fi ended_at_unix_ns="$(phase_now_unix_ns)" append_phase_record "$phase" "$phase_status" "$model_json" "$phase_dir" \ - "$log_paths_json" "$split_evidence_json" "$started_at_unix_ns" "$ended_at_unix_ns" \ - "$phase_exit_code" + "$log_paths_json" "$split_evidence_json" "$durable_l3_evidence_json" \ + "$started_at_unix_ns" "$ended_at_unix_ns" "$phase_exit_code" if [[ "$phase_exit_code" -ne 0 ]]; then write_phase_manifest failed "$phase" 1 || true SUITE_FINALIZED=1 @@ -567,6 +676,7 @@ run_phase() { write_phase_manifest in-progress "" 0 } +if [[ "$DURABLE_ONLY" != "1" ]]; then run_phase dense-standalone dense \ "[\"$PHASE_ROOT/dense-standalone/server.log\",\"$PHASE_ROOT/dense-standalone/headless.log\"]" env \ MESH_CI_DEVICE="$DEVICE" \ @@ -610,6 +720,24 @@ run_phase recurrent-split-kv recurrent \ MESH_TWO_NODE_SPLIT_EXPECTED_EXACT_PAYLOAD_KIND=kv-recurrent \ MESH_TWO_NODE_SPLIT_WORK_DIR="$PHASE_ROOT/recurrent-split-kv" \ scripts/ci-two-node-split-smoke.sh "$MESH_LLM" "$ARTIFACT_DIR" "$RECURRENT_MODEL" +fi + +run_phase durable-l3 dense \ + "[\"$PHASE_ROOT/durable-l3/dense-seed.log\",\"$PHASE_ROOT/durable-l3/dense-worker.log\",\"$PHASE_ROOT/durable-l3/dense-restart-seed.log\",\"$PHASE_ROOT/durable-l3/dense-restart-worker.log\",\"$PHASE_ROOT/durable-l3/recurrent-seed.log\",\"$PHASE_ROOT/durable-l3/recurrent-worker.log\",\"$PHASE_ROOT/durable-l3/recurrent-restart-seed.log\",\"$PHASE_ROOT/durable-l3/recurrent-restart-worker.log\"]" env \ + MESH_TWO_NODE_SPLIT_DEVICE="$DEVICE" \ + MESH_TWO_NODE_SPLIT_MODEL="$DENSE_MODEL" \ + MESH_TWO_NODE_SPLIT_MODEL_LABEL=dense \ + MESH_TWO_NODE_SPLIT_RECURRENT_MODEL="$RECURRENT_MODEL" \ + MESH_TWO_NODE_SPLIT_RECURRENT_CTX_SIZE=4096 \ + MESH_TWO_NODE_SPLIT_RECURRENT_EXPECTED_EXACT_PAYLOAD_KIND=kv-recurrent \ + MESH_TWO_NODE_SPLIT_DURABLE_L3=1 \ + MESH_TWO_NODE_SPLIT_DURABLE_L3_ROOT="$PHASE_ROOT/durable-l3/cache-roots" \ + MESH_TWO_NODE_SPLIT_DENSE_ARTIFACT_ID="$DENSE_ARTIFACT_ID" \ + MESH_TWO_NODE_SPLIT_DENSE_SHA256="$DENSE_SHA256" \ + MESH_TWO_NODE_SPLIT_RECURRENT_ARTIFACT_ID="$RECURRENT_ARTIFACT_ID" \ + MESH_TWO_NODE_SPLIT_RECURRENT_SHA256="$RECURRENT_SHA256" \ + MESH_TWO_NODE_SPLIT_WORK_DIR="$PHASE_ROOT/durable-l3" \ + scripts/ci-two-node-split-smoke.sh "$MESH_LLM" "$ARTIFACT_DIR" "$DENSE_MODEL" write_phase_manifest passed "" 1 SUITE_FINALIZED=1 diff --git a/scripts/ci-two-node-split-smoke.sh b/scripts/ci-two-node-split-smoke.sh index 5990f8709d..d0d0322de8 100755 --- a/scripts/ci-two-node-split-smoke.sh +++ b/scripts/ci-two-node-split-smoke.sh @@ -60,6 +60,14 @@ mkdir -p "$WORK_DIR" # must fit platform SUN_LEN limits, especially on macOS where TMPDIR is long. PROCESS_ROOT="${MESH_TWO_NODE_SPLIT_PROCESS_ROOT:-$(mktemp -d "/tmp/m2split.XXXXXX")}" CLIENT_ROUTING="${MESH_TWO_NODE_SPLIT_CLIENT_ROUTING:-0}" +DURABLE_L3="${MESH_TWO_NODE_SPLIT_DURABLE_L3:-0}" +DURABLE_L3_ROOT="${MESH_TWO_NODE_SPLIT_DURABLE_L3_ROOT:-${WORK_DIR}/durable-l3-roots}" +DURABLE_L3_EVIDENCE_PATH="${WORK_DIR}/durable-l3-evidence.json" +DURABLE_L3_RECORDS="${WORK_DIR}/.durable-l3-evidence.jsonl" +DENSE_ARTIFACT_ID="${MESH_TWO_NODE_SPLIT_DENSE_ARTIFACT_ID:-unspecified}" +DENSE_MODEL_SHA256="${MESH_TWO_NODE_SPLIT_DENSE_SHA256:-unspecified}" +RECURRENT_ARTIFACT_ID="${MESH_TWO_NODE_SPLIT_RECURRENT_ARTIFACT_ID:-unspecified}" +RECURRENT_MODEL_SHA256="${MESH_TWO_NODE_SPLIT_RECURRENT_SHA256:-unspecified}" CLIENT_API_PORT="${MESH_TWO_NODE_SPLIT_CLIENT_API_PORT:-9369}" CLIENT_CONSOLE_PORT="${MESH_TWO_NODE_SPLIT_CLIENT_CONSOLE_PORT:-3163}" PRIMARY_MODEL_LABEL="${MESH_TWO_NODE_SPLIT_MODEL_LABEL:-}" @@ -78,6 +86,15 @@ SPLIT_EVIDENCE_PATH="" SPLIT_SNAPSHOT_DIR="" SPLIT_RECONCILE_LOG="" +if [[ "$DURABLE_L3" != "0" && "$DURABLE_L3" != "1" ]]; then + echo "MESH_TWO_NODE_SPLIT_DURABLE_L3 must be 0 or 1" >&2 + exit 2 +fi +if [[ "$DURABLE_L3" == "1" ]]; then + mkdir -p "$DURABLE_L3_ROOT" + : >"$DURABLE_L3_RECORDS" +fi + echo "=== CI Two-Node Split Smoke ===" echo " mesh-llm: $MESH_LLM" echo " bin-dir: $BIN_DIR (compatibility placeholder)" @@ -97,6 +114,7 @@ echo " ctx size: ${CTX_SIZE:-model default}" echo " max vram: ${MAX_VRAM}GB" echo " device: $DEVICE" echo " client routing: $CLIENT_ROUTING" +echo " durable L3: $DURABLE_L3" if [[ ! -x "$MESH_LLM" ]]; then echo "Missing executable mesh-llm binary: $MESH_LLM" >&2 @@ -293,6 +311,11 @@ descendant_pids() { kill_tree() { local pid="${1:-}" [[ -n "$pid" ]] || return 0 + if command -v taskkill.exe >/dev/null 2>&1; then + taskkill.exe //PID "$pid" //T //F >/dev/null 2>&1 || true + wait "$pid" 2>/dev/null || true + return 0 + fi local children children="$(descendant_pids "$pid" | sort -u || true)" kill "$pid" 2>/dev/null || true @@ -652,12 +675,24 @@ start_node() { if [[ -n "$CTX_SIZE" ]]; then args+=(--ctx-size "$CTX_SIZE") fi + if [[ "$DURABLE_L3" == "1" ]]; then + args+=( + --kv-cache-disk 2GiB + --kv-cache-disk-dir "${DURABLE_L3_ROOT}/${label}" + --kv-cache-min-free 1GiB + ) + fi - HOME="$home" \ - MESH_LLM_RUNTIME_ROOT="$runtime" \ - MESH_LLM_EPHEMERAL_KEY=1 \ - SKIPPY_TELEMETRY_STDERR=1 \ - "$MESH_LLM" "${args[@]}" >"$log_file" 2>&1 & + local -a node_env=( + env + "HOME=$home" + "MESH_LLM_RUNTIME_ROOT=$runtime" + "SKIPPY_TELEMETRY_STDERR=1" + ) + if [[ "$DURABLE_L3" != "1" ]]; then + node_env+=("MESH_LLM_EPHEMERAL_KEY=1") + fi + "${node_env[@]}" "$MESH_LLM" "${args[@]}" >"$log_file" 2>&1 & printf '%s\n' "$!" } @@ -975,6 +1010,297 @@ raise SystemExit( PY } +capture_kv_cache_statuses() { + local prefix="$1" + "$MESH_LLM" kv-cache status --port "$SEED_CONSOLE_PORT" --json \ + >"${prefix}-seed.json" + "$MESH_LLM" kv-cache status --port "$WORKER_CONSOLE_PORT" --json \ + >"${prefix}-worker.json" +} + +durable_population_ready() { + python3 - "$1-seed.json" "$1-worker.json" <<'PY' +import json +import sys + +statuses = [json.load(open(path, encoding="utf-8")) for path in sys.argv[1:]] +if any(status.get("effective", {}).get("state") != "active" for status in statuses): + raise SystemExit(1) +if sum(len(status.get("inventory") or []) for status in statuses) == 0: + raise SystemExit(1) +if sum((status.get("activity") or {}).get("writes", 0) for status in statuses) == 0: + raise SystemExit(1) +PY +} + +wait_for_durable_population() { + local prefix="$1" + local deadline=$(( $(date +%s) + READINESS_TIMEOUT_SECONDS )) + while [[ "$(date +%s)" -lt "$deadline" ]]; do + if capture_kv_cache_statuses "$prefix" 2>/dev/null && \ + durable_population_ready "$prefix" 2>/dev/null; then + return 0 + fi + sleep 1 + done + echo "durable L3 did not publish restorable state before restart" >&2 + capture_kv_cache_statuses "$prefix" 2>/dev/null || true + return 1 +} + +record_durable_restart() { + local evidence_dir="$1" + local artifact_id model_sha256 + case "$MODEL_LABEL" in + dense) + artifact_id="$DENSE_ARTIFACT_ID" + model_sha256="$DENSE_MODEL_SHA256" + ;; + recurrent) + artifact_id="$RECURRENT_ARTIFACT_ID" + model_sha256="$RECURRENT_MODEL_SHA256" + ;; + *) + echo "unsupported durable L3 model label: $MODEL_LABEL" >&2 + return 1 + ;; + esac + python3 - "$DURABLE_L3_RECORDS" "$MODEL_LABEL" "$MODEL" "$artifact_id" \ + "$model_sha256" "$EXPECTED_EXACT_PAYLOAD_KIND" \ + "${DURABLE_L3_ROOT}/seed" "${DURABLE_L3_ROOT}/worker" \ + "$evidence_dir" <<'PY' +import json +import os +from pathlib import Path +import sys + +( + records_path, + model_label, + model_path, + artifact_id, + model_sha256, + expected_payload_kind, + seed_root, + worker_root, + evidence_dir, +) = sys.argv[1:] +root = Path(evidence_dir) + +def load(name): + with (root / name).open(encoding="utf-8") as handle: + return json.load(handle) + +before = {node: load(f"before-{node}.json") for node in ("seed", "worker")} +restart_before = { + node: load(f"restart-before-{node}.json") for node in ("seed", "worker") +} +after = {node: load(f"after-{node}.json") for node in ("seed", "worker")} +cleared = {node: load(f"cleared-{node}.json") for node in ("seed", "worker")} +warm_response = load("warm-response.json") +restored_response = load("restored-response.json") + +for stage, statuses in ( + ("before", before), + ("restart-before", restart_before), + ("after", after), + ("cleared", cleared), +): + for node, status in statuses.items(): + effective = status.get("effective") or {} + if effective.get("state") != "active" or effective.get("reason") is not None: + raise SystemExit(f"{stage} {node} disk tier is not active: {effective!r}") + configured = status.get("configured") or {} + sources = configured.get("sources") or {} + if configured.get("mode") != "fixed" or configured.get("budget_bytes") != 2 * 1024**3: + raise SystemExit(f"{stage} {node} has unexpected disk configuration: {configured!r}") + if configured.get("minimum_free_bytes") != 1024**3: + raise SystemExit(f"{stage} {node} has unexpected free-space reserve: {configured!r}") + if any(sources.get(field) != "cli" for field in ("mode", "budget", "directory", "minimum_free")): + raise SystemExit(f"{stage} {node} configuration source is not CLI: {sources!r}") + +before_inventory = sum(len(status.get("inventory") or []) for status in before.values()) +restart_inventory = sum( + len(status.get("inventory") or []) for status in restart_before.values() +) +if before_inventory == 0 or restart_inventory == 0: + raise SystemExit("durable inventory was missing before or after process restart") + +before_writes = sum((status.get("activity") or {}).get("writes", 0) for status in before.values()) +restart_fills = sum((status.get("activity") or {}).get("fills", 0) for status in after.values()) +restart_initial_fills = sum( + (status.get("activity") or {}).get("fills", 0) for status in restart_before.values() +) +if before_writes <= 0: + raise SystemExit("the population process recorded no durable L3 writes") +if restart_initial_fills != 0 or restart_fills <= 0: + raise SystemExit( + f"restart did not prove an L3 fill: before={restart_initial_fills}, after={restart_fills}" + ) + +usage = restored_response.get("usage") or {} +prompt_tokens = usage.get("prompt_tokens") +cached_tokens = (usage.get("prompt_tokens_details") or {}).get("cached_tokens") +if not isinstance(prompt_tokens, int) or not isinstance(cached_tokens, int) or cached_tokens <= 0: + raise SystemExit(f"restart response did not report restored tokens: {usage!r}") + +def output_text(response): + try: + return response["choices"][0]["message"]["content"] + except (KeyError, IndexError, TypeError): + raise SystemExit(f"response omitted assistant output: {response!r}") + +warm_output = output_text(warm_response) +restored_output = output_text(restored_response) +if warm_output != restored_output: + raise SystemExit( + f"restart output diverged: warm={warm_output!r}, restored={restored_output!r}" + ) + +payload_kinds = sorted({ + entry.get("payload_kind") + for status in restart_before.values() + for entry in status.get("inventory") or [] + if isinstance(entry, dict) and isinstance(entry.get("payload_kind"), str) +}) +if not payload_kinds: + raise SystemExit("durable inventory did not identify a payload kind") +if expected_payload_kind and expected_payload_kind not in payload_kinds: + raise SystemExit( + f"expected payload kind {expected_payload_kind!r}, observed {payload_kinds!r}" + ) + +for node, status in cleared.items(): + if status.get("inventory"): + raise SystemExit(f"clear left {node} inventory behind") + if (status.get("usage") or {}).get("used_bytes") != 0: + raise SystemExit(f"clear left {node} managed bytes behind: {status.get('usage')!r}") + +record = { + "model": { + "label": model_label, + "artifact_id": artifact_id, + "sha256": model_sha256, + "path": model_path, + }, + "configuration": { + "source": "cli", + "mode": "fixed", + "budget_bytes": 2 * 1024**3, + "minimum_free_bytes": 1024**3, + "roots": {"seed": seed_root, "worker": worker_root}, + }, + "cache_root_lifecycle": "preserved", + "process_boundary": True, + "payload_kinds": payload_kinds, + "statuses": { + "before_stop": before, + "after_restart_before_request": restart_before, + "after_restore": after, + "after_clear": cleared, + }, + "restored_prompt_tokens": prompt_tokens, + "restored_cached_tokens": cached_tokens, + "l3_fill_count": restart_fills, + "exact_output_match": True, + "status_command": "mesh-llm kv-cache status --json", + "clear_command": "mesh-llm kv-cache clear --yes --json", +} +with open(records_path, "a", encoding="utf-8") as handle: + handle.write(json.dumps(record, sort_keys=True, separators=(",", ":")) + "\n") +PY +} + +run_durable_restart_probe() { + [[ "$DURABLE_L3" == "1" ]] || return 0 + local request_path="$1" + local warm_response_path="$2" + local evidence_dir="${WORK_DIR}/durable-${MODEL_LABEL}" + mkdir -p "$evidence_dir" + cp "$request_path" "$evidence_dir/request.json" + cp "$warm_response_path" "$evidence_dir/warm-response.json" + wait_for_durable_population "$evidence_dir/before" + + kill_tree "$CLIENT_PID" + CLIENT_PID="" + kill_tree "$WORKER_PID" + WORKER_PID="" + kill_tree "$SEED_PID" + SEED_PID="" + + SEED_LOG="${WORK_DIR}/${MODEL_LABEL}-restart-seed.log" + WORKER_LOG="${WORK_DIR}/${MODEL_LABEL}-restart-worker.log" + SEED_PID="$(start_node seed "" "$SEED_API_PORT" "$SEED_CONSOLE_PORT" "$SEED_BIND_PORT" "$SEED_LOG")" + wait_for_seed_token "${MODEL_LABEL} durable restart: " + WORKER_PID="$(start_node worker "$TOKEN" "$WORKER_API_PORT" "$WORKER_CONSOLE_PORT" "$WORKER_BIND_PORT" "$WORKER_LOG")" + DRIVER_LABEL="" + DRIVER_API_PORT="" + wait_for_split_topology "${MODEL_LABEL} durable restart: " + MODEL_ID="$(python3 -c 'import json,sys; print(json.load(open(sys.argv[1], encoding="utf-8")).get("model_id", ""))' "$SPLIT_EVIDENCE_PATH")" + [[ -n "$MODEL_ID" ]] || { + echo "durable restart split evidence did not return a model id" >&2 + return 1 + } + sleep "$REQUEST_SETTLE_SECONDS" + capture_kv_cache_statuses "$evidence_dir/restart-before" + python3 - "$MODEL_ID" "$evidence_dir/request.json" <<'PY' +import json +import os +import sys + +model, path = sys.argv[1:] +with open(path, encoding="utf-8") as handle: + payload = json.load(handle) +payload["model"] = model +temporary = f"{path}.tmp" +with open(temporary, "w", encoding="utf-8") as handle: + json.dump(payload, handle) +os.replace(temporary, path) +PY + curl -fsS --max-time 180 \ + "http://127.0.0.1:${DRIVER_API_PORT}/v1/chat/completions" \ + -H 'content-type: application/json' \ + -d @"$evidence_dir/request.json" \ + -o "$evidence_dir/restored-response.json" + sleep "$REQUEST_SETTLE_SECONDS" + capture_kv_cache_statuses "$evidence_dir/after" + "$MESH_LLM" kv-cache clear --port "$SEED_CONSOLE_PORT" --yes --json \ + >"$evidence_dir/clear-seed.json" + "$MESH_LLM" kv-cache clear --port "$WORKER_CONSOLE_PORT" --yes --json \ + >"$evidence_dir/clear-worker.json" + capture_kv_cache_statuses "$evidence_dir/cleared" + record_durable_restart "$evidence_dir" +} + +write_durable_l3_evidence() { + [[ "$DURABLE_L3" == "1" ]] || return 0 + python3 - "$DURABLE_L3_RECORDS" "$DURABLE_L3_EVIDENCE_PATH" \ + "$([[ -n "$RECURRENT_MODEL" ]] && printf 'dense,recurrent' || printf '%s' "$PRIMARY_MODEL_LABEL")" <<'PY' +import json +import os +import sys + +records_path, output_path, expected_labels = sys.argv[1:] +with open(records_path, encoding="utf-8") as handle: + records = [json.loads(line) for line in handle if line.strip()] +expected = expected_labels.split(",") +observed = [record.get("model", {}).get("label") for record in records] +if observed != expected: + raise SystemExit(f"durable L3 evidence models differ: expected {expected}, observed {observed}") +evidence = { + "schema_version": 1, + "kind": "mesh-llm-durable-l3-restart", + "status": "passed", + "models": records, +} +temporary = f"{output_path}.tmp" +with open(temporary, "w", encoding="utf-8") as handle: + json.dump(evidence, handle, indent=2, sort_keys=True) + handle.write("\n") +os.replace(temporary, output_path) +PY +} + prefix_validated=0 for attempt in $(seq 1 "$PREFIX_ATTEMPTS"); do payload_dir="${PREFIX_PAYLOAD_ROOT}/attempt-${attempt}" @@ -1026,6 +1352,9 @@ fi assert_expected_stage_payload +run_durable_restart_probe "${payload_dir}/prompt-${PREFIX_REQUEST_COUNT}.json" \ + "${response_dir}/response-${PREFIX_REQUEST_COUNT}.json" + echo "Two-node split smoke passed for model leg: ${MODEL_LABEL:-default}" # Optional recurrent leg: rerun the identical flow against a second model in @@ -1114,5 +1443,10 @@ if [[ -n "$RECURRENT_MODEL" ]]; then assert_expected_stage_payload + run_durable_restart_probe "${payload_dir}/prompt-${PREFIX_REQUEST_COUNT}.json" \ + "${response_dir}/response-${PREFIX_REQUEST_COUNT}.json" + echo "Two-node split smoke passed for model leg: recurrent" fi + +write_durable_l3_evidence diff --git a/scripts/remote-handoff-sweep.sh b/scripts/remote-handoff-sweep.sh new file mode 100755 index 0000000000..315440ef72 --- /dev/null +++ b/scripts/remote-handoff-sweep.sh @@ -0,0 +1,65 @@ +#!/usr/bin/env bash +# Sweep remote-handoff sender runs across prefix lengths against a receiver +# started with --accept-count matching the number of runs, e.g.: +# +# receiver$ target/release/skippy-correctness remote-handoff --role recv \ +# --listen 0.0.0.0:19081 --model M --layer-end N --ctx-size 16384 \ +# --n-gpu-layers 99 --decode-tokens 32 --accept-count 4 \ +# --allow-mismatch --report-out recv.json +# +# sender$ scripts/remote-handoff-sweep.sh :19081 \ +# [prefix counts...] +set -euo pipefail + +PEER="${1:?receiver address}" +MODEL="${2:?model path}" +LAYER_END="${3:?layer end}" +OUT_DIR="${4:?output directory}" +shift 4 +PREFIXES=("${@:-512 2048 4096 8192}") +if [[ $# -eq 0 ]]; then PREFIXES=(512 2048 4096 8192); fi + +CTX_SIZE="${CTX_SIZE:-16384}" +DECODE_TOKENS="${DECODE_TOKENS:-32}" +BIN="${BIN:-target/release/skippy-correctness}" + +mkdir -p "$OUT_DIR" +for prefix in "${PREFIXES[@]}"; do + echo "== prefix ${prefix}" + "$BIN" remote-handoff --role send --peer "$PEER" \ + --model "$MODEL" --layer-end "$LAYER_END" --ctx-size "$CTX_SIZE" \ + --n-gpu-layers 99 --prefix-token-count "$prefix" \ + --decode-tokens "$DECODE_TOKENS" --baseline \ + --report-out "$OUT_DIR/send-${prefix}.json" \ + > "$OUT_DIR/send-${prefix}.log" 2>&1 \ + || echo " prefix ${prefix} FAILED (see $OUT_DIR/send-${prefix}.log)" +done + +python3 - "$OUT_DIR" <<'EOF' +import glob, json, sys + +rows = [] +for path in sorted(glob.glob(f"{sys.argv[1]}/send-*.json")): + r = json.load(open(path)) + rows.append(( + r["prompt_token_count"], + r["state_bytes"] / 2**20, + r["transfer_gbps"], + r["source_prefill_ms"], + r["state_export_ms"], + r["transfer_ms"], + r["receiver"]["kv_attach_ms"], + r["ttft_disaggregated_ms"], + r.get("ttft_local_ms"), + r.get("ttft_speedup"), + r["matches"], + )) +rows.sort() +print(f"{'prefix':>7} {'MiB':>7} {'Gbps':>6} {'prefill':>8} {'export':>7} " + f"{'xfer':>7} {'attach':>7} {'ttft-pd':>8} {'ttft-lo':>8} {'speedup':>7} match") +for r in rows: + local = f"{r[8]:8.0f}" if r[8] is not None else " -" + speedup = f"{r[9]:7.2f}" if r[9] is not None else " -" + print(f"{r[0]:>7} {r[1]:7.1f} {r[2]:6.2f} {r[3]:8.0f} {r[4]:7.0f} " + f"{r[5]:7.0f} {r[6]:7.0f} {r[7]:8.0f} {local} {speedup} {r[10]}") +EOF diff --git a/scripts/tests/test_ci_lane_workflows.py b/scripts/tests/test_ci_lane_workflows.py index 18b3ad4fe1..a52b549e26 100644 --- a/scripts/tests/test_ci_lane_workflows.py +++ b/scripts/tests/test_ci_lane_workflows.py @@ -138,7 +138,7 @@ def test_dispatched_lanes_pass_source_sha_only_to_product_workflows( "ci-website-lane.yml": 2, "ci-linux-lane.yml": 10, "ci-macos-lane.yml": 9, - "ci-windows-lane.yml": 6, + "ci-windows-lane.yml": 7, } for workflow_name, expected_calls in lane_workflows.items(): with self.subTest(workflow=workflow_name): @@ -215,11 +215,15 @@ def test_lane_plans_are_bounded_platform_projections(self) -> None: for platform in ("linux", "macos", "windows"): self.assertIn(f'select(.platform == "{platform}")', action) self.assertIn( - 'smoke: [.matrices.smoke[] | select(.id != "metal-model-load")]', + 'smoke: [.matrices.smoke[] | select(.id != "metal-model-load" and .id != "product-integration-metal" and .id != "product-integration-windows-cpu")]', action, ) self.assertIn( - 'smoke: [.matrices.smoke[] | select(.id == "metal-model-load")]', + 'smoke: [.matrices.smoke[] | select(.id == "metal-model-load" or .id == "product-integration-metal")]', + action, + ) + self.assertIn( + 'smoke: [.matrices.smoke[] | select(.id == "product-integration-windows-cpu")]', action, ) @@ -328,6 +332,9 @@ def test_product_smoke_jobs_parse_formatted_matrix_json(self) -> None: "metal-model-load", "product-integration-metal", ), + "ci-windows-product-smoke-slice.yml": ( + "product-integration-windows-cpu", + ), } for workflow_name, smoke_ids in smoke_workflows.items(): workflow = self.workflow(workflow_name) diff --git a/scripts/tests/test_ci_product_integration_smoke.py b/scripts/tests/test_ci_product_integration_smoke.py index 86e955a0e5..0c24967288 100644 --- a/scripts/tests/test_ci_product_integration_smoke.py +++ b/scripts/tests/test_ci_product_integration_smoke.py @@ -19,6 +19,7 @@ "dense-constrained-tokio-restart", "dense-split-kv", "recurrent-split-kv", + "durable-l3", ] DENSE_ARTIFACT_ID = "smollm2-q8-inference" RECURRENT_ARTIFACT_ID = "family-granite-hybrid" @@ -84,6 +85,42 @@ def write_stub(self, path: Path) -> None: missing-*|tampered-*) ;; *) exit 64 ;; esac +if [[ "$phase" == durable-l3 ]]; then +python3 - "$evidence_root/durable-l3-evidence.json" <<'PY' +import json +import os +import sys + +rows = [] +for label, kind, artifact_var, sha_var in ( + ("dense", "full-state", "MESH_TWO_NODE_SPLIT_DENSE_ARTIFACT_ID", "MESH_TWO_NODE_SPLIT_DENSE_SHA256"), + ("recurrent", "kv-recurrent", "MESH_TWO_NODE_SPLIT_RECURRENT_ARTIFACT_ID", "MESH_TWO_NODE_SPLIT_RECURRENT_SHA256"), +): + rows.append({ + "model": { + "label": label, + "artifact_id": os.environ[artifact_var], + "sha256": os.environ[sha_var], + "path": f"/{label}.gguf", + }, + "configuration": {"source": "cli"}, + "cache_root_lifecycle": "preserved", + "process_boundary": True, + "payload_kinds": [kind], + "restored_cached_tokens": 128, + "l3_fill_count": 1, + "exact_output_match": True, + }) +with open(sys.argv[1], "w", encoding="utf-8") as handle: + json.dump({ + "schema_version": 1, + "kind": "mesh-llm-durable-l3-restart", + "status": "passed", + "models": rows, + }, handle) + handle.write("\\n") +PY +fi """, encoding="utf-8", ) @@ -199,6 +236,8 @@ def run_suite( "STUB_CUDA_NORMAL_PROBE_STATUS": str(cuda_normal_probe_status), "STUB_CUDA_STRICT_PROBE_STATUS": str(cuda_strict_probe_status), } + if platform == "windows": + env["MESH_PRODUCT_INTEGRATION_DURABLE_ONLY"] = "1" if failure_phase is not None: env["STUB_FAIL_PHASE"] = failure_phase @@ -229,7 +268,7 @@ def run_suite( ) return result, manifest - def test_success_reconciles_the_exact_five_phases(self) -> None: + def test_success_reconciles_the_required_phases(self) -> None: result, manifest = self.run_suite() self.assertEqual(result.returncode, 0, result.stderr) @@ -295,6 +334,17 @@ def test_success_reconciles_the_exact_five_phases(self) -> None: ) else: self.assertIsNone(phase["split_evidence"]) + if phase["phase"] == "durable-l3": + self.assertRegex( + phase["durable_l3_evidence"]["sha256"], r"^[0-9a-f]{64}$" + ) + self.assertTrue( + phase["durable_l3_evidence"]["path"].endswith( + "durable-l3-evidence.json" + ) + ) + else: + self.assertIsNone(phase["durable_l3_evidence"]) self.assertLessEqual( phase["started_at_unix_ns"], phase["ended_at_unix_ns"] ) @@ -306,6 +356,7 @@ def test_typed_backend_selectors_drive_the_expected_device(self) -> None: ("linux", "vulkan"): "Vulkan0", ("linux", "rocm"): "ROCm0", ("macos", "metal"): "MTL0", + ("windows", "cpu"): "CPU", } for (platform, backend), expected_device in cases.items(): with self.subTest(platform=platform, backend=backend): diff --git a/scripts/tests/test_ci_two_node_split_smoke.py b/scripts/tests/test_ci_two_node_split_smoke.py index 2999bc777d..311013fde1 100644 --- a/scripts/tests/test_ci_two_node_split_smoke.py +++ b/scripts/tests/test_ci_two_node_split_smoke.py @@ -174,7 +174,7 @@ def test_prefix_validator_exit_behavior(self): def test_readiness_reconciles_persisted_snapshots_from_both_observers(self): script = SMOKE_SCRIPT.read_text(encoding="utf-8") - self.assertEqual(script.count('wait_for_split_topology "'), 2) + self.assertEqual(script.count('wait_for_split_topology "'), 3) for observer in ("seed", "worker"): for snapshot in ("status", "stages", "models"): self.assertIn(f"{observer}-{snapshot}.json", script) diff --git a/scripts/tests/test_ci_workflow_artifacts.py b/scripts/tests/test_ci_workflow_artifacts.py index c69a8ded75..247adaa7d1 100644 --- a/scripts/tests/test_ci_workflow_artifacts.py +++ b/scripts/tests/test_ci_workflow_artifacts.py @@ -55,6 +55,7 @@ def test_orchestrator_calls_same_slices_for_pr_and_main(self): "ci-platform-checks-slice.yml", "ci-linux-product-smoke-slice.yml", "ci-macos-product-smoke-slice.yml", + "ci-windows-product-smoke-slice.yml", "ci-linux-sdk-slice.yml", "ci-macos-sdk-slice.yml", ): @@ -243,7 +244,8 @@ def test_two_node_split_smoke_covers_dense_and_recurrent_models(self): self.assertIn("run_phase dense-split-kv", product_script) self.assertIn("run_phase recurrent-split-kv", product_script) self.assertIn("MESH_TWO_NODE_SPLIT_EXPECTED_EXACT_PAYLOAD_KIND=kv-recurrent", product_script) - self.assertNotIn("MESH_TWO_NODE_SPLIT_RECURRENT_MODEL=", product_script) + self.assertEqual(product_script.count("MESH_TWO_NODE_SPLIT_RECURRENT_MODEL="), 1) + self.assertIn("run_phase durable-l3", product_script) self.assertIn("run_client_routing_probe", smoke_script) self.assertIn("Passive client routing and streaming validated", smoke_script) self.assertIn( @@ -262,6 +264,7 @@ def test_product_integration_uploads_reconciled_phase_evidence_on_every_outcome( self.assertIn("phase-results.json", workflow) self.assertIn("*/split-evidence.json", workflow) self.assertIn("*/split-evidence-snapshots/*.json", workflow) + self.assertIn("*/durable-l3-evidence.json", workflow) self.assertIn("*/*.log", workflow) self.assertIn("-evidence", workflow) self.assertIn("if-no-files-found: error", workflow) @@ -270,12 +273,21 @@ def test_protected_catalog_defers_product_integration_rollout(self): slices = json.loads(SLICES.read_text()) smoke_ids = {row["id"] for row in slices["smoke_rows"]} linux = (WORKFLOWS / "ci-linux-product-smoke-slice.yml").read_text() + windows = (WORKFLOWS / "ci-windows-product-smoke-slice.yml").read_text() + product_workflow = (WORKFLOWS / "product-integration-smoke.yml").read_text() self.assertNotIn("product-integration-cpu", smoke_ids) + self.assertNotIn("product-integration-metal", smoke_ids) + self.assertNotIn("product-integration-windows-cpu", smoke_ids) self.assertNotIn("qwen-recurrent-gate", smoke_ids) self.assertIn("core", smoke_ids) self.assertIn("two-node-client", smoke_ids) self.assertIn("two-node-split", smoke_ids) + self.assertIn( + "binary_name: ${{ inputs.platform == 'windows' && 'mesh-llm.exe' || 'mesh-llm' }}", + product_workflow, + ) + self.assertIn("product-integration-windows-cpu", windows) for smoke_id in ("core", "two-node-client", "two-node-split"): self.assertIn( f"contains(fromJson(inputs.smoke_matrix).*.id, '{smoke_id}')", @@ -284,7 +296,6 @@ def test_protected_catalog_defers_product_integration_rollout(self): self.assertIn("Qwen3.5-0.8B-Q4_K_M.gguf", linux) self.assertIn("expected_exact_payload_kind: kv-recurrent", linux) self.assertNotIn("product-integration-cuda", smoke_ids) - self.assertNotIn("product-integration-metal", smoke_ids) self.assertIn("core-cuda", smoke_ids) self.assertIn("metal-model-load", smoke_ids) diff --git a/scripts/tests/test_kv_restart_replay.py b/scripts/tests/test_kv_restart_replay.py new file mode 100644 index 0000000000..82019b01ab --- /dev/null +++ b/scripts/tests/test_kv_restart_replay.py @@ -0,0 +1,282 @@ +from __future__ import annotations + +import importlib.util +import json +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest import mock + + +REPO = Path(__file__).resolve().parents[2] +SCRIPT = REPO / "evals/kv-restart-replay.py" + + +def load_module(): + spec = importlib.util.spec_from_file_location("kv_restart_replay", SCRIPT) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot import {SCRIPT}") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +BENCH = load_module() + + +class KvRestartReplayTest(unittest.TestCase): + def test_server_command_rejects_endpoint_overrides(self) -> None: + binary = Path("/tmp/mesh-llm") + + with self.assertRaisesRegex(AssertionError, "--port"): + BENCH.server_command(binary, "model.gguf", ["--port=9447"]) + with self.assertRaisesRegex(AssertionError, "--host"): + BENCH.server_command(binary, "model.gguf", ["--host", "0.0.0.0"]) + + def test_manifest_uses_distinct_deterministic_assistant_responses(self) -> None: + manifest = BENCH.build_manifest(2, 32, 16) + + self.assertNotEqual(manifest["turns"][0]["request"], manifest["turns"][0]["response"]) + self.assertEqual(manifest, BENCH.build_manifest(2, 32, 16)) + + def test_macos_memory_uses_sysctl(self) -> None: + results = [ + subprocess.CompletedProcess([], 0, stdout="Apple M2\n"), + subprocess.CompletedProcess([], 0, stdout="Mac14,6\n"), + subprocess.CompletedProcess([], 0, stdout="17179869184\n"), + ] + with ( + mock.patch.object(BENCH.sys, "platform", "darwin"), + mock.patch.object(BENCH.subprocess, "run", side_effect=results) as run, + ): + fingerprint = BENCH.hardware_fingerprint() + + self.assertEqual(fingerprint["physical_memory_bytes"], 17179869184) + self.assertEqual(run.call_args_list[-1].args[0], ["sysctl", "-n", "hw.memsize"]) + + def test_binary_provenance_keeps_unknown_source_sha_without_git(self) -> None: + with tempfile.TemporaryDirectory() as directory: + binary = Path(directory) / "mesh-llm" + binary.write_bytes(b"binary") + with mock.patch.object(BENCH.subprocess, "run", side_effect=FileNotFoundError): + provenance = BENCH.binary_provenance(binary) + + self.assertEqual(provenance["git_describe"], "unknown") + self.assertEqual(provenance["source_sha"], "unknown") + + def test_stream_request_rejects_truncated_stream_with_usage(self) -> None: + class TruncatedResponse: + status = 200 + + def __iter__(self): + return iter( + [ + b'data: {"choices":[{"delta":{"content":"partial"}}]}\n', + b'data: {"choices":[],"usage":{"completion_tokens":1,"prompt_tokens":10}}\n', + ] + ) + + class TruncatedConnection: + def __init__(self, *_args, **_kwargs): + pass + + def request(self, *_args, **_kwargs): + pass + + def getresponse(self): + return TruncatedResponse() + + def close(self): + pass + + with mock.patch.object(BENCH.http.client, "HTTPConnection", TruncatedConnection): + result = BENCH.stream_request( + "request-1", + [{"role": "user", "content": "task"}], + "model", + 8, + 10, + ) + + self.assertEqual(result["error"], "stream ended without terminal [DONE] marker") + + def test_stream_request_requires_prompt_and_cached_usage(self) -> None: + class Response: + status = 200 + + def __init__(self, usage): + self.usage = usage + + def __iter__(self): + return iter( + [ + b'data: {"choices":[{"delta":{"content":"ok"}}]}\n', + f'data: {json.dumps({"choices": [], "usage": self.usage})}\n'.encode(), + b"data: [DONE]\n", + ] + ) + + class Connection: + usage = {} + + def __init__(self, *_args, **_kwargs): + pass + + def request(self, *_args, **_kwargs): + pass + + def getresponse(self): + return Response(self.usage) + + def close(self): + pass + + for usage, expected in [ + ({"prompt_tokens_details": {"cached_tokens": 0}}, "prompt token usage"), + ({"prompt_tokens": 10}, "cached token usage"), + ]: + with self.subTest(expected=expected): + Connection.usage = usage + with mock.patch.object(BENCH.http.client, "HTTPConnection", Connection): + result = BENCH.stream_request("request", [], "model", 8, 10) + self.assertIn(expected, result["error"]) + + def test_single_restore_sample_does_not_report_p95(self) -> None: + summary = BENCH.summarize_cohort( + "restore", + [ + { + "ttft_seconds": 0.2, + "total_seconds": 0.3, + "prompt_tokens": 10, + "cached_tokens": 5, + "decode_tokens_per_second": 10.0, + } + ], + ) + + self.assertEqual(summary["ttft_p50_seconds"], 0.2) + self.assertIsNone(summary["ttft_p95_seconds"]) + + def test_run_arm_replays_the_last_fill_prompt_once_per_post_restart_request(self) -> None: + calls: list[tuple[str, list[dict[str, str]]]] = [] + + def fake_stream(request_id, messages, *_args): + calls.append((request_id, list(messages))) + return { + "request_id": request_id, + "ttft_seconds": 0.1, + "total_seconds": 0.2, + "prompt_tokens": 10, + "completion_tokens": 1, + "cached_tokens": 5, + "decode_tokens_per_second": 10.0, + } + + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + binary = root / "mesh-llm" + model = root / "model.gguf" + binary.write_bytes(b"binary") + model.write_bytes(b"model") + args = SimpleNamespace( + binary=str(binary), + model=str(model), + turns=2, + turn_target_tokens=32, + system_tokens=16, + restore_repeats=3, + max_output_tokens=8, + request_timeout=10.0, + ready_timeout=10.0, + serve_extra_args=[], + ) + with ( + mock.patch.object( + BENCH, + "start_server", + side_effect=lambda *_args: (SimpleNamespace(), ["mesh-llm", "serve"]), + ), + mock.patch.object(BENCH, "stop_server"), + mock.patch.object(BENCH, "wait_for_model", return_value="model"), + mock.patch.object(BENCH, "stream_request", side_effect=fake_stream), + mock.patch.object(BENCH, "binary_provenance", return_value={"source_sha": "a" * 40}), + mock.patch.object(BENCH, "hardware_fingerprint", return_value={"platform": "test"}), + ): + run = BENCH.run_arm(args, root / "output") + + self.assertEqual([row["requests"] for row in run["cohorts"]], [2, 1, 2]) + self.assertEqual([request_id for request_id, _ in calls], [ + "fill-1", + "fill-2", + "restore-1", + "warm-1", + "warm-2", + ]) + self.assertTrue(all(messages[-1]["role"] == "user" for _, messages in calls)) + self.assertEqual(calls[1][1], calls[2][1]) + self.assertEqual(calls[2][1], calls[3][1]) + + def test_run_arm_propagates_final_server_shutdown_failure(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + binary = root / "mesh-llm" + model = root / "model.gguf" + binary.write_bytes(b"binary") + model.write_bytes(b"model") + args = SimpleNamespace( + binary=str(binary), + model=str(model), + turns=1, + turn_target_tokens=32, + system_tokens=16, + restore_repeats=1, + max_output_tokens=8, + request_timeout=10.0, + ready_timeout=10.0, + serve_extra_args=[], + ) + with ( + mock.patch.object(BENCH, "start_server", return_value=(SimpleNamespace(), [])), + mock.patch.object(BENCH, "wait_for_model", return_value="model"), + mock.patch.object(BENCH, "stop_server", side_effect=[None, RuntimeError("did not stop")]), + mock.patch.object( + BENCH, + "stream_request", + return_value={ + "request_id": "fill-1", + "ttft_seconds": 0.1, + "total_seconds": 0.2, + "prompt_tokens": 10, + "completion_tokens": 1, + "cached_tokens": 0, + "decode_tokens_per_second": 10.0, + }, + ), + mock.patch.object(BENCH, "binary_provenance", return_value={"source_sha": "a" * 40}), + mock.patch.object(BENCH, "hardware_fingerprint", return_value={"platform": "test"}), + ): + with self.assertRaisesRegex(RuntimeError, "did not stop"): + BENCH.run_arm(args, root / "output") + + def test_main_rejects_any_nonempty_output_directory(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + output = root / "output" + output.mkdir() + (output / "partial.log").write_text("partial", encoding="utf-8") + with mock.patch.object( + sys, + "argv", + ["kv-restart-replay.py", "--model", str(root / "model.gguf"), "--output", str(output)], + ): + with self.assertRaisesRegex(SystemExit, "not empty"): + BENCH.main() + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/tests/test_llama_native_full_replay.py b/scripts/tests/test_llama_native_full_replay.py index 77b0d6e45d..a550da944c 100644 --- a/scripts/tests/test_llama_native_full_replay.py +++ b/scripts/tests/test_llama_native_full_replay.py @@ -39,6 +39,7 @@ def run_build( full_replay: bool, upstream_tests: bool = False, repeat_cached: bool = False, + backend: str = "cpu", ) -> list[dict[str, object]]: with tempfile.TemporaryDirectory() as temp_dir: root = Path(temp_dir) @@ -123,7 +124,7 @@ def run_build( { "LLAMA_WORKDIR": str(llama), "LLAMA_STAGE_BUILD_DIR": str(build), - "LLAMA_STAGE_BACKEND": "cpu", + "LLAMA_STAGE_BACKEND": backend, "LLAMA_STAGE_LINK_MODE": "static", "LLAMA_STAGE_FORCE_BUILD": "1", "LLAMA_STAGE_USE_SCCACHE": "0", @@ -196,6 +197,12 @@ def test_cached_full_replay_rebuilds_and_runs_skippy_gates(self) -> None: self.assertEqual(len(build_calls), 2) self.assertEqual(len(ctest_calls), 2) + def test_metal_full_replay_builds_cachegen_fixture(self) -> None: + trace = self.run_build(full_replay=True, backend="metal") + build = next(call for call in trace if call["args"][0] == "--build") + + self.assertIn("test-skippy-cachegen-metal", build["args"]) + def test_cached_standard_build_keeps_cache_shortcut(self) -> None: trace = self.run_build(full_replay=False, repeat_cached=True) diff --git a/scripts/tests/test_validate_ci_lane_results.py b/scripts/tests/test_validate_ci_lane_results.py index 0e54d1a434..14b391a33b 100644 --- a/scripts/tests/test_validate_ci_lane_results.py +++ b/scripts/tests/test_validate_ci_lane_results.py @@ -81,7 +81,7 @@ def test_macos_lane_maps_swift_sdk_and_platform_jobs(self) -> None: self.assertEqual(VALIDATOR._required_jobs(plan), expected) VALIDATOR.validate(plan, {job: state("success") for job in expected}) - def test_windows_runtime_products_do_not_require_host_rows(self) -> None: + def test_windows_runtime_products_require_selected_product_smoke(self) -> None: plan = { "lane": "windows", "required": True, @@ -90,9 +90,10 @@ def test_windows_runtime_products_do_not_require_host_rows(self) -> None: "hosts": [], "runtime_products": [{"id": "windows-cpu"}], "platform_checks": [], + "smoke": [{"id": "product-integration-windows-cpu"}], }, } - expected = {"native_runtimes", "runtime_product"} + expected = {"native_runtimes", "runtime_product", "product_smoke"} self.assertEqual(VALIDATOR._required_jobs(plan), expected) VALIDATOR.validate(plan, {job: state("success") for job in expected}) diff --git a/scripts/validate-ci-lane-results.py b/scripts/validate-ci-lane-results.py index 9c629288c5..7eabe7dd7e 100644 --- a/scripts/validate-ci-lane-results.py +++ b/scripts/validate-ci-lane-results.py @@ -99,6 +99,8 @@ def _required_jobs(lane_plan: dict[str, Any]) -> set[str]: jobs.update({"native_runtimes", "runtime_product"}) if _ids(lane_plan, "platform_checks"): jobs.add("platform_checks") + if _ids(lane_plan, "smoke"): + jobs.add("product_smoke") else: raise LaneResultError(f"unknown CI lane {lane!r}") return jobs diff --git a/third_party/llama.cpp/patches/0003-skippy-implement-model-lifecycle-and-package-loading.patch b/third_party/llama.cpp/patches/0003-skippy-implement-model-lifecycle-and-package-loading.patch index 4c6c7df905..3602c9d412 100644 --- a/third_party/llama.cpp/patches/0003-skippy-implement-model-lifecycle-and-package-loading.patch +++ b/third_party/llama.cpp/patches/0003-skippy-implement-model-lifecycle-and-package-loading.patch @@ -359,7 +359,7 @@ new file mode 100644 index 000000000..a6ec5b2eb --- /dev/null +++ b/src/skippy/model_load_config.cpp -@@ -0,0 +1,159 @@ +@@ -0,0 +1,167 @@ +#include "skippy.h" +#include "skippy/errors.h" +#include "skippy/environment.h" @@ -514,10 +514,18 @@ index 000000000..a6ec5b2eb +void skippy_apply_context_hardware_config( + const struct skippy_runtime_config * config, + llama_context_params & params) { -+ if (config == nullptr || config->op_offload == SKIPPY_TRISTATE_AUTO) { ++ if (config == nullptr) { + return; + } -+ params.op_offload = config->op_offload == SKIPPY_TRISTATE_TRUE; ++ if (config->cache_type_k >= 0) { ++ params.type_k = static_cast(config->cache_type_k); ++ } ++ if (config->cache_type_v >= 0) { ++ params.type_v = static_cast(config->cache_type_v); ++ } ++ if (config->op_offload != SKIPPY_TRISTATE_AUTO) { ++ params.op_offload = config->op_offload == SKIPPY_TRISTATE_TRUE; ++ } +} diff --git a/src/skippy/model_load_config_internal.h b/src/skippy/model_load_config_internal.h new file mode 100644 @@ -555,7 +563,7 @@ new file mode 100644 index 000000000..7eb4d5e8c --- /dev/null +++ b/src/skippy/model_loading.cpp -@@ -0,0 +1,950 @@ +@@ -0,0 +1,948 @@ +#include "skippy.h" +#include "skippy/errors.h" +#include "skippy/environment.h" @@ -1052,6 +1060,8 @@ index 000000000..7eb4d5e8c + +struct skippy_runtime_config skippy_runtime_config_default(void) { + struct skippy_runtime_config config = {}; ++ config.cache_type_k = GGML_TYPE_F16; ++ config.cache_type_v = GGML_TYPE_F16; + config.kv_offload = SKIPPY_TRISTATE_AUTO; + config.kv_unified = SKIPPY_TRISTATE_AUTO; + config.swa_full = SKIPPY_TRISTATE_AUTO; @@ -1130,8 +1140,6 @@ index 000000000..7eb4d5e8c + params.swa_full = skippy_resolve_tristate(config->swa_full, params.swa_full); + skippy_apply_context_hardware_config(config, params); + } -+ params.type_k = config != nullptr && config->cache_type_k > 0 ? static_cast(config->cache_type_k) : GGML_TYPE_F16; -+ params.type_v = config != nullptr && config->cache_type_v > 0 ? static_cast(config->cache_type_v) : GGML_TYPE_F16; + params.flash_attn_type = config != nullptr ? static_cast(config->flash_attn_type) : LLAMA_FLASH_ATTN_TYPE_AUTO; + params.embeddings = config != nullptr && config->filter_tensors_on_load && !config->include_output; + const bool glm_dsa_op_timing_enabled = skippy_glm_dsa_op_timing_enabled(); @@ -1488,8 +1496,6 @@ index 000000000..7eb4d5e8c + ctx_params.swa_full = skippy_resolve_tristate(config->swa_full, ctx_params.swa_full); + skippy_apply_context_hardware_config(config, ctx_params); + } -+ ctx_params.type_k = config != nullptr && config->cache_type_k > 0 ? static_cast(config->cache_type_k) : GGML_TYPE_F16; -+ ctx_params.type_v = config != nullptr && config->cache_type_v > 0 ? static_cast(config->cache_type_v) : GGML_TYPE_F16; + ctx_params.flash_attn_type = config != nullptr ? static_cast(config->flash_attn_type) : LLAMA_FLASH_ATTN_TYPE_AUTO; + + target_model->mtp_ctx = llama_init_from_model(draft_model, ctx_params); @@ -3123,4 +3129,3 @@ index 000000000..a82b64408 +}; -- 2.55.0 - diff --git a/third_party/llama.cpp/patches/0006-test-skippy-cover-graph-planning-and-stage-contracts.patch b/third_party/llama.cpp/patches/0006-test-skippy-cover-graph-planning-and-stage-contracts.patch index 649c0531d2..2c3e44bdc9 100644 --- a/third_party/llama.cpp/patches/0006-test-skippy-cover-graph-planning-and-stage-contracts.patch +++ b/third_party/llama.cpp/patches/0006-test-skippy-cover-graph-planning-and-stage-contracts.patch @@ -682,7 +682,7 @@ new file mode 100644 index 000000000..1a1bb0ca5 --- /dev/null +++ b/src/skippy/tests/hardware_application_probe.cpp -@@ -0,0 +1,64 @@ +@@ -0,0 +1,76 @@ +#include "skippy/model_load_config_internal.h" +#include "skippy/runtime.h" + @@ -698,6 +698,10 @@ index 000000000..1a1bb0ca5 +int main() { + skippy_runtime_config config = skippy_runtime_config_default(); + bool passed = true; ++ passed &= require(config.cache_type_k == GGML_TYPE_F16, ++ "default K cache type must be F16"); ++ passed &= require(config.cache_type_v == GGML_TYPE_F16, ++ "default V cache type must be F16"); + passed &= require(config.kv_offload == SKIPPY_TRISTATE_AUTO, + "default kv_offload must preserve the derived default"); + passed &= require(config.kv_unified == SKIPPY_TRISTATE_AUTO, @@ -737,6 +741,14 @@ index 000000000..1a1bb0ca5 + skippy_apply_context_hardware_config(&config, context_params); + passed &= require(context_params.op_offload, "auto op_offload must preserve the derived default"); + ++ config.cache_type_k = GGML_TYPE_F32; ++ config.cache_type_v = GGML_TYPE_F32; ++ skippy_apply_context_hardware_config(&config, context_params); ++ passed &= require(context_params.type_k == GGML_TYPE_F32, ++ "explicit F32 K cache type must not fall back to F16"); ++ passed &= require(context_params.type_v == GGML_TYPE_F32, ++ "explicit F32 V cache type must not fall back to F16"); ++ + config.op_offload = SKIPPY_TRISTATE_FALSE; + skippy_apply_context_hardware_config(&config, context_params); + passed &= require(!context_params.op_offload, "false op_offload override was not applied"); @@ -3786,4 +3798,3 @@ index 000000000..380cf4976 +} -- 2.55.0 - diff --git a/third_party/llama.cpp/patches/0034-skippy-define-CacheGen-page-import-ABI.patch b/third_party/llama.cpp/patches/0034-skippy-define-CacheGen-page-import-ABI.patch new file mode 100644 index 0000000000..fd806a0b2a --- /dev/null +++ b/third_party/llama.cpp/patches/0034-skippy-define-CacheGen-page-import-ABI.patch @@ -0,0 +1,87 @@ +From 216d925e721653b7134b46c175e679cc5d0f2eeb Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 11 Sep 2026 16:02:11 +1000 +Subject: [PATCH] skippy: define CacheGen page import ABI + +Expose validated record descriptors for direct decode into resident KV storage. + +Assisted-by: scama +--- + include/skippy/common.h | 4 +++- + include/skippy/state.h | 31 +++++++++++++++++++++++++++++++ + 2 files changed, 34 insertions(+), 1 deletion(-) + +diff --git a/include/skippy/common.h b/include/skippy/common.h +index 65cfa2ce8..be996397a 100644 +--- a/include/skippy/common.h ++++ b/include/skippy/common.h +@@ -46,7 +46,7 @@ extern "C" { + + #define SKIPPY_ABI_VERSION_MAJOR 0 + #define SKIPPY_ABI_VERSION_MINOR 1 +-#define SKIPPY_ABI_VERSION_PATCH 54 ++#define SKIPPY_ABI_VERSION_PATCH 55 + + #if defined(_MSC_VER) + #define SKIPPY_DEPRECATED(message) __declspec(deprecated(message)) +@@ -104,6 +104,8 @@ enum skippy_feature { + #define SKIPPY_FEATURE_DEVICE_EVENTS ((uint64_t)1 << 34) + #define SKIPPY_FEATURE_DIAGNOSTIC_EVENTS ((uint64_t)1 << 35) + #define SKIPPY_FEATURE_UNLOAD_EVENTS ((uint64_t)1 << 36) ++ ++#define SKIPPY_FEATURE_CACHEGEN_KV_PAGE (UINT64_C(1) << 37) + + enum skippy_status { + SKIPPY_STATUS_OK = 0, + SKIPPY_STATUS_ERROR = 1, +diff --git a/include/skippy/state.h b/include/skippy/state.h +index aa698c124..7504e573c 100644 +--- a/include/skippy/state.h ++++ b/include/skippy/state.h +@@ -27,6 +27,29 @@ enum skippy_kv_page_component_role { + SKIPPY_KV_PAGE_COMPONENT_SWA = 2, + }; + ++#define SKIPPY_CACHEGEN_RECORD_V1_ABI_VERSION 1 ++ ++enum skippy_cachegen_record_kind { ++ SKIPPY_CACHEGEN_RECORD_F16 = 0, ++ SKIPPY_CACHEGEN_RECORD_EXACT = 1, ++ SKIPPY_CACHEGEN_RECORD_F16_TRANSPOSED = 2, ++}; ++ ++/** @brief Describes one validated CacheGen record and its logical page destination. */ ++struct skippy_cachegen_record_v1 { ++ uint32_t abi_version; ++ uint32_t kind; ++ uint32_t element_bytes; ++ uint32_t reserved0; ++ uint64_t output_offset; ++ uint64_t decoded_bytes; ++ uint64_t token_count; ++ uint64_t token_start; ++ uint64_t total_tokens; ++ const void * payload; ++ size_t payload_bytes; ++}; ++ + /** @brief Describes one independently ranged cache in a composite KV page. */ + struct skippy_kv_page_component_desc { + uint32_t version; +@@ -153,6 +176,14 @@ LLAMA_API enum skippy_status skippy_import_kv_page( + size_t input_bytes, + struct skippy_error ** out_error); + ++/** @brief Imports validated CacheGen records directly into resident KV storage. */ ++LLAMA_API enum skippy_status skippy_import_cachegen_kv_page_v1( ++ struct skippy_session * session, ++ const struct skippy_kv_page_desc * desc, ++ const struct skippy_cachegen_record_v1 * records, ++ size_t record_count, ++ struct skippy_error ** out_error); ++ + /** @brief Saves a session prefix into a resident cache sequence. */ + LLAMA_API enum skippy_status skippy_session_save_prefix( + struct skippy_session * session, +-- +2.54.0 (Apple Git-157) diff --git a/third_party/llama.cpp/patches/0035-skippy-expose-CacheGen-backend-capability.patch b/third_party/llama.cpp/patches/0035-skippy-expose-CacheGen-backend-capability.patch new file mode 100644 index 0000000000..bf63211633 --- /dev/null +++ b/third_party/llama.cpp/patches/0035-skippy-expose-CacheGen-backend-capability.patch @@ -0,0 +1,63 @@ +From c95fb1f1001f0e7bce75676e4a6dc3b82b430d14 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 11 Sep 2026 16:02:47 +1000 +Subject: [PATCH] skippy: expose CacheGen backend capability + +Advertise the compressed-page contract and return an explicit unsupported status until a resident backend decoder is available. + +Assisted-by: scama +--- + src/skippy/abi.cpp | 3 ++- + src/skippy/state.cpp | 23 +++++++++++++++++++++++ + 2 files changed, 25 insertions(+), 1 deletion(-) + +diff --git a/src/skippy/abi.cpp b/src/skippy/abi.cpp +index ec2b47444..2a0b1b097 100644 +--- a/src/skippy/abi.cpp ++++ b/src/skippy/abi.cpp +@@ -46,7 +46,8 @@ uint64_t skippy_abi_features(void) { + SKIPPY_FEATURE_KV_EVENTS | + SKIPPY_FEATURE_DEVICE_EVENTS | + SKIPPY_FEATURE_DIAGNOSTIC_EVENTS | +- SKIPPY_FEATURE_UNLOAD_EVENTS; ++ SKIPPY_FEATURE_UNLOAD_EVENTS | ++ SKIPPY_FEATURE_CACHEGEN_KV_PAGE; + } + + void skippy_error_free(struct skippy_error * error) { +diff --git a/src/skippy/state.cpp b/src/skippy/state.cpp +index 620ee3578..af3a56601 100644 +--- a/src/skippy/state.cpp ++++ b/src/skippy/state.cpp +@@ -657,6 +657,29 @@ enum skippy_status skippy_import_kv_page( + return skippy_success(out_error); + } + ++enum skippy_status skippy_import_cachegen_kv_page_v1( ++ struct skippy_session * session, ++ const struct skippy_kv_page_desc * desc, ++ const struct skippy_cachegen_record_v1 * records, ++ size_t record_count, ++ struct skippy_error ** out_error) { ++ if (session == nullptr || session->ctx == nullptr || desc == nullptr || records == nullptr || record_count == 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "session, descriptor, and CacheGen records are required"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ for (size_t i = 0; i < record_count; ++i) { ++ const auto & record = records[i]; ++ if (record.abi_version != SKIPPY_CACHEGEN_RECORD_V1_ABI_VERSION || record.reserved0 != 0 || ++ record.payload == nullptr || record.payload_bytes == 0 || record.decoded_bytes == 0 || record.token_count == 0 || ++ record.kind > SKIPPY_CACHEGEN_RECORD_F16_TRANSPOSED) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "invalid CacheGen record descriptor"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ } ++ skippy_set_error(out_error, SKIPPY_STATUS_UNSUPPORTED, "CacheGen device import is unavailable on this backend"); ++ return SKIPPY_STATUS_UNSUPPORTED; ++} ++ + enum skippy_status skippy_retire_verify_checkpoint( + skippy_session * session, + uint64_t token_start, +-- +2.54.0 (Apple Git-157) diff --git a/third_party/llama.cpp/patches/0036-skippy-dispatch-CacheGen-pages-into-resident-KV.patch b/third_party/llama.cpp/patches/0036-skippy-dispatch-CacheGen-pages-into-resident-KV.patch new file mode 100644 index 0000000000..6312aed656 --- /dev/null +++ b/third_party/llama.cpp/patches/0036-skippy-dispatch-CacheGen-pages-into-resident-KV.patch @@ -0,0 +1,554 @@ +From 562610033432885cdbae80b4fe7d038f2e070ae1 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 11 Sep 2026 16:42:03 +1000 +Subject: [PATCH] skippy: dispatch CacheGen pages into resident KV + +Assisted-by: scama +--- + ggml/include/ggml-backend.h | 27 +++ + src/llama-kv-cache.cpp | 338 ++++++++++++++++++++++++++++++++++++ + src/llama-kv-cache.h | 12 ++ + src/skippy/state.cpp | 105 +++++++++-- + 4 files changed, 466 insertions(+), 16 deletions(-) + +diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h +index cc3f8cd36..e2460fcda 100644 +--- a/ggml/include/ggml-backend.h ++++ b/ggml/include/ggml-backend.h +@@ -224,6 +224,33 @@ extern "C" { + }; + typedef struct ggml_backend_feature * (*ggml_backend_get_features_t)(ggml_backend_reg_t reg); + ++ // Optional backend entry point for decoding validated CacheGen F16 tiles ++ // directly into resident tensors. Callers resolve ++ // `ggml_backend_cachegen_decode_f16` through get_proc_address. ++ struct ggml_backend_cachegen_tile { ++ const void * payload; ++ size_t payload_bytes; ++ uint32_t token_offset; ++ uint32_t token_count; ++ }; ++ ++ struct ggml_backend_cachegen_job { ++ struct ggml_tensor * dst; ++ const struct ggml_backend_cachegen_tile * tiles; ++ size_t tile_count; ++ const uint32_t * cells; ++ size_t cell_count; ++ uint32_t channels; ++ uint64_t token_stride; ++ uint64_t channel_stride; ++ }; ++ ++ typedef bool (*ggml_backend_cachegen_decode_f16_t)( ++ const struct ggml_backend_cachegen_job * jobs, ++ size_t job_count, ++ char * error, ++ size_t error_capacity); ++ + // + // Backend registry + // +diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp +index a7d65f23f..dfd00139b 100644 +--- a/src/llama-kv-cache.cpp ++++ b/src/llama-kv-cache.cpp +@@ -1958,6 +1958,344 @@ bool llama_kv_cache::stage_import_kv_page( + return src == static_cast(input) + input_bytes; + } + ++static uint16_t skippy_cachegen_read_u16(const uint8_t * bytes) { ++ return static_cast(bytes[0]) | static_cast(bytes[1]) << 8; ++} ++ ++static uint32_t skippy_cachegen_read_u32(const uint8_t * bytes) { ++ return static_cast(bytes[0]) | static_cast(bytes[1]) << 8 | ++ static_cast(bytes[2]) << 16 | static_cast(bytes[3]) << 24; ++} ++ ++static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & record, ++ size_t expected_rows, ++ size_t expected_channels, ++ std::string & error) { ++ constexpr size_t header_bytes = 16; ++ constexpr size_t cdf_entries = 33; ++ constexpr size_t max_rows = 256; ++ const auto * payload = static_cast(record.payload); ++ if (record.abi_version != SKIPPY_CACHEGEN_RECORD_V1_ABI_VERSION || record.reserved0 != 0 || ++ record.element_bytes != 2 || payload == nullptr || record.payload_bytes < header_bytes || ++ record.decoded_bytes == 0 || record.token_count != expected_rows || expected_rows == 0 || ++ expected_rows > max_rows || expected_channels == 0 || ++ expected_channels > std::numeric_limits::max()) { ++ error = "invalid CacheGen F16 record descriptor"; ++ return false; ++ } ++ if (std::memcmp(payload, "LCG1", 4) != 0 || payload[5] != 0 || (payload[4] != 16 && payload[4] != 32) || ++ skippy_cachegen_read_u16(payload + 6) != expected_rows || ++ skippy_cachegen_read_u32(payload + 8) != expected_channels) { ++ error = "CacheGen record disagrees with its segment header"; ++ return false; ++ } ++ if (expected_rows > std::numeric_limits::max() / expected_channels || ++ expected_rows * expected_channels > std::numeric_limits::max() / 2 || ++ record.decoded_bytes != expected_rows * expected_channels * 2) { ++ error = "CacheGen record decoded geometry is inconsistent"; ++ return false; ++ } ++ const size_t max_bytes = expected_rows * sizeof(float); ++ if (expected_channels > std::numeric_limits::max() / (cdf_entries * sizeof(uint16_t)) || ++ expected_channels > std::numeric_limits::max() / sizeof(uint16_t)) { ++ error = "CacheGen segment metadata overflows"; ++ return false; ++ } ++ const size_t cdf_bytes = expected_channels * cdf_entries * sizeof(uint16_t); ++ const size_t length_bytes = expected_channels * sizeof(uint16_t); ++ if (header_bytes > std::numeric_limits::max() - max_bytes || ++ header_bytes + max_bytes > std::numeric_limits::max() - cdf_bytes || ++ header_bytes + max_bytes + cdf_bytes > std::numeric_limits::max() - length_bytes) { ++ error = "CacheGen segment metadata range overflows"; ++ return false; ++ } ++ const size_t cdf_offset = header_bytes + max_bytes; ++ const size_t length_offset = cdf_offset + cdf_bytes; ++ const size_t stream_offset = length_offset + length_bytes; ++ const size_t stream_bytes = skippy_cachegen_read_u32(payload + 12); ++ if (stream_offset > record.payload_bytes || stream_bytes != record.payload_bytes - stream_offset) { ++ error = "CacheGen segment stream length is inconsistent"; ++ return false; ++ } ++ for (size_t row = 0; row < expected_rows; ++row) { ++ const uint32_t bits = skippy_cachegen_read_u32(payload + header_bytes + row * sizeof(float)); ++ float value; ++ std::memcpy(&value, &bits, sizeof(value)); ++ if (!std::isfinite(value) || value < 0.0f) { ++ error = "CacheGen segment contains an invalid row maximum"; ++ return false; ++ } ++ } ++ for (size_t channel = 0; channel < expected_channels; ++channel) { ++ const uint8_t * cdf = payload + cdf_offset + channel * cdf_entries * sizeof(uint16_t); ++ uint16_t previous = skippy_cachegen_read_u16(cdf); ++ if (previous != 0) { ++ error = "CacheGen segment has an invalid CDF endpoint"; ++ return false; ++ } ++ for (size_t entry = 1; entry < cdf_entries; ++entry) { ++ const uint16_t current = skippy_cachegen_read_u16(cdf + entry * sizeof(uint16_t)); ++ if (current <= previous) { ++ error = "CacheGen segment CDF is not strictly increasing"; ++ return false; ++ } ++ previous = current; ++ } ++ if (previous != std::numeric_limits::max()) { ++ error = "CacheGen segment has an invalid CDF endpoint"; ++ return false; ++ } ++ } ++ size_t summed_stream_bytes = 0; ++ for (size_t channel = 0; channel < expected_channels; ++channel) { ++ const size_t length = skippy_cachegen_read_u16(payload + length_offset + channel * sizeof(uint16_t)); ++ if (length == 0 || length > max_rows || summed_stream_bytes > stream_bytes || ++ length > stream_bytes - summed_stream_bytes) { ++ error = "CacheGen segment has an invalid channel stream length"; ++ return false; ++ } ++ summed_stream_bytes += length; ++ } ++ if (summed_stream_bytes != stream_bytes) { ++ error = "CacheGen channel lengths do not cover the segment stream"; ++ return false; ++ } ++ return true; ++} ++ ++struct skippy_cachegen_job_storage { ++ ggml_tensor * dst = nullptr; ++ std::vector tiles; ++ uint32_t channels = 0; ++ uint64_t token_stride = 0; ++ uint64_t channel_stride = 0; ++}; ++ ++bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id seq_id, ++ const skippy_kv_page_desc & desc, ++ const skippy_cachegen_record_v1 * records, ++ size_t record_count, ++ size_t output_base, ++ std::string & error, ++ bool & unsupported, ++ bool & invalid_argument, ++ bool validate_only) { ++ unsupported = false; ++ invalid_argument = true; ++ uint8_t validation_sentinel = 0; ++ if (!stage_import_kv_page(seq_id, desc, &validation_sentinel, desc.payload_bytes, error, true)) { ++ return false; ++ } ++ if (records == nullptr || record_count == 0 || desc.k_type != GGML_TYPE_F16 || desc.v_type != GGML_TYPE_F16) { ++ error = "CacheGen import requires non-empty F16 K/V records"; ++ return false; ++ } ++ ++ const uint32_t strm = seq_to_stream[seq_id]; ++ std::vector selected; ++ for (const auto & layer : layers) { ++ if (layer.il >= static_cast(desc.layer_start) && layer.il < static_cast(desc.layer_end) && ++ layer.k_stream.size() > strm && layer.k_stream[strm] != nullptr) { ++ selected.push_back(&layer); ++ } ++ } ++ ++ size_t record_index = 0; ++ size_t logical_offset = output_base; ++ std::vector storage; ++ storage.reserve(selected.size() * 2); ++ auto consume_tiles = [&](ggml_tensor * dst, size_t row_bytes, bool transposed) -> bool { ++ skippy_cachegen_job_storage entry; ++ entry.dst = dst; ++ entry.token_stride = transposed ? desc.v_element_bytes : row_bytes; ++ entry.channel_stride = transposed ? static_cast(v_cells[strm].size()) * desc.v_element_bytes : 2; ++ const size_t channels = row_bytes / 2; ++ entry.channels = static_cast(channels); ++ size_t token_offset = 0; ++ while (token_offset < desc.token_count) { ++ if (record_index >= record_count) { ++ error = "CacheGen records do not cover the native KV page"; ++ return false; ++ } ++ const auto & record = records[record_index++]; ++ const size_t rows = static_cast(record.token_count); ++ const uint32_t expected_kind = ++ transposed ? SKIPPY_CACHEGEN_RECORD_F16_TRANSPOSED : SKIPPY_CACHEGEN_RECORD_F16; ++ if (record.kind != expected_kind || rows == 0 || rows > desc.token_count - token_offset || ++ record.output_offset != logical_offset + (transposed ? 0 : token_offset * row_bytes) || ++ record.token_start != (transposed ? token_offset : 0) || ++ record.total_tokens != (transposed ? desc.token_count : 0) || ++ !skippy_cachegen_validate_segment(record, rows, channels, error)) { ++ if (error.empty()) { ++ error = "CacheGen record order or destination geometry is invalid"; ++ } ++ return false; ++ } ++ entry.tiles.push_back({ record.payload, record.payload_bytes, static_cast(token_offset), ++ static_cast(rows) }); ++ token_offset += rows; ++ } ++ storage.push_back(std::move(entry)); ++ logical_offset += static_cast(desc.token_count) * row_bytes; ++ return true; ++ }; ++ ++ for (const auto * layer : selected) { ++ if (!consume_tiles(layer->k_stream[strm], desc.k_row_bytes, false)) { ++ return false; ++ } ++ } ++ for (const auto * layer : selected) { ++ const size_t v_row_bytes = ++ v_trans ? static_cast(hparams.n_embd_v_gqa(layer->il)) * desc.v_element_bytes : desc.v_row_bytes; ++ if (!consume_tiles(layer->v_stream[strm], v_row_bytes, v_trans)) { ++ return false; ++ } ++ } ++ ++ const skippy_cachegen_record_v1 * exact = nullptr; ++ if (desc.k_idx_row_bytes != 0 && (desc.token_count > std::numeric_limits::max() / desc.k_idx_row_bytes || ++ static_cast(desc.token_count) * desc.k_idx_row_bytes > ++ std::numeric_limits::max() / selected.size())) { ++ error = "CacheGen exact indexer geometry overflows"; ++ return false; ++ } ++ const size_t exact_bytes = static_cast(desc.token_count) * desc.k_idx_row_bytes * selected.size(); ++ if (exact_bytes > 0) { ++ if (record_index >= record_count) { ++ error = "CacheGen page is missing exact indexer state"; ++ return false; ++ } ++ exact = &records[record_index++]; ++ if (exact->abi_version != SKIPPY_CACHEGEN_RECORD_V1_ABI_VERSION || ++ exact->kind != SKIPPY_CACHEGEN_RECORD_EXACT || exact->element_bytes != 1 || exact->reserved0 != 0 || ++ exact->output_offset != logical_offset || exact->decoded_bytes != exact_bytes || ++ exact->token_count != desc.token_count || exact->token_start != 0 || exact->total_tokens != 0 || ++ exact->payload == nullptr || exact->payload_bytes != exact_bytes) { ++ error = "CacheGen exact indexer record is invalid"; ++ return false; ++ } ++ logical_offset += exact_bytes; ++ } ++ if (record_index != record_count || logical_offset != output_base + desc.payload_bytes) { ++ error = "CacheGen records do not exactly cover the native KV page"; ++ return false; ++ } ++ ++ ggml_backend_cachegen_decode_f16_t decode = nullptr; ++ for (const auto & entry : storage) { ++ if (entry.dst == nullptr || entry.dst->buffer == nullptr) { ++ error = "CacheGen destination tensor has no backend buffer"; ++ return false; ++ } ++ const auto buft = ggml_backend_buffer_get_type(entry.dst->buffer); ++ const auto device = ggml_backend_buft_get_device(buft); ++ const auto reg = device == nullptr ? nullptr : ggml_backend_dev_backend_reg(device); ++ auto candidate = reg == nullptr ? ++ nullptr : ++ reinterpret_cast( ++ ggml_backend_reg_get_proc_address(reg, "ggml_backend_cachegen_decode_f16")); ++ if (candidate == nullptr || (decode != nullptr && candidate != decode)) { ++ invalid_argument = false; ++ unsupported = true; ++ error = "resident KV backend does not provide CacheGen F16 decode"; ++ return false; ++ } ++ decode = candidate; ++ } ++ invalid_argument = false; ++ if (validate_only) { ++ return true; ++ } ++ ++ auto & stream_cells = v_cells[strm]; ++ for (uint32_t i = 0; i < stream_cells.size(); ++i) { ++ if (stream_cells.seq_has(i, seq_id)) { ++ invalid_argument = true; ++ error = "CacheGen page import requires an empty session KV sequence"; ++ return false; ++ } ++ } ++ ++ llama_ubatch ubatch = {}; ++ auto udata = std::make_shared(); ++ const uint32_t n_tokens = static_cast(desc.token_count); ++ const uint32_t n_pos = hparams.n_pos_per_embd(); ++ if (n_pos != 0 && n_tokens > std::numeric_limits::max() / n_pos) { ++ invalid_argument = true; ++ error = "CacheGen page position geometry overflows"; ++ return false; ++ } ++ udata->token.resize(n_tokens); ++ udata->pos.resize(static_cast(n_tokens) * n_pos); ++ udata->n_seq_id.resize(n_tokens, 1); ++ udata->seq_id.resize(n_tokens); ++ udata->seq_id_unq = { seq_id }; ++ udata->seq_idx.resize(LLAMA_MAX_SEQ, -1); ++ udata->seq_idx[seq_id] = 0; ++ udata->output.resize(n_tokens, 0); ++ udata->seq_id_data.resize(n_tokens, 0); ++ for (uint32_t i = 0; i < n_tokens; ++i) { ++ const llama_pos pos = static_cast(desc.token_start + i); ++ for (uint32_t j = 0; j < n_pos; ++j) { ++ udata->pos[static_cast(j) * n_tokens + i] = pos; ++ } ++ udata->seq_id_data[i] = seq_id; ++ udata->seq_id[i] = &udata->seq_id_data[i]; ++ } ++ ubatch.b_equal_seqs = true; ++ ubatch.n_tokens = n_tokens; ++ ubatch.n_seq_tokens = n_tokens; ++ ubatch.n_seqs = 1; ++ ubatch.n_seqs_unq = 1; ++ ubatch.n_pos = n_pos; ++ ubatch.token = udata->token.data(); ++ ubatch.pos = udata->pos.data(); ++ ubatch.n_seq_id = udata->n_seq_id.data(); ++ ubatch.seq_id = udata->seq_id.data(); ++ ubatch.seq_id_unq = udata->seq_id_unq.data(); ++ ubatch.seq_idx = udata->seq_idx.data(); ++ ubatch.output = udata->output.data(); ++ ubatch.data = std::move(udata); ++ ++ const uint32_t head_before = v_heads[strm]; ++ const slot_info sinfo = find_slot(ubatch, false); ++ if (sinfo.empty()) { ++ error = "failed to allocate KV cache cells for CacheGen page"; ++ return false; ++ } ++ apply_ubatch(sinfo, ubatch); ++ ++ std::vector jobs; ++ jobs.reserve(storage.size()); ++ for (const auto & entry : storage) { ++ jobs.push_back({ entry.dst, entry.tiles.data(), entry.tiles.size(), sinfo.idxs[0].data(), sinfo.idxs[0].size(), ++ entry.channels, entry.token_stride, entry.channel_stride }); ++ } ++ char backend_error[256] = {}; ++ if (!decode(jobs.data(), jobs.size(), backend_error, sizeof(backend_error))) { ++ seq_rm(seq_id, static_cast(desc.token_start), ++ static_cast(desc.token_start + desc.token_count)); ++ v_heads[strm] = head_before; ++ error = backend_error[0] == '\0' ? "CacheGen backend decode failed" : backend_error; ++ return false; ++ } ++ ++ if (exact != nullptr) { ++ const char * src = static_cast(exact->payload); ++ const auto cell_runs = llama_kv_cache_cell_runs(sinfo.idxs[0]); ++ for (const auto * layer : selected) { ++ auto * k_idx = layer->k_idx_stream[strm]; ++ for (const auto & run : cell_runs) { ++ const size_t run_bytes = static_cast(run.count) * desc.k_idx_row_bytes; ++ ggml_backend_tensor_set(k_idx, src, static_cast(run.first) * desc.k_idx_row_bytes, run_bytes); ++ src += run_bytes; ++ } ++ } ++ } ++ return true; ++} ++ + uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const { + uint32_t result = 0; + +diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h +index 8368e0ee2..ba95e63dc 100644 +--- a/src/llama-kv-cache.h ++++ b/src/llama-kv-cache.h +@@ -14,6 +14,7 @@ struct llama_hparams; + struct llama_model; + struct llama_context; + struct skippy_kv_page_desc; ++struct skippy_cachegen_record_v1; + + // + // llama_kv_cache +@@ -209,6 +210,17 @@ public: + std::string & error, + bool validate_only = false); + ++ bool stage_import_cachegen_kv_page( ++ llama_seq_id seq_id, ++ const skippy_kv_page_desc & desc, ++ const skippy_cachegen_record_v1 * records, ++ size_t record_count, ++ size_t output_base, ++ std::string & error, ++ bool & unsupported, ++ bool & invalid_argument, ++ bool validate_only = false); ++ + // + // graph_build API + // +diff --git a/src/skippy/state.cpp b/src/skippy/state.cpp +index af3a56601..dc2203735 100644 +--- a/src/skippy/state.cpp ++++ b/src/skippy/state.cpp +@@ -657,27 +657,100 @@ enum skippy_status skippy_import_kv_page( + return skippy_success(out_error); + } + +-enum skippy_status skippy_import_cachegen_kv_page_v1( +- struct skippy_session * session, +- const struct skippy_kv_page_desc * desc, +- const struct skippy_cachegen_record_v1 * records, +- size_t record_count, +- struct skippy_error ** out_error) { ++enum skippy_status skippy_import_cachegen_kv_page_v1(struct skippy_session * session, ++ const struct skippy_kv_page_desc * desc, ++ const struct skippy_cachegen_record_v1 * records, ++ size_t record_count, ++ struct skippy_error ** out_error) { + if (session == nullptr || session->ctx == nullptr || desc == nullptr || records == nullptr || record_count == 0) { +- skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "session, descriptor, and CacheGen records are required"); ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, ++ "session, descriptor, and CacheGen records are required"); + return SKIPPY_STATUS_INVALID_ARGUMENT; + } +- for (size_t i = 0; i < record_count; ++i) { +- const auto & record = records[i]; +- if (record.abi_version != SKIPPY_CACHEGEN_RECORD_V1_ABI_VERSION || record.reserved0 != 0 || +- record.payload == nullptr || record.payload_bytes == 0 || record.decoded_bytes == 0 || record.token_count == 0 || +- record.kind > SKIPPY_CACHEGEN_RECORD_F16_TRANSPOSED) { +- skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "invalid CacheGen record descriptor"); +- return SKIPPY_STATUS_INVALID_ARGUMENT; ++ if (session->n_past != 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "CacheGen page import requires a fresh session"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ ++ std::string error; ++ bool ok = false; ++ bool unsupported = false; ++ bool invalid_argument = false; ++ if (desc->codec == SKIPPY_KV_PAGE_CODEC_ISWA_COMPOSITE_V1) { ++ auto * iswa = skippy_get_iswa_cache(session); ++ if (!skippy_validate_composite_desc(*desc, desc->payload_bytes, error)) { ++ invalid_argument = true; ++ } else if (iswa == nullptr) { ++ unsupported = true; ++ error = "composite CacheGen page requires ISWA memory"; ++ } else { ++ const auto base = skippy_page_from_component(*desc, desc->components[0]); ++ const auto swa = skippy_page_from_component(*desc, desc->components[1]); ++ const uint64_t split_offset = desc->components[1].payload_offset; ++ size_t split = 0; ++ while (split < record_count && records[split].output_offset < split_offset) { ++ ++split; ++ } ++ if (split == 0 || split == record_count) { ++ invalid_argument = true; ++ error = "CacheGen records do not cover both ISWA components"; ++ } else { ++ bool base_unsupported = false; ++ bool swa_unsupported = false; ++ bool base_invalid = false; ++ bool swa_invalid = false; ++ const bool base_valid = iswa->get_base()->stage_import_cachegen_kv_page( ++ session->seq_id, base, records, split, 0, error, base_unsupported, base_invalid, true); ++ const bool swa_valid = ++ base_valid && iswa->get_swa()->stage_import_cachegen_kv_page( ++ session->seq_id, swa, records + split, record_count - split, ++ static_cast(split_offset), error, swa_unsupported, swa_invalid, true); ++ unsupported = base_unsupported || swa_unsupported; ++ invalid_argument = base_invalid || swa_invalid; ++ if (base_valid && swa_valid) { ++ bool ignored = false; ++ bool ignored_invalid = false; ++ ok = iswa->get_base()->stage_import_cachegen_kv_page(session->seq_id, base, records, split, 0, ++ error, ignored, ignored_invalid) && ++ iswa->get_swa()->stage_import_cachegen_kv_page( ++ session->seq_id, swa, records + split, record_count - split, ++ static_cast(split_offset), error, ignored, ignored_invalid); ++ invalid_argument = ignored_invalid; ++ } ++ } ++ } ++ if (!ok && iswa != nullptr) { ++ iswa->get_base()->seq_rm(session->seq_id, -1, -1); ++ iswa->get_swa()->seq_rm(session->seq_id, -1, -1); + } ++ } else if (desc->codec == 0 || desc->codec == SKIPPY_KV_PAGE_CODEC_SINGLE_V1) { ++ if (auto * kv = skippy_get_kv_cache(session, out_error)) { ++ ok = kv->stage_import_cachegen_kv_page(session->seq_id, *desc, records, record_count, 0, error, unsupported, ++ invalid_argument); ++ } else { ++ return SKIPPY_STATUS_UNSUPPORTED; ++ } ++ } else { ++ invalid_argument = true; ++ error = "unsupported native KV page codec"; ++ } ++ if (!ok) { ++ const enum skippy_status status = invalid_argument ? ++ SKIPPY_STATUS_INVALID_ARGUMENT : ++ (unsupported ? SKIPPY_STATUS_UNSUPPORTED : SKIPPY_STATUS_RUNTIME_ERROR); ++ skippy_set_error(out_error, status, error.c_str()); ++ return status; ++ } ++ if (llm_graph_result * graph = session->ctx->get_gf_res_prev()) { ++ graph->reset(); + } +- skippy_set_error(out_error, SKIPPY_STATUS_UNSUPPORTED, "CacheGen device import is unavailable on this backend"); +- return SKIPPY_STATUS_UNSUPPORTED; ++ session->n_past = std::max(session->n_past, ++ static_cast(std::min(desc->token_start + desc->token_count, ++ std::numeric_limits::max()))); ++ skippy_mtp_clear_session_state(session); ++ session->verify_checkpoints.clear(); ++ session->ctx->synchronize(); ++ return skippy_success(out_error); + } + + enum skippy_status skippy_retire_verify_checkpoint( +-- +2.54.0 (Apple Git-157) + diff --git a/third_party/llama.cpp/patches/0037-ggml-metal-decode-CacheGen-pages-into-resident-KV.patch b/third_party/llama.cpp/patches/0037-ggml-metal-decode-CacheGen-pages-into-resident-KV.patch new file mode 100644 index 0000000000..df7f68fddc --- /dev/null +++ b/third_party/llama.cpp/patches/0037-ggml-metal-decode-CacheGen-pages-into-resident-KV.patch @@ -0,0 +1,651 @@ +From 32d3370571b5b73f136e29f2b1c05bb137493954 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 11 Sep 2026 16:42:09 +1000 +Subject: [PATCH] ggml-metal: decode CacheGen pages into resident KV + +Assisted-by: scama +--- + ggml/src/ggml-metal/CMakeLists.txt | 1 + + ggml/src/ggml-metal/ggml-metal-device.h | 7 + + ggml/src/ggml-metal/ggml-metal-device.m | 208 +++++++++++++++++++ + ggml/src/ggml-metal/ggml-metal.cpp | 3 + + ggml/src/ggml-metal/kernels/cachegen.metal | 98 +++++++++ + tests/CMakeLists.txt | 4 + + tests/test-skippy-cachegen-metal.cpp | 224 +++++++++++++++++++++ + 7 files changed, 545 insertions(+) + create mode 100644 ggml/src/ggml-metal/kernels/cachegen.metal + create mode 100644 tests/test-skippy-cachegen-metal.cpp + +diff --git a/ggml/src/ggml-metal/CMakeLists.txt b/ggml/src/ggml-metal/CMakeLists.txt +index a661e710a..208c8e888 100644 +--- a/ggml/src/ggml-metal/CMakeLists.txt ++++ b/ggml/src/ggml-metal/CMakeLists.txt +@@ -30,6 +30,7 @@ set(METALLIB_KERNELS_DEQUANTIZE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/dequantize. + set(METALLIB_KERNELS_QUANTIZE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/quantize.h") + + set(METALLIB_KERNEL_SOURCES ++ kernels/cachegen.metal + kernels/fa.metal + kernels/mul_mv.metal + kernels/mul_mm.metal +diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h +index 31fc07d44..22e4fb436 100644 +--- a/ggml/src/ggml-metal/ggml-metal-device.h ++++ b/ggml/src/ggml-metal/ggml-metal-device.h +@@ -1,6 +1,7 @@ + #pragma once + + #include "ggml.h" ++#include "ggml-backend.h" + + #ifdef __cplusplus + extern "C" { +@@ -350,6 +351,12 @@ void ggml_metal_buffer_clear (ggml_metal_buffer_t buf, uint8_t value); + // + struct ggml_metal_buffer_id ggml_metal_buffer_get_id(ggml_metal_buffer_t buf, const struct ggml_tensor * t); + ++bool ggml_metal_cachegen_decode_f16( ++ const struct ggml_backend_cachegen_job * jobs, ++ size_t job_count, ++ char * error, ++ size_t error_capacity); ++ + #ifdef __cplusplus + } + #endif +diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m +index d6775211e..c6e22a863 100644 +--- a/ggml/src/ggml-metal/ggml-metal-device.m ++++ b/ggml/src/ggml-metal/ggml-metal-device.m +@@ -11,6 +11,7 @@ + + #include + #include ++#include + + #ifndef TARGET_OS_VISION + #define TARGET_OS_VISION 0 +@@ -110,6 +111,7 @@ int ggml_metal_pipeline_max_theads_per_threadgroup(struct ggml_metal_pipeline_wi + // X(suffix, name): name is both the kernels/.metal basename and the + // ggml_metallib__{start,end} embed-symbol stem. + #define GGML_METAL_LIBS \ ++ X(CACHEGEN, cachegen) \ + X(FA, fa) \ + X(MUL_MV, mul_mv) \ + X(MUL_MM, mul_mm) \ +@@ -2641,3 +2643,209 @@ struct ggml_metal_buffer_id ggml_metal_buffer_get_id(ggml_metal_buffer_t buf, co + + return res; + } ++ ++struct ggml_metal_cachegen_tile { ++ uint32_t payload_offset; ++ uint32_t prefix_offset; ++ uint32_t token_offset; ++ uint32_t rows; ++}; ++ ++struct ggml_metal_cachegen_params { ++ uint32_t channels; ++ uint32_t tile_count; ++ uint64_t token_stride; ++ uint64_t channel_stride; ++}; ++ ++static bool ggml_metal_cachegen_error(char * error, size_t capacity, const char * message) { ++ if (error != NULL && capacity > 0) { ++ snprintf(error, capacity, "%s", message); ++ } ++ return false; ++} ++ ++bool ggml_metal_cachegen_decode_f16( ++ const struct ggml_backend_cachegen_job * jobs, ++ size_t job_count, ++ char * error, ++ size_t error_capacity) { ++ if (jobs == NULL || job_count == 0) { ++ return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen jobs are required"); ++ } ++ ++ @autoreleasepool { ++ ggml_metal_device_t device = NULL; ++ id cmd_buf = nil; ++ id encoder = nil; ++ NSMutableArray * temporary_buffers = [[NSMutableArray alloc] init]; ++ for (size_t job_index = 0; job_index < job_count; ++job_index) { ++ const struct ggml_backend_cachegen_job * job = &jobs[job_index]; ++ if (job->dst == NULL || job->dst->buffer == NULL || job->tiles == NULL || job->tile_count == 0 || ++ job->tile_count > UINT32_MAX || job->cells == NULL || job->cell_count == 0 || job->channels == 0 || ++ job->dst->type != GGML_TYPE_F16 || job->token_stride == 0 || job->channel_stride == 0 || ++ job->cell_count > SIZE_MAX / sizeof(uint32_t)) { ++ [temporary_buffers release]; ++ return ggml_metal_cachegen_error(error, error_capacity, "invalid Metal CacheGen job"); ++ } ++ uint64_t max_cell = 0; ++ for (size_t cell_index = 0; cell_index < job->cell_count; ++cell_index) { ++ max_cell = MAX(max_cell, job->cells[cell_index]); ++ } ++ const uint64_t max_channel = (uint64_t) job->channels - 1; ++ if (max_cell > (UINT64_MAX - sizeof(uint16_t)) / job->token_stride || ++ max_channel > (UINT64_MAX - max_cell * job->token_stride - sizeof(uint16_t)) / job->channel_stride || ++ max_cell * job->token_stride + max_channel * job->channel_stride + sizeof(uint16_t) > ggml_nbytes(job->dst)) { ++ [temporary_buffers release]; ++ return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen destination geometry is out of bounds"); ++ } ++ ggml_metal_buffer_t dst_buffer = (ggml_metal_buffer_t) job->dst->buffer->context; ++ if (dst_buffer == NULL || (device != NULL && device != dst_buffer->dev)) { ++ [temporary_buffers release]; ++ return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen jobs span incompatible devices"); ++ } ++ if (device == NULL) { ++ device = dst_buffer->dev; ++ cmd_buf = [device->mtl_queue commandBufferWithUnretainedReferences]; ++ encoder = [cmd_buf computeCommandEncoder]; ++ if (cmd_buf == nil || encoder == nil) { ++ [temporary_buffers release]; ++ return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen command allocation failed"); ++ } ++ } ++ ++ NSMutableData * payload_data = [[NSMutableData alloc] init]; ++ NSMutableData * tile_data = [[NSMutableData alloc] init]; ++ NSMutableData * prefix_data = [[NSMutableData alloc] init]; ++ for (size_t tile_index = 0; tile_index < job->tile_count; ++tile_index) { ++ const struct ggml_backend_cachegen_tile * tile = &job->tiles[tile_index]; ++ if (tile->payload == NULL || tile->payload_bytes < 16 || tile->token_count == 0 || ++ tile->token_offset > job->cell_count || tile->token_count > job->cell_count - tile->token_offset || ++ payload_data.length > UINT32_MAX || tile->payload_bytes > UINT32_MAX - payload_data.length || ++ prefix_data.length / sizeof(uint32_t) > UINT32_MAX - ((size_t) job->channels + 1)) { ++ [payload_data release]; ++ [tile_data release]; ++ [prefix_data release]; ++ [encoder endEncoding]; ++ [temporary_buffers release]; ++ return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen tile exceeds bounded geometry"); ++ } ++ const uint8_t * payload = (const uint8_t *) tile->payload; ++ const size_t lengths_offset = 16 + (size_t) tile->token_count * sizeof(float) + ++ (size_t) job->channels * 33 * sizeof(uint16_t); ++ const size_t streams_offset = lengths_offset + (size_t) job->channels * sizeof(uint16_t); ++ if (streams_offset > tile->payload_bytes) { ++ [payload_data release]; ++ [tile_data release]; ++ [prefix_data release]; ++ [encoder endEncoding]; ++ [temporary_buffers release]; ++ return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen tile metadata is truncated"); ++ } ++ struct ggml_metal_cachegen_tile encoded_tile = { ++ (uint32_t) payload_data.length, ++ (uint32_t) (prefix_data.length / sizeof(uint32_t)), ++ tile->token_offset, ++ tile->token_count, ++ }; ++ uint32_t prefix = 0; ++ [prefix_data appendBytes:&prefix length:sizeof(prefix)]; ++ for (uint32_t channel = 0; channel < job->channels; ++channel) { ++ const size_t offset = lengths_offset + (size_t) channel * sizeof(uint16_t); ++ const uint32_t length = (uint32_t) payload[offset] | (uint32_t) payload[offset + 1] << 8; ++ if (prefix > UINT32_MAX - length) { ++ [payload_data release]; ++ [tile_data release]; ++ [prefix_data release]; ++ [encoder endEncoding]; ++ [temporary_buffers release]; ++ return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen stream offsets overflow"); ++ } ++ prefix += length; ++ [prefix_data appendBytes:&prefix length:sizeof(prefix)]; ++ } ++ const uint32_t declared_stream_bytes = (uint32_t) payload[12] | ++ (uint32_t) payload[13] << 8 | (uint32_t) payload[14] << 16 | (uint32_t) payload[15] << 24; ++ if (prefix != declared_stream_bytes || prefix != tile->payload_bytes - streams_offset) { ++ [payload_data release]; ++ [tile_data release]; ++ [prefix_data release]; ++ [encoder endEncoding]; ++ [temporary_buffers release]; ++ return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen stream length is inconsistent"); ++ } ++ [tile_data appendBytes:&encoded_tile length:sizeof(encoded_tile)]; ++ [payload_data appendBytes:tile->payload length:tile->payload_bytes]; ++ } ++ ++ id payload_buffer = [device->mtl_device newBufferWithBytes:payload_data.bytes ++ length:payload_data.length options:MTLResourceStorageModeShared]; ++ id tile_buffer = [device->mtl_device newBufferWithBytes:tile_data.bytes ++ length:tile_data.length options:MTLResourceStorageModeShared]; ++ id prefix_buffer = [device->mtl_device newBufferWithBytes:prefix_data.bytes ++ length:prefix_data.length options:MTLResourceStorageModeShared]; ++ id cell_buffer = [device->mtl_device newBufferWithBytes:job->cells ++ length:job->cell_count * sizeof(uint32_t) options:MTLResourceStorageModeShared]; ++ [payload_data release]; ++ [tile_data release]; ++ [prefix_data release]; ++ if (payload_buffer == nil || tile_buffer == nil || prefix_buffer == nil || cell_buffer == nil) { ++ [payload_buffer release]; ++ [tile_buffer release]; ++ [prefix_buffer release]; ++ [cell_buffer release]; ++ [encoder endEncoding]; ++ [temporary_buffers release]; ++ return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen staging allocation failed"); ++ } ++ [temporary_buffers addObject:payload_buffer]; ++ [temporary_buffers addObject:tile_buffer]; ++ [temporary_buffers addObject:prefix_buffer]; ++ [temporary_buffers addObject:cell_buffer]; ++ [payload_buffer release]; ++ [tile_buffer release]; ++ [prefix_buffer release]; ++ [cell_buffer release]; ++ ++ const struct ggml_metal_pipeline_with_params pipeline = ++ ggml_metal_library_compile_pipeline(device->library, ++ "kernel_cachegen_decode_f16", "kernel_cachegen_decode_f16", NULL); ++ if (pipeline.pipeline == NULL) { ++ [encoder endEncoding]; ++ [temporary_buffers release]; ++ return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen pipeline is unavailable"); ++ } ++ const struct ggml_metal_buffer_id dst = ggml_metal_buffer_get_id(dst_buffer, job->dst); ++ if (dst.metal == NULL) { ++ [encoder endEncoding]; ++ [temporary_buffers release]; ++ return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen destination buffer is unavailable"); ++ } ++ const struct ggml_metal_cachegen_params params = { ++ job->channels, ++ (uint32_t) job->tile_count, ++ job->token_stride, ++ job->channel_stride, ++ }; ++ [encoder setComputePipelineState:pipeline.pipeline->obj]; ++ [encoder setBuffer:temporary_buffers[temporary_buffers.count - 4] offset:0 atIndex:0]; ++ [encoder setBuffer:temporary_buffers[temporary_buffers.count - 3] offset:0 atIndex:1]; ++ [encoder setBuffer:temporary_buffers[temporary_buffers.count - 2] offset:0 atIndex:2]; ++ [encoder setBuffer:temporary_buffers[temporary_buffers.count - 1] offset:0 atIndex:3]; ++ [encoder setBuffer:dst.metal offset:dst.offs atIndex:4]; ++ [encoder setBytes:¶ms length:sizeof(params) atIndex:5]; ++ [encoder dispatchThreadgroups:MTLSizeMake((job->channels - 1) / 64 + 1, job->tile_count, 1) ++ threadsPerThreadgroup:MTLSizeMake(64, 1, 1)]; ++ } ++ [encoder endEncoding]; ++ [cmd_buf commit]; ++ [cmd_buf waitUntilCompleted]; ++ if (cmd_buf.status == MTLCommandBufferStatusError) { ++ const char * message = cmd_buf.error == nil ? "Metal CacheGen command failed" : cmd_buf.error.localizedDescription.UTF8String; ++ [temporary_buffers release]; ++ return ggml_metal_cachegen_error(error, error_capacity, message); ++ } ++ [temporary_buffers release]; ++ } ++ return true; ++} +diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp +index 3bd6abd06..6bec83efd 100644 +--- a/ggml/src/ggml-metal/ggml-metal.cpp ++++ b/ggml/src/ggml-metal/ggml-metal.cpp +@@ -910,6 +910,9 @@ static void * ggml_backend_metal_get_proc_address(ggml_backend_reg_t reg, const + if (strcmp(name, "ggml_backend_get_features") == 0) { + return (void *)ggml_backend_metal_get_features; + } ++ if (strcmp(name, "ggml_backend_cachegen_decode_f16") == 0) { ++ return (void *)ggml_metal_cachegen_decode_f16; ++ } + if (strcmp(name, "ggml_backend_metal_tuning_set_fa_vec_override") == 0) { + return (void *)ggml_backend_metal_tuning_set_fa_vec_override; + } +diff --git a/ggml/src/ggml-metal/kernels/cachegen.metal b/ggml/src/ggml-metal/kernels/cachegen.metal +new file mode 100644 +index 000000000..82684a71e +--- /dev/null ++++ b/ggml/src/ggml-metal/kernels/cachegen.metal +@@ -0,0 +1,98 @@ ++#include "common.h" ++ ++struct cachegen_tile { ++ uint payload_offset; ++ uint prefix_offset; ++ uint token_offset; ++ uint rows; ++}; ++ ++struct cachegen_params { ++ uint channels; ++ uint tile_count; ++ ulong token_stride; ++ ulong channel_stride; ++}; ++ ++static inline uint cachegen_read_bit(const device uchar * stream, uint stream_bytes, thread uint & bit) { ++ const uint byte_index = bit >> 3; ++ const uint value = byte_index < stream_bytes ? stream[byte_index] : 0; ++ const uint result = (value >> (7 - (bit & 7))) & 1; ++ ++bit; ++ return result; ++} ++ ++kernel void kernel_cachegen_decode_f16(const device uchar * payload [[buffer(0)]], ++ const device cachegen_tile * tiles [[buffer(1)]], ++ const device uint * prefixes [[buffer(2)]], ++ const device uint * cells [[buffer(3)]], ++ device uchar * dst [[buffer(4)]], ++ constant cachegen_params & params [[buffer(5)]], ++ uint2 gid [[thread_position_in_grid]]) { ++ const uint channel = gid.x; ++ const uint tile_index = gid.y; ++ if (channel >= params.channels || tile_index >= params.tile_count) { ++ return; ++ } ++ ++ const cachegen_tile tile = tiles[tile_index]; ++ const device uchar * segment = payload + tile.payload_offset; ++ const uint bins = segment[4]; ++ const device float * maxes = reinterpret_cast(segment + 16); ++ const device ushort * cdfs = reinterpret_cast(segment + 16 + tile.rows * sizeof(float)); ++ const device ushort * cdf = cdfs + channel * 33; ++ const uint prefix_index = tile.prefix_offset + channel; ++ const uint stream_start = prefixes[prefix_index]; ++ const uint stream_bytes = prefixes[prefix_index + 1] - stream_start; ++ const uint lengths_offset = 16 + tile.rows * sizeof(float) + params.channels * 33 * sizeof(ushort); ++ const uint streams_offset = lengths_offset + params.channels * sizeof(ushort); ++ const device uchar * stream = segment + streams_offset + stream_start; ++ ++ uint bit = 0; ++ uint value = 0; ++ for (uint i = 0; i < 32; ++i) { ++ value = (value << 1) | cachegen_read_bit(stream, stream_bytes, bit); ++ } ++ uint low = 0; ++ uint high = 0xffffffffu; ++ const float center = float(bins / 2 - 1); ++ for (uint row = 0; row < tile.rows; ++row) { ++ const ulong span = ulong(high) - ulong(low) + 1; ++ const ushort count = ushort((((ulong(value) - ulong(low) + 1) * 65536ul - 1) / span)); ++ uint symbol = 0; ++ while (symbol + 1 < 33 && cdf[symbol + 1] <= count) { ++ ++symbol; ++ } ++ if (symbol >= 32) { ++ return; ++ } ++ ++ const float normalized = (float(symbol) - center) / center; ++ const half decoded = half(normalized * maxes[row]); ++ const ulong destination = ++ ulong(cells[tile.token_offset + row]) * params.token_stride + ulong(channel) * params.channel_stride; ++ *reinterpret_cast(dst + destination) = decoded; ++ ++ if (row + 1 == tile.rows) { ++ break; ++ } ++ const ulong cdf_low = cdf[symbol]; ++ const ulong cdf_high = symbol == 31 ? 65536ul : cdf[symbol + 1]; ++ high = low - 1 + uint((span * cdf_high) >> 16); ++ low = low + uint((span * cdf_low) >> 16); ++ while (true) { ++ if (low >= 0x80000000u || high < 0x80000000u) { ++ low <<= 1; ++ high = (high << 1) | 1; ++ value = (value << 1) | cachegen_read_bit(stream, stream_bytes, bit); ++ } else if (low >= 0x40000000u && high < 0xc0000000u) { ++ low = (low << 1) & 0x7fffffffu; ++ high = (high << 1) | 0x80000001u; ++ value -= 0x40000000u; ++ value = (value << 1) | cachegen_read_bit(stream, stream_bytes, bit); ++ } else { ++ break; ++ } ++ } ++ } ++} +diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt +index 5a45037bb..4f57384ca 100644 +--- a/tests/CMakeLists.txt ++++ b/tests/CMakeLists.txt +@@ -279,6 +279,10 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) + set_tests_properties(test-save-load-state PROPERTIES FIXTURES_REQUIRED generate-models) + endif() + ++if (GGML_METAL) ++ llama_build_and_test(test-skippy-cachegen-metal.cpp) ++endif() ++ + llama_build_and_test(test-chat-peg-parser.cpp peg-parser/simple-tokenize.cpp) + llama_build_and_test(test-jinja.cpp) + llama_test(test-jinja NAME test-jinja-py ARGS -py LABEL python) +diff --git a/tests/test-skippy-cachegen-metal.cpp b/tests/test-skippy-cachegen-metal.cpp +new file mode 100644 +index 000000000..b64ba373c +--- /dev/null ++++ b/tests/test-skippy-cachegen-metal.cpp +@@ -0,0 +1,224 @@ ++#include "ggml-backend.h" ++#include "ggml-metal.h" ++#include "ggml.h" ++ ++#include ++#include ++#include ++ ++static const uint8_t encoded16[] = { ++ 0x4c, 0x43, 0x47, 0x31, 0x10, 0x00, 0x11, 0x00, 0x08, 0x00, 0x00, 0x00, 0x2b, 0x00, 0x00, 0x00, 0x00, 0xa0, 0x66, ++ 0x3e, 0x00, 0x80, 0xcd, 0x3e, 0x00, 0x60, 0x0e, 0x3f, 0x00, 0x40, 0x2e, 0x3f, 0x00, 0x40, 0x44, 0x3f, 0x00, 0x00, ++ 0x4f, 0x3f, 0x00, 0xe0, 0x4d, 0x3f, 0x00, 0x00, 0x41, 0x3f, 0x00, 0x20, 0x29, 0x3f, 0x00, 0xa0, 0x09, 0x3f, 0x00, ++ 0xe0, 0xcd, 0x3e, 0x00, 0x20, 0x75, 0x3e, 0x00, 0x80, 0x71, 0x3d, 0x00, 0xc0, 0x36, 0x3e, 0x00, 0x80, 0xb1, 0x3e, ++ 0x00, 0xa0, 0xfa, 0x3e, 0x00, 0x40, 0x1d, 0x3f, 0x00, 0x00, 0x1b, 0x1e, 0x29, 0x2d, 0x37, 0x3c, 0x38, 0x3c, 0x39, ++ 0x3c, 0x3a, 0x3c, 0x48, 0x4b, 0x49, 0x4b, 0x4a, 0x4b, 0x58, 0x5a, 0x59, 0x5a, 0x67, 0x69, 0x83, 0x87, 0xab, 0xb4, ++ 0xee, 0xff, 0xef, 0xff, 0xf0, 0xff, 0xf1, 0xff, 0xf2, 0xff, 0xf3, 0xff, 0xf4, 0xff, 0xf5, 0xff, 0xf6, 0xff, 0xf7, ++ 0xff, 0xf8, 0xff, 0xf9, 0xff, 0xfa, 0xff, 0xfb, 0xff, 0xfc, 0xff, 0xfd, 0xff, 0xfe, 0xff, 0xff, 0xff, 0x00, 0x00, ++ 0x1b, 0x1e, 0x29, 0x2d, 0x37, 0x3c, 0x38, 0x3c, 0x39, 0x3c, 0x3a, 0x3c, 0x3b, 0x3c, 0x49, 0x4b, 0x4a, 0x4b, 0x4b, ++ 0x4b, 0x59, 0x5a, 0x67, 0x69, 0x75, 0x78, 0xab, 0xb4, 0xee, 0xff, 0xef, 0xff, 0xf0, 0xff, 0xf1, 0xff, 0xf2, 0xff, ++ 0xf3, 0xff, 0xf4, 0xff, 0xf5, 0xff, 0xf6, 0xff, 0xf7, 0xff, 0xf8, 0xff, 0xf9, 0xff, 0xfa, 0xff, 0xfb, 0xff, 0xfc, ++ 0xff, 0xfd, 0xff, 0xfe, 0xff, 0xff, 0xff, 0x00, 0x00, 0x28, 0x2d, 0x36, 0x3c, 0x37, 0x3c, 0x38, 0x3c, 0x39, 0x3c, ++ 0x3a, 0x3c, 0x3b, 0x3c, 0x3c, 0x3c, 0x4a, 0x4b, 0x4b, 0x4b, 0x4c, 0x4b, 0x5a, 0x5a, 0x83, 0x87, 0xab, 0xb4, 0xee, ++ 0xff, 0xef, 0xff, 0xf0, 0xff, 0xf1, 0xff, 0xf2, 0xff, 0xf3, 0xff, 0xf4, 0xff, 0xf5, 0xff, 0xf6, 0xff, 0xf7, 0xff, ++ 0xf8, 0xff, 0xf9, 0xff, 0xfa, 0xff, 0xfb, 0xff, 0xfc, 0xff, 0xfd, 0xff, 0xfe, 0xff, 0xff, 0xff, 0x00, 0x00, 0x28, ++ 0x2d, 0x36, 0x3c, 0x37, 0x3c, 0x38, 0x3c, 0x39, 0x3c, 0x3a, 0x3c, 0x3b, 0x3c, 0x3c, 0x3c, 0x3d, 0x3c, 0x4b, 0x4b, ++ 0x4c, 0x4b, 0x67, 0x69, 0x75, 0x78, 0xab, 0xb4, 0xee, 0xff, 0xef, 0xff, 0xf0, 0xff, 0xf1, 0xff, 0xf2, 0xff, 0xf3, ++ 0xff, 0xf4, 0xff, 0xf5, 0xff, 0xf6, 0xff, 0xf7, 0xff, 0xf8, 0xff, 0xf9, 0xff, 0xfa, 0xff, 0xfb, 0xff, 0xfc, 0xff, ++ 0xfd, 0xff, 0xfe, 0xff, 0xff, 0xff, 0x00, 0x00, 0x28, 0x2d, 0x36, 0x3c, 0x37, 0x3c, 0x38, 0x3c, 0x39, 0x3c, 0x3a, ++ 0x3c, 0x3b, 0x3c, 0x3c, 0x3c, 0x3d, 0x3c, 0x3e, 0x3c, 0x59, 0x5a, 0x5a, 0x5a, 0x68, 0x69, 0x9e, 0xa5, 0xee, 0xff, ++ 0xef, 0xff, 0xf0, 0xff, 0xf1, 0xff, 0xf2, 0xff, 0xf3, 0xff, 0xf4, 0xff, 0xf5, 0xff, 0xf6, 0xff, 0xf7, 0xff, 0xf8, ++ 0xff, 0xf9, 0xff, 0xfa, 0xff, 0xfb, 0xff, 0xfc, 0xff, 0xfd, 0xff, 0xfe, 0xff, 0xff, 0xff, 0x00, 0x00, 0x28, 0x2d, ++ 0x36, 0x3c, 0x37, 0x3c, 0x38, 0x3c, 0x39, 0x3c, 0x3a, 0x3c, 0x3b, 0x3c, 0x3c, 0x3c, 0x3d, 0x3c, 0x4b, 0x4b, 0x4c, ++ 0x4b, 0x4d, 0x4b, 0x5b, 0x5a, 0x91, 0x96, 0xee, 0xff, 0xef, 0xff, 0xf0, 0xff, 0xf1, 0xff, 0xf2, 0xff, 0xf3, 0xff, ++ 0xf4, 0xff, 0xf5, 0xff, 0xf6, 0xff, 0xf7, 0xff, 0xf8, 0xff, 0xf9, 0xff, 0xfa, 0xff, 0xfb, 0xff, 0xfc, 0xff, 0xfd, ++ 0xff, 0xfe, 0xff, 0xff, 0xff, 0x00, 0x00, 0x35, 0x3c, 0x36, 0x3c, 0x37, 0x3c, 0x38, 0x3c, 0x39, 0x3c, 0x3a, 0x3c, ++ 0x3b, 0x3c, 0x3c, 0x3c, 0x4a, 0x4b, 0x4b, 0x4b, 0x4c, 0x4b, 0x4d, 0x4b, 0x4e, 0x4b, 0x76, 0x78, 0xee, 0xff, 0xef, ++ 0xff, 0xf0, 0xff, 0xf1, 0xff, 0xf2, 0xff, 0xf3, 0xff, 0xf4, 0xff, 0xf5, 0xff, 0xf6, 0xff, 0xf7, 0xff, 0xf8, 0xff, ++ 0xf9, 0xff, 0xfa, 0xff, 0xfb, 0xff, 0xfc, 0xff, 0xfd, 0xff, 0xfe, 0xff, 0xff, 0xff, 0x00, 0x00, 0x35, 0x3c, 0x36, ++ 0x3c, 0x37, 0x3c, 0x38, 0x3c, 0x39, 0x3c, 0x3a, 0x3c, 0x3b, 0x3c, 0x3c, 0x3c, 0x4a, 0x4b, 0x4b, 0x4b, 0x4c, 0x4b, ++ 0x4d, 0x4b, 0x4e, 0x4b, 0x69, 0x69, 0xee, 0xff, 0xef, 0xff, 0xf0, 0xff, 0xf1, 0xff, 0xf2, 0xff, 0xf3, 0xff, 0xf4, ++ 0xff, 0xf5, 0xff, 0xf6, 0xff, 0xf7, 0xff, 0xf8, 0xff, 0xf9, 0xff, 0xfa, 0xff, 0xfb, 0xff, 0xfc, 0xff, 0xfd, 0xff, ++ 0xfe, 0xff, 0xff, 0xff, 0x07, 0x00, 0x06, 0x00, 0x06, 0x00, 0x06, 0x00, 0x05, 0x00, 0x05, 0x00, 0x04, 0x00, 0x04, ++ 0x00, 0x40, 0xfd, 0xb1, 0x42, 0x5b, 0x3d, 0x00, 0x40, 0xfe, 0x79, 0x1b, 0x7a, 0x40, 0x41, 0x0c, 0x38, 0x22, 0x9c, ++ 0xc8, 0x41, 0x75, 0xb5, 0x6f, 0x07, 0xc0, 0x47, 0xd4, 0xca, 0xf4, 0x9b, 0x52, 0x5a, 0x84, 0x9c, 0xb4, 0x60, 0x87, ++ 0xf8, 0x88, 0xfe, 0xfc, 0x5d, 0xf0, ++}; ++ ++static const uint8_t expected16[] = { ++ 0x1e, 0xa8, 0x00, 0x00, 0x1e, 0x28, 0x1e, 0x2c, 0x2d, 0x2e, 0x26, 0x31, 0x2d, 0x32, 0x35, 0x33, 0x57, 0x2f, 0x81, ++ 0x31, 0x57, 0x33, 0x57, 0x33, 0x96, 0x34, 0x81, 0x35, 0x81, 0x35, 0x6c, 0x36, 0x16, 0x35, 0x16, 0x35, 0x5b, 0x36, ++ 0x5b, 0x36, 0xa1, 0x37, 0xa1, 0x37, 0x73, 0x38, 0x73, 0x38, 0xc7, 0x37, 0xc7, 0x37, 0xc7, 0x37, 0xab, 0x38, 0xab, ++ 0x38, 0xab, 0x38, 0x72, 0x39, 0x72, 0x39, 0x61, 0x38, 0x42, 0x39, 0x42, 0x39, 0x42, 0x39, 0x42, 0x39, 0x22, 0x3a, ++ 0x22, 0x3a, 0x22, 0x3a, 0x8b, 0x39, 0x8b, 0x39, 0x8b, 0x39, 0x8b, 0x39, 0x78, 0x3a, 0x78, 0x3a, 0x78, 0x3a, 0x78, ++ 0x3a, 0x84, 0x39, 0x84, 0x39, 0x84, 0x39, 0x6f, 0x3a, 0x6f, 0x3a, 0x6f, 0x3a, 0x6f, 0x3a, 0x6f, 0x3a, 0x2b, 0x39, ++ 0x08, 0x3a, 0x08, 0x3a, 0x08, 0x3a, 0x08, 0x3a, 0x08, 0x3a, 0x08, 0x3a, 0x08, 0x3a, 0x49, 0x39, 0x49, 0x39, 0x49, ++ 0x39, 0x49, 0x39, 0x49, 0x39, 0x49, 0x39, 0x49, 0x39, 0x49, 0x39, 0x4d, 0x38, 0x4d, 0x38, 0x4d, 0x38, 0x4d, 0x38, ++ 0x4d, 0x38, 0x4d, 0x38, 0x4d, 0x38, 0x4d, 0x38, 0x6f, 0x36, 0x6f, 0x36, 0x6f, 0x36, 0x6f, 0x36, 0x6f, 0x36, 0x6f, ++ 0x36, 0x6f, 0x36, 0x84, 0x35, 0xa9, 0x33, 0xa9, 0x33, 0xa9, 0x33, 0x91, 0x32, 0x91, 0x32, 0x91, 0x32, 0x91, 0x32, ++ 0x91, 0x32, 0x8c, 0x2b, 0x78, 0x2a, 0x64, 0x29, 0x50, 0x28, 0x78, 0x26, 0x50, 0x24, 0x50, 0x20, 0x50, 0x20, 0x14, ++ 0xb0, 0x14, 0xb0, 0xe5, 0xb0, 0xe5, 0xb0, 0xe5, 0xb0, 0xe5, 0xb0, 0xb6, 0xb1, 0xb6, 0xb1, 0xc1, 0xb4, 0xc1, 0xb4, ++ 0x8c, 0xb5, 0x8c, 0xb5, 0x8c, 0xb5, 0x8c, 0xb5, 0x8c, 0xb5, 0x8c, 0xb5, 0xd5, 0xb7, 0xd5, 0xb7, 0xd5, 0xb7, 0xd5, ++ 0xb7, 0xd5, 0xb7, 0xd5, 0xb7, 0xd5, 0xb7, 0xd5, 0xb7, 0xea, 0xb8, 0xea, 0xb8, 0xea, 0xb8, 0xea, 0xb8, 0xea, 0xb8, ++ 0xea, 0xb8, 0xea, 0xb8, 0xea, 0xb8, ++}; ++ ++static const uint8_t encoded32[] = { ++ 0x4c, 0x43, 0x47, 0x31, 0x20, 0x00, 0x11, 0x00, 0x08, 0x00, 0x00, 0x00, 0x35, 0x00, 0x00, 0x00, ++ 0x00, 0xa0, 0x66, 0x3e, 0x00, 0x80, 0xcd, 0x3e, 0x00, 0x60, 0x0e, 0x3f, 0x00, 0x40, 0x2e, 0x3f, ++ 0x00, 0x40, 0x44, 0x3f, 0x00, 0x00, 0x4f, 0x3f, 0x00, 0xe0, 0x4d, 0x3f, 0x00, 0x00, 0x41, 0x3f, ++ 0x00, 0x20, 0x29, 0x3f, 0x00, 0xa0, 0x09, 0x3f, 0x00, 0xe0, 0xcd, 0x3e, 0x00, 0x20, 0x75, 0x3e, ++ 0x00, 0x80, 0x71, 0x3d, 0x00, 0xc0, 0x36, 0x3e, 0x00, 0x80, 0xb1, 0x3e, 0x00, 0xa0, 0xfa, 0x3e, ++ 0x00, 0x40, 0x1d, 0x3f, 0x00, 0x00, 0x1b, 0x1e, 0x1c, 0x1e, 0x2a, 0x2d, 0x2b, 0x2d, 0x39, 0x3c, ++ 0x3a, 0x3c, 0x3b, 0x3c, 0x3c, 0x3c, 0x3d, 0x3c, 0x3e, 0x3c, 0x3f, 0x3c, 0x40, 0x3c, 0x4e, 0x4b, ++ 0x4f, 0x4b, 0x50, 0x4b, 0x51, 0x4b, 0x52, 0x4b, 0x53, 0x4b, 0x54, 0x4b, 0x55, 0x4b, 0x63, 0x5a, ++ 0x64, 0x5a, 0x65, 0x5a, 0x73, 0x69, 0x74, 0x69, 0x82, 0x78, 0x91, 0x87, 0x9f, 0x96, 0xad, 0xa5, ++ 0xc8, 0xc3, 0xfe, 0xff, 0xff, 0xff, 0x00, 0x00, 0x1b, 0x1e, 0x29, 0x2d, 0x2a, 0x2d, 0x2b, 0x2d, ++ 0x39, 0x3c, 0x3a, 0x3c, 0x3b, 0x3c, 0x3c, 0x3c, 0x3d, 0x3c, 0x3e, 0x3c, 0x3f, 0x3c, 0x40, 0x3c, ++ 0x41, 0x3c, 0x42, 0x3c, 0x50, 0x4b, 0x51, 0x4b, 0x52, 0x4b, 0x53, 0x4b, 0x54, 0x4b, 0x55, 0x4b, ++ 0x56, 0x4b, 0x57, 0x4b, 0x65, 0x5a, 0x66, 0x5a, 0x74, 0x69, 0x75, 0x69, 0x83, 0x78, 0x92, 0x87, ++ 0xba, 0xb4, 0xc8, 0xc3, 0xfe, 0xff, 0xff, 0xff, 0x00, 0x00, 0x1b, 0x1e, 0x29, 0x2d, 0x2a, 0x2d, ++ 0x38, 0x3c, 0x39, 0x3c, 0x3a, 0x3c, 0x3b, 0x3c, 0x3c, 0x3c, 0x3d, 0x3c, 0x3e, 0x3c, 0x3f, 0x3c, ++ 0x40, 0x3c, 0x41, 0x3c, 0x42, 0x3c, 0x43, 0x3c, 0x44, 0x3c, 0x45, 0x3c, 0x53, 0x4b, 0x54, 0x4b, ++ 0x55, 0x4b, 0x56, 0x4b, 0x57, 0x4b, 0x58, 0x4b, 0x66, 0x5a, 0x67, 0x5a, 0x75, 0x69, 0x83, 0x78, ++ 0x92, 0x87, 0xad, 0xa5, 0xd5, 0xd2, 0xfe, 0xff, 0xff, 0xff, 0x00, 0x00, 0x1b, 0x1e, 0x29, 0x2d, ++ 0x37, 0x3c, 0x38, 0x3c, 0x39, 0x3c, 0x3a, 0x3c, 0x3b, 0x3c, 0x3c, 0x3c, 0x3d, 0x3c, 0x3e, 0x3c, ++ 0x3f, 0x3c, 0x40, 0x3c, 0x41, 0x3c, 0x42, 0x3c, 0x43, 0x3c, 0x44, 0x3c, 0x45, 0x3c, 0x46, 0x3c, ++ 0x47, 0x3c, 0x48, 0x3c, 0x56, 0x4b, 0x57, 0x4b, 0x58, 0x4b, 0x59, 0x4b, 0x74, 0x69, 0x75, 0x69, ++ 0x83, 0x78, 0x92, 0x87, 0xa0, 0x96, 0xd5, 0xd2, 0xfe, 0xff, 0xff, 0xff, 0x00, 0x00, 0x1b, 0x1e, ++ 0x29, 0x2d, 0x37, 0x3c, 0x38, 0x3c, 0x39, 0x3c, 0x3a, 0x3c, 0x3b, 0x3c, 0x3c, 0x3c, 0x3d, 0x3c, ++ 0x3e, 0x3c, 0x3f, 0x3c, 0x40, 0x3c, 0x41, 0x3c, 0x42, 0x3c, 0x43, 0x3c, 0x44, 0x3c, 0x45, 0x3c, ++ 0x46, 0x3c, 0x47, 0x3c, 0x48, 0x3c, 0x49, 0x3c, 0x4a, 0x3c, 0x65, 0x5a, 0x66, 0x5a, 0x67, 0x5a, ++ 0x68, 0x5a, 0x76, 0x69, 0x84, 0x78, 0xa0, 0x96, 0xd5, 0xd2, 0xfe, 0xff, 0xff, 0xff, 0x00, 0x00, ++ 0x28, 0x2d, 0x36, 0x3c, 0x37, 0x3c, 0x38, 0x3c, 0x39, 0x3c, 0x3a, 0x3c, 0x3b, 0x3c, 0x3c, 0x3c, ++ 0x3d, 0x3c, 0x3e, 0x3c, 0x3f, 0x3c, 0x40, 0x3c, 0x41, 0x3c, 0x42, 0x3c, 0x43, 0x3c, 0x44, 0x3c, ++ 0x45, 0x3c, 0x46, 0x3c, 0x47, 0x3c, 0x48, 0x3c, 0x56, 0x4b, 0x57, 0x4b, 0x58, 0x4b, 0x59, 0x4b, ++ 0x5a, 0x4b, 0x68, 0x5a, 0x69, 0x5a, 0x77, 0x69, 0x93, 0x87, 0xc8, 0xc3, 0xfe, 0xff, 0xff, 0xff, ++ 0x00, 0x00, 0x28, 0x2d, 0x36, 0x3c, 0x37, 0x3c, 0x38, 0x3c, 0x39, 0x3c, 0x3a, 0x3c, 0x3b, 0x3c, ++ 0x3c, 0x3c, 0x3d, 0x3c, 0x3e, 0x3c, 0x3f, 0x3c, 0x40, 0x3c, 0x41, 0x3c, 0x42, 0x3c, 0x43, 0x3c, ++ 0x44, 0x3c, 0x45, 0x3c, 0x46, 0x3c, 0x54, 0x4b, 0x55, 0x4b, 0x56, 0x4b, 0x57, 0x4b, 0x58, 0x4b, ++ 0x59, 0x4b, 0x5a, 0x4b, 0x5b, 0x4b, 0x5c, 0x4b, 0x77, 0x69, 0x78, 0x69, 0xae, 0xa5, 0xfe, 0xff, ++ 0xff, 0xff, 0x00, 0x00, 0x35, 0x3c, 0x36, 0x3c, 0x37, 0x3c, 0x38, 0x3c, 0x39, 0x3c, 0x3a, 0x3c, ++ 0x3b, 0x3c, 0x3c, 0x3c, 0x3d, 0x3c, 0x3e, 0x3c, 0x3f, 0x3c, 0x40, 0x3c, 0x41, 0x3c, 0x42, 0x3c, ++ 0x43, 0x3c, 0x44, 0x3c, 0x52, 0x4b, 0x53, 0x4b, 0x54, 0x4b, 0x55, 0x4b, 0x56, 0x4b, 0x57, 0x4b, ++ 0x58, 0x4b, 0x59, 0x4b, 0x5a, 0x4b, 0x5b, 0x4b, 0x5c, 0x4b, 0x6a, 0x5a, 0x6b, 0x5a, 0x79, 0x69, ++ 0xfe, 0xff, 0xff, 0xff, 0x08, 0x00, 0x07, 0x00, 0x08, 0x00, 0x07, 0x00, 0x07, 0x00, 0x06, 0x00, ++ 0x06, 0x00, 0x04, 0x00, 0x41, 0x04, 0x03, 0xa5, 0xc8, 0x5a, 0x02, 0xc0, 0x41, 0x06, 0x23, 0xa4, ++ 0x7a, 0x6e, 0xeb, 0x41, 0x0a, 0x0f, 0x21, 0xb2, 0xd0, 0xef, 0x00, 0x41, 0x7e, 0x7c, 0xc9, 0x9a, ++ 0xa4, 0x08, 0x47, 0xb5, 0x66, 0xc1, 0x5d, 0x23, 0x70, 0x51, 0x1b, 0x08, 0x21, 0x46, 0xbb, 0x5b, ++ 0xc1, 0xcd, 0xf2, 0x3a, 0x80, 0xff, 0x2f, 0xba, 0xac, ++}; ++ ++static const uint8_t expected32[] = { ++ 0xc4, 0xa9, 0xb0, 0xa3, 0xb0, 0x27, 0xce, 0x2c, 0xba, 0x2e, 0xce, 0x30, 0xc4, 0x31, 0x35, 0x33, ++ 0x48, 0x30, 0xfe, 0x31, 0xda, 0x32, 0xb5, 0x33, 0xb6, 0x34, 0x23, 0x35, 0xfe, 0x35, 0x6c, 0x36, ++ 0xbf, 0x34, 0x57, 0x35, 0xef, 0x35, 0x87, 0x36, 0x1e, 0x37, 0xb6, 0x37, 0x27, 0x38, 0x73, 0x38, ++ 0x43, 0x37, 0xfd, 0x37, 0x5b, 0x38, 0x5b, 0x38, 0xb8, 0x38, 0x15, 0x39, 0x15, 0x39, 0x72, 0x39, ++ 0x7f, 0x38, 0xe8, 0x38, 0x51, 0x39, 0x51, 0x39, 0xb9, 0x39, 0xb9, 0x39, 0x22, 0x3a, 0x22, 0x3a, ++ 0x2d, 0x39, 0x9b, 0x39, 0x9b, 0x39, 0x0a, 0x3a, 0x0a, 0x3a, 0x0a, 0x3a, 0x78, 0x3a, 0x78, 0x3a, ++ 0x93, 0x39, 0x93, 0x39, 0x01, 0x3a, 0x01, 0x3a, 0x01, 0x3a, 0x6f, 0x3a, 0x6f, 0x3a, 0x6f, 0x3a, ++ 0xa1, 0x39, 0xa1, 0x39, 0xa1, 0x39, 0xa1, 0x39, 0x08, 0x3a, 0x08, 0x3a, 0x08, 0x3a, 0x08, 0x3a, ++ 0xef, 0x38, 0x49, 0x39, 0x49, 0x39, 0x49, 0x39, 0x49, 0x39, 0x49, 0x39, 0x49, 0x39, 0x49, 0x39, ++ 0x4d, 0x38, 0x4d, 0x38, 0x4d, 0x38, 0x4d, 0x38, 0x4d, 0x38, 0x4d, 0x38, 0x4d, 0x38, 0x4d, 0x38, ++ 0x6f, 0x36, 0x6f, 0x36, 0x6f, 0x36, 0x6f, 0x36, 0x01, 0x36, 0x01, 0x36, 0x01, 0x36, 0x01, 0x36, ++ 0xa9, 0x33, 0xa9, 0x33, 0x26, 0x33, 0x26, 0x33, 0xa4, 0x32, 0xa4, 0x32, 0x21, 0x32, 0x21, 0x32, ++ 0x8c, 0x2b, 0x8a, 0x2a, 0x89, 0x29, 0x87, 0x28, 0x0b, 0x27, 0x08, 0x25, 0x0a, 0x22, 0x06, 0x1c, ++ 0x30, 0xb0, 0x30, 0xb0, 0x92, 0xb0, 0xf3, 0xb0, 0xf3, 0xb0, 0x55, 0xb1, 0x55, 0xb1, 0xb6, 0xb1, ++ 0xcf, 0xb4, 0x2d, 0xb5, 0x2d, 0xb5, 0x2d, 0xb5, 0x2d, 0xb5, 0x8c, 0xb5, 0x8c, 0xb5, 0x8c, 0xb5, ++ 0xd5, 0xb7, 0xd5, 0xb7, 0xd5, 0xb7, 0xd5, 0xb7, 0xd5, 0xb7, 0xd5, 0xb7, 0xd5, 0xb7, 0xd5, 0xb7, ++ 0xea, 0xb8, 0xea, 0xb8, 0xea, 0xb8, 0xea, 0xb8, 0xea, 0xb8, 0xea, 0xb8, 0xea, 0xb8, 0xea, 0xb8, ++}; ++ ++int main() { ++ constexpr uint32_t rows = 17; ++ constexpr uint32_t channels = 8; ++ constexpr uint32_t capacity = rows + 5; ++ ++ ggml_backend_t backend = ggml_backend_metal_init(); ++ if (backend == nullptr) { ++ return 1; ++ } ++ ggml_init_params params = { ggml_tensor_overhead() * 5, nullptr, true }; ++ ggml_context * ctx = ggml_init(params); ++ if (ctx == nullptr) { ++ ggml_backend_free(backend); ++ return 2; ++ } ++ ggml_tensor * row_major16 = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, channels, capacity); ++ ggml_tensor * transposed16 = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, capacity, channels); ++ ggml_tensor * row_major32 = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, channels, capacity); ++ ggml_tensor * transposed32 = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, capacity, channels); ++ ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); ++ if (buffer == nullptr) { ++ ggml_free(ctx); ++ ggml_backend_free(backend); ++ return 3; ++ } ++ const auto finish = [&](int status) { ++ ggml_backend_buffer_free(buffer); ++ ggml_free(ctx); ++ ggml_backend_free(backend); ++ return status; ++ }; ++ ++ std::vector cells(rows); ++ for (uint32_t row = 0; row < rows; ++row) { ++ cells[row] = capacity - 1 - row; ++ } ++ const ggml_backend_cachegen_tile tile16 = { encoded16, sizeof(encoded16), 0, rows }; ++ const ggml_backend_cachegen_tile tile32 = { encoded32, sizeof(encoded32), 0, rows }; ++ const ggml_backend_cachegen_job jobs[] = { ++ { row_major16, &tile16, 1, cells.data(), cells.size(), channels, channels * 2, 2 }, ++ { transposed16, &tile16, 1, cells.data(), cells.size(), channels, 2, capacity * 2 }, ++ { row_major32, &tile32, 1, cells.data(), cells.size(), channels, channels * 2, 2 }, ++ { transposed32, &tile32, 1, cells.data(), cells.size(), channels, 2, capacity * 2 }, ++ }; ++ ggml_backend_dev_t device = ggml_backend_get_device(backend); ++ auto decode = reinterpret_cast( ++ ggml_backend_reg_get_proc_address( ++ ggml_backend_dev_backend_reg(device), "ggml_backend_cachegen_decode_f16")); ++ if (decode == nullptr) { ++ return finish(4); ++ } ++ char error[256] = {}; ++ if (!decode(jobs, 4, error, sizeof(error))) { ++ return finish(5); ++ } ++ ggml_backend_cachegen_job invalid_job = jobs[0]; ++ invalid_job.token_stride = UINT64_MAX; ++ if (decode(&invalid_job, 1, error, sizeof(error))) { ++ return finish(8); ++ } ++ ++ const auto check_fixture = [&](ggml_tensor * row_major, ggml_tensor * transposed, ++ const uint8_t * expected, int status) { ++ std::vector row_bytes(ggml_nbytes(row_major)); ++ std::vector transposed_bytes(ggml_nbytes(transposed)); ++ ggml_backend_tensor_get(row_major, row_bytes.data(), 0, row_bytes.size()); ++ ggml_backend_tensor_get(transposed, transposed_bytes.data(), 0, transposed_bytes.size()); ++ for (uint32_t row = 0; row < rows; ++row) { ++ if (std::memcmp( ++ row_bytes.data() + static_cast(cells[row]) * channels * 2, ++ expected + static_cast(row) * channels * 2, ++ channels * 2) != 0) { ++ return status; ++ } ++ for (uint32_t channel = 0; channel < channels; ++channel) { ++ if (std::memcmp( ++ transposed_bytes.data() + (cells[row] + static_cast(channel) * capacity) * 2, ++ expected + (static_cast(row) * channels + channel) * 2, ++ 2) != 0) { ++ return status; ++ } ++ } ++ } ++ return 0; ++ }; ++ if (const int status = check_fixture(row_major16, transposed16, expected16, 6)) { ++ return finish(status); ++ } ++ if (const int status = check_fixture(row_major32, transposed32, expected32, 7)) { ++ return finish(status); ++ } ++ return finish(0); ++} +-- +2.54.0 (Apple Git-157) + diff --git a/third_party/llama.cpp/patches/0038-ggml-cuda-decode-CacheGen-pages-into-resident-KV.patch b/third_party/llama.cpp/patches/0038-ggml-cuda-decode-CacheGen-pages-into-resident-KV.patch new file mode 100644 index 0000000000..111a06bf73 --- /dev/null +++ b/third_party/llama.cpp/patches/0038-ggml-cuda-decode-CacheGen-pages-into-resident-KV.patch @@ -0,0 +1,481 @@ +From 3391262f8620a3d1a7ddb6a174e75e0c6129c804 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 11 Sep 2026 16:58:10 +1000 +Subject: [PATCH] ggml-cuda: decode CacheGen pages into resident KV + +Share the F16 decoder across CUDA and HIP, batch launches before synchronization, and write directly into the allocated Skippy cell layout. + +Assisted-by: scama +--- + ggml/src/ggml-cuda/cachegen.cu | 111 +++++++++++++ + ggml/src/ggml-cuda/cachegen.cuh | 26 +++ + ggml/src/ggml-cuda/ggml-cuda.cu | 281 ++++++++++++++++++++++++++++++++ + 3 files changed, 418 insertions(+) + create mode 100644 ggml/src/ggml-cuda/cachegen.cu + create mode 100644 ggml/src/ggml-cuda/cachegen.cuh + +diff --git a/ggml/src/ggml-cuda/cachegen.cu b/ggml/src/ggml-cuda/cachegen.cu +new file mode 100644 +index 000000000..5d91aa0ca +--- /dev/null ++++ b/ggml/src/ggml-cuda/cachegen.cu +@@ -0,0 +1,111 @@ ++#include "cachegen.cuh" ++ ++#if !defined(GGML_USE_MUSA) ++ ++static __device__ __forceinline__ uint32_t cachegen_read_bit( ++ const uint8_t * stream, uint32_t stream_bytes, uint32_t & bit) { ++ const uint32_t byte_index = bit >> 3; ++ const uint32_t value = byte_index < stream_bytes ? stream[byte_index] : 0; ++ const uint32_t result = (value >> (7 - (bit & 7))) & 1; ++ ++bit; ++ return result; ++} ++ ++static __global__ void cachegen_decode_f16( ++ const uint8_t * payload, ++ const ggml_cuda_cachegen_tile * tiles, ++ const uint32_t * prefixes, ++ const uint32_t * cells, ++ uint8_t * dst, ++ uint32_t channels, ++ uint32_t tile_count, ++ uint64_t token_stride, ++ uint64_t channel_stride) { ++ const uint32_t channel = blockIdx.x * blockDim.x + threadIdx.x; ++ const uint32_t tile_index = blockIdx.y; ++ if (channel >= channels || tile_index >= tile_count) { ++ return; ++ } ++ ++ const ggml_cuda_cachegen_tile tile = tiles[tile_index]; ++ const uint8_t * segment = payload + tile.payload_offset; ++ const uint32_t bins = segment[4]; ++ const float * maxes = reinterpret_cast(segment + 16); ++ const uint16_t * cdfs = reinterpret_cast(segment + 16 + tile.rows * sizeof(float)); ++ const uint16_t * cdf = cdfs + channel * 33; ++ const uint32_t prefix_index = tile.prefix_offset + channel; ++ const uint32_t stream_start = prefixes[prefix_index]; ++ const uint32_t stream_bytes = prefixes[prefix_index + 1] - stream_start; ++ const uint32_t lengths_offset = 16 + tile.rows * sizeof(float) + channels * 33 * sizeof(uint16_t); ++ const uint32_t streams_offset = lengths_offset + channels * sizeof(uint16_t); ++ const uint8_t * stream = segment + streams_offset + stream_start; ++ ++ uint32_t bit = 0; ++ uint32_t value = 0; ++ for (uint32_t i = 0; i < 32; ++i) { ++ value = (value << 1) | cachegen_read_bit(stream, stream_bytes, bit); ++ } ++ uint32_t low = 0; ++ uint32_t high = 0xffffffffu; ++ const float center = float(bins / 2 - 1); ++ for (uint32_t row = 0; row < tile.rows; ++row) { ++ const uint64_t span = uint64_t(high) - uint64_t(low) + 1; ++ const uint16_t count = uint16_t((((uint64_t(value) - uint64_t(low) + 1) * 65536ull - 1) / span)); ++ uint32_t symbol = 0; ++ while (symbol + 1 < 33 && cdf[symbol + 1] <= count) { ++ ++symbol; ++ } ++ if (symbol >= 32) { ++ return; ++ } ++ ++ const float normalized = (float(symbol) - center) / center; ++ const half decoded = __float2half(normalized * maxes[row]); ++ const uint64_t destination = ++ uint64_t(cells[tile.token_offset + row]) * token_stride + uint64_t(channel) * channel_stride; ++ *reinterpret_cast(dst + destination) = decoded; ++ ++ if (row + 1 == tile.rows) { ++ break; ++ } ++ const uint64_t cdf_low = cdf[symbol]; ++ const uint64_t cdf_high = symbol == 31 ? 65536ull : cdf[symbol + 1]; ++ high = low - 1 + uint32_t((span * cdf_high) >> 16); ++ low = low + uint32_t((span * cdf_low) >> 16); ++ while (true) { ++ if (low >= 0x80000000u || high < 0x80000000u) { ++ low <<= 1; ++ high = (high << 1) | 1; ++ value = (value << 1) | cachegen_read_bit(stream, stream_bytes, bit); ++ } else if (low >= 0x40000000u && high < 0xc0000000u) { ++ low = (low << 1) & 0x7fffffffu; ++ high = (high << 1) | 0x80000001u; ++ value -= 0x40000000u; ++ value = (value << 1) | cachegen_read_bit(stream, stream_bytes, bit); ++ } else { ++ break; ++ } ++ } ++ } ++} ++ ++cudaError_t ggml_cuda_cachegen_decode_f16_launch( ++ const uint8_t * payload, ++ const ggml_cuda_cachegen_tile * tiles, ++ const uint32_t * prefixes, ++ const uint32_t * cells, ++ uint8_t * dst, ++ uint32_t channels, ++ uint32_t tile_count, ++ uint64_t token_stride, ++ uint64_t channel_stride, ++ cudaStream_t stream) { ++ constexpr uint32_t block_size = 64; ++ const dim3 block(block_size, 1, 1); ++ const dim3 grid((channels + block_size - 1) / block_size, tile_count, 1); ++ cachegen_decode_f16<<>>( ++ payload, tiles, prefixes, cells, dst, channels, tile_count, token_stride, channel_stride); ++ return cudaGetLastError(); ++} ++ ++#endif // !defined(GGML_USE_MUSA) +diff --git a/ggml/src/ggml-cuda/cachegen.cuh b/ggml/src/ggml-cuda/cachegen.cuh +new file mode 100644 +index 000000000..d6d216675 +--- /dev/null ++++ b/ggml/src/ggml-cuda/cachegen.cuh +@@ -0,0 +1,26 @@ ++#pragma once ++ ++#include "common.cuh" ++ ++#if !defined(GGML_USE_MUSA) ++ ++struct ggml_cuda_cachegen_tile { ++ uint32_t payload_offset; ++ uint32_t prefix_offset; ++ uint32_t token_offset; ++ uint32_t rows; ++}; ++ ++cudaError_t ggml_cuda_cachegen_decode_f16_launch( ++ const uint8_t * payload, ++ const ggml_cuda_cachegen_tile * tiles, ++ const uint32_t * prefixes, ++ const uint32_t * cells, ++ uint8_t * dst, ++ uint32_t channels, ++ uint32_t tile_count, ++ uint64_t token_stride, ++ uint64_t channel_stride, ++ cudaStream_t stream); ++ ++#endif // !defined(GGML_USE_MUSA) +diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu +index 38bd4c9a0..e955c4394 100644 +--- a/ggml/src/ggml-cuda/ggml-cuda.cu ++++ b/ggml/src/ggml-cuda/ggml-cuda.cu +@@ -10,6 +10,7 @@ + #include "ggml-cuda/argmax.cuh" + #include "ggml-cuda/argsort.cuh" + #include "ggml-cuda/binbcast.cuh" ++#include "ggml-cuda/cachegen.cuh" + #include "ggml-cuda/clamp.cuh" + #include "ggml-cuda/col2im-1d.cuh" + #include "ggml-cuda/concat.cuh" +@@ -85,6 +86,7 @@ + #include + #include + #include ++#include + #include + #include + #include +@@ -5667,6 +5669,280 @@ static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t + GGML_UNUSED(reg); + } + ++#if !defined(GGML_USE_MUSA) ++ ++struct ggml_cuda_cachegen_staged_job { ++ int physical_device = -1; ++ uint8_t * dst = nullptr; ++ uint32_t channels = 0; ++ uint64_t token_stride = 0; ++ uint64_t channel_stride = 0; ++ std::vector payload; ++ std::vector tiles; ++ std::vector prefixes; ++ std::vector cells; ++ uint8_t * payload_device = nullptr; ++ ggml_cuda_cachegen_tile * tiles_device = nullptr; ++ uint32_t * prefixes_device = nullptr; ++ uint32_t * cells_device = nullptr; ++}; ++ ++static bool ggml_cuda_cachegen_error(char * error, size_t capacity, const char * message) { ++ if (error != nullptr && capacity > 0) { ++ snprintf(error, capacity, "%s", message); ++ } ++ return false; ++} ++ ++static bool ggml_cuda_cachegen_cuda_error( ++ char * error, size_t capacity, const char * operation, cudaError_t status) { ++ if (error != nullptr && capacity > 0) { ++ snprintf(error, capacity, "%s: %s", operation, cudaGetErrorString(status)); ++ } ++ return false; ++} ++ ++static cudaError_t ggml_cuda_cachegen_synchronize(const std::vector & jobs) { ++ std::vector devices; ++ for (const ggml_cuda_cachegen_staged_job & job : jobs) { ++ if (std::find(devices.begin(), devices.end(), job.physical_device) != devices.end()) { ++ continue; ++ } ++ const cudaError_t set_status = cudaSetDevice(job.physical_device); ++ if (set_status != cudaSuccess) { ++ return set_status; ++ } ++ const cudaError_t sync_status = cudaStreamSynchronize(cudaStreamPerThread); ++ if (sync_status != cudaSuccess) { ++ return sync_status; ++ } ++ devices.push_back(job.physical_device); ++ } ++ return cudaSuccess; ++} ++ ++static cudaError_t ggml_cuda_cachegen_release(std::vector & jobs) { ++ cudaError_t first_error = cudaSuccess; ++ for (ggml_cuda_cachegen_staged_job & job : jobs) { ++ cudaError_t status = cudaSetDevice(job.physical_device); ++ if (status != cudaSuccess) { ++ if (first_error == cudaSuccess) { ++ first_error = status; ++ } ++ continue; ++ } ++ void * allocations[] = { job.payload_device, job.tiles_device, job.prefixes_device, job.cells_device }; ++ for (void * allocation : allocations) { ++ if (allocation == nullptr) { ++ continue; ++ } ++ status = cudaFree(allocation); ++ if (status != cudaSuccess && first_error == cudaSuccess) { ++ first_error = status; ++ } ++ } ++ job.payload_device = nullptr; ++ job.tiles_device = nullptr; ++ job.prefixes_device = nullptr; ++ job.cells_device = nullptr; ++ } ++ return first_error; ++} ++ ++static bool ggml_cuda_cachegen_decode_f16( ++ const struct ggml_backend_cachegen_job * jobs, ++ size_t job_count, ++ char * error, ++ size_t error_capacity) { ++ if (jobs == nullptr || job_count == 0) { ++ return ggml_cuda_cachegen_error(error, error_capacity, "CUDA/HIP CacheGen jobs are required"); ++ } ++ ++ std::vector staged_jobs; ++ try { ++ staged_jobs.reserve(job_count); ++ for (size_t job_index = 0; job_index < job_count; ++job_index) { ++ const struct ggml_backend_cachegen_job * job = &jobs[job_index]; ++ if (job->dst == nullptr || job->dst->data == nullptr || job->dst->buffer == nullptr || ++ !ggml_backend_buffer_is_cuda(job->dst->buffer) || job->dst->buffer->context == nullptr || ++ job->tiles == nullptr || job->tile_count == 0 || job->tile_count > UINT32_MAX || ++ job->cells == nullptr || job->cell_count == 0 || job->cell_count > UINT32_MAX || ++ job->cell_count > SIZE_MAX / sizeof(uint32_t) || job->channels == 0 || ++ job->channels == UINT32_MAX || ++ job->dst->type != GGML_TYPE_F16 || job->token_stride == 0 || job->channel_stride == 0 || ++ job->token_stride % sizeof(uint16_t) != 0 || job->channel_stride % sizeof(uint16_t) != 0) { ++ return ggml_cuda_cachegen_error(error, error_capacity, "invalid CUDA/HIP CacheGen job"); ++ } ++ ++ staged_jobs.emplace_back(); ++ ggml_cuda_cachegen_staged_job & staged = staged_jobs.back(); ++ auto * buffer_context = static_cast(job->dst->buffer->context); ++ staged.physical_device = ggml_cuda_get_physical_device(buffer_context->device); ++ staged.dst = static_cast(job->dst->data); ++ staged.channels = job->channels; ++ staged.token_stride = job->token_stride; ++ staged.channel_stride = job->channel_stride; ++ staged.cells.assign(job->cells, job->cells + job->cell_count); ++ ++ uint64_t max_cell = 0; ++ for (uint32_t cell : staged.cells) { ++ max_cell = std::max(max_cell, static_cast(cell)); ++ } ++ const uint64_t max_channel = static_cast(job->channels) - 1; ++ if (max_cell > (UINT64_MAX - sizeof(uint16_t)) / job->token_stride || ++ max_channel > (UINT64_MAX - max_cell * job->token_stride - sizeof(uint16_t)) / ++ job->channel_stride || ++ max_cell * job->token_stride + max_channel * job->channel_stride + sizeof(uint16_t) > ++ ggml_nbytes(job->dst)) { ++ return ggml_cuda_cachegen_error( ++ error, error_capacity, "CUDA/HIP CacheGen destination geometry is out of bounds"); ++ } ++ ++ if (job->tile_count > SIZE_MAX / sizeof(ggml_cuda_cachegen_tile)) { ++ return ggml_cuda_cachegen_error(error, error_capacity, "CUDA/HIP CacheGen tile geometry overflows"); ++ } ++ staged.tiles.reserve(job->tile_count); ++ for (size_t tile_index = 0; tile_index < job->tile_count; ++tile_index) { ++ const struct ggml_backend_cachegen_tile * tile = &job->tiles[tile_index]; ++ if (tile->payload == nullptr || tile->payload_bytes < 16 || tile->payload_bytes > UINT32_MAX || ++ tile->token_count == 0 || tile->token_offset > job->cell_count || ++ tile->token_count > job->cell_count - tile->token_offset || ++ staged.payload.size() > UINT32_MAX - 3 || ++ staged.prefixes.size() > UINT32_MAX - (static_cast(job->channels) + 1)) { ++ return ggml_cuda_cachegen_error( ++ error, error_capacity, "CUDA/HIP CacheGen tile exceeds bounded geometry"); ++ } ++ ++ const size_t aligned_payload_size = (staged.payload.size() + 3) & ~size_t(3); ++ if (tile->payload_bytes > UINT32_MAX - aligned_payload_size) { ++ return ggml_cuda_cachegen_error( ++ error, error_capacity, "CUDA/HIP CacheGen payload offsets overflow"); ++ } ++ staged.payload.resize(aligned_payload_size, 0); ++ ++ const uint8_t * payload = static_cast(tile->payload); ++ const uint64_t streams_offset_u64 = 16ull + static_cast(tile->token_count) * sizeof(float) + ++ static_cast(job->channels) * 33 * sizeof(uint16_t) + ++ static_cast(job->channels) * sizeof(uint16_t); ++ if (streams_offset_u64 > tile->payload_bytes || streams_offset_u64 > UINT32_MAX) { ++ return ggml_cuda_cachegen_error( ++ error, error_capacity, "CUDA/HIP CacheGen tile metadata is truncated"); ++ } ++ const size_t lengths_offset = static_cast(streams_offset_u64) - ++ static_cast(job->channels) * sizeof(uint16_t); ++ const size_t streams_offset = static_cast(streams_offset_u64); ++ const ggml_cuda_cachegen_tile encoded_tile = { ++ static_cast(staged.payload.size()), ++ static_cast(staged.prefixes.size()), ++ tile->token_offset, ++ tile->token_count, ++ }; ++ uint32_t prefix = 0; ++ staged.prefixes.push_back(prefix); ++ for (uint32_t channel = 0; channel < job->channels; ++channel) { ++ const size_t offset = lengths_offset + static_cast(channel) * sizeof(uint16_t); ++ const uint32_t length = static_cast(payload[offset]) | ++ static_cast(payload[offset + 1]) << 8; ++ if (prefix > UINT32_MAX - length) { ++ return ggml_cuda_cachegen_error( ++ error, error_capacity, "CUDA/HIP CacheGen stream offsets overflow"); ++ } ++ prefix += length; ++ staged.prefixes.push_back(prefix); ++ } ++ const uint32_t declared_stream_bytes = static_cast(payload[12]) | ++ static_cast(payload[13]) << 8 | static_cast(payload[14]) << 16 | ++ static_cast(payload[15]) << 24; ++ if (prefix != declared_stream_bytes || prefix != tile->payload_bytes - streams_offset) { ++ return ggml_cuda_cachegen_error( ++ error, error_capacity, "CUDA/HIP CacheGen stream length is inconsistent"); ++ } ++ staged.tiles.push_back(encoded_tile); ++ staged.payload.insert(staged.payload.end(), payload, payload + tile->payload_bytes); ++ } ++ } ++ } catch (const std::bad_alloc &) { ++ return ggml_cuda_cachegen_error(error, error_capacity, "CUDA/HIP CacheGen host staging allocation failed"); ++ } ++ ++ int original_device = -1; ++ cudaError_t status = cudaGetDevice(&original_device); ++ if (status != cudaSuccess) { ++ return ggml_cuda_cachegen_cuda_error(error, error_capacity, "CUDA/HIP CacheGen device query failed", status); ++ } ++ ++ const auto fail = [&](const char * operation, cudaError_t failure) { ++ (void) ggml_cuda_cachegen_synchronize(staged_jobs); ++ (void) ggml_cuda_cachegen_release(staged_jobs); ++ (void) cudaSetDevice(original_device); ++ return ggml_cuda_cachegen_cuda_error(error, error_capacity, operation, failure); ++ }; ++ ++ for (ggml_cuda_cachegen_staged_job & job : staged_jobs) { ++ status = cudaSetDevice(job.physical_device); ++ if (status != cudaSuccess) { ++ return fail("CUDA/HIP CacheGen device selection failed", status); ++ } ++#define GGML_CACHEGEN_CUDA_CALL(call, operation) \ ++ do { \ ++ status = (call); \ ++ if (status != cudaSuccess) { \ ++ return fail((operation), status); \ ++ } \ ++ } while (0) ++ GGML_CACHEGEN_CUDA_CALL(cudaMalloc(reinterpret_cast(&job.payload_device), job.payload.size()), ++ "CUDA/HIP CacheGen payload allocation failed"); ++ GGML_CACHEGEN_CUDA_CALL(cudaMalloc(reinterpret_cast(&job.tiles_device), ++ job.tiles.size() * sizeof(ggml_cuda_cachegen_tile)), ++ "CUDA/HIP CacheGen tile allocation failed"); ++ GGML_CACHEGEN_CUDA_CALL(cudaMalloc(reinterpret_cast(&job.prefixes_device), ++ job.prefixes.size() * sizeof(uint32_t)), ++ "CUDA/HIP CacheGen prefix allocation failed"); ++ GGML_CACHEGEN_CUDA_CALL(cudaMalloc(reinterpret_cast(&job.cells_device), ++ job.cells.size() * sizeof(uint32_t)), ++ "CUDA/HIP CacheGen cell allocation failed"); ++ GGML_CACHEGEN_CUDA_CALL(cudaMemcpyAsync(job.payload_device, job.payload.data(), job.payload.size(), ++ cudaMemcpyHostToDevice, cudaStreamPerThread), ++ "CUDA/HIP CacheGen payload upload failed"); ++ GGML_CACHEGEN_CUDA_CALL(cudaMemcpyAsync(job.tiles_device, job.tiles.data(), ++ job.tiles.size() * sizeof(ggml_cuda_cachegen_tile), ++ cudaMemcpyHostToDevice, cudaStreamPerThread), ++ "CUDA/HIP CacheGen tile upload failed"); ++ GGML_CACHEGEN_CUDA_CALL(cudaMemcpyAsync(job.prefixes_device, job.prefixes.data(), ++ job.prefixes.size() * sizeof(uint32_t), ++ cudaMemcpyHostToDevice, cudaStreamPerThread), ++ "CUDA/HIP CacheGen prefix upload failed"); ++ GGML_CACHEGEN_CUDA_CALL(cudaMemcpyAsync(job.cells_device, job.cells.data(), ++ job.cells.size() * sizeof(uint32_t), ++ cudaMemcpyHostToDevice, cudaStreamPerThread), ++ "CUDA/HIP CacheGen cell upload failed"); ++ GGML_CACHEGEN_CUDA_CALL( ++ ggml_cuda_cachegen_decode_f16_launch(job.payload_device, job.tiles_device, job.prefixes_device, ++ job.cells_device, job.dst, job.channels, ++ static_cast(job.tiles.size()), job.token_stride, ++ job.channel_stride, cudaStreamPerThread), ++ "CUDA/HIP CacheGen kernel launch failed"); ++#undef GGML_CACHEGEN_CUDA_CALL ++ } ++ ++ status = ggml_cuda_cachegen_synchronize(staged_jobs); ++ if (status != cudaSuccess) { ++ return fail("CUDA/HIP CacheGen kernel execution failed", status); ++ } ++ status = ggml_cuda_cachegen_release(staged_jobs); ++ if (status != cudaSuccess) { ++ (void) cudaSetDevice(original_device); ++ return ggml_cuda_cachegen_cuda_error(error, error_capacity, "CUDA/HIP CacheGen staging release failed", status); ++ } ++ status = cudaSetDevice(original_device); ++ if (status != cudaSuccess) { ++ return ggml_cuda_cachegen_cuda_error(error, error_capacity, "CUDA/HIP CacheGen device restore failed", status); ++ } ++ return true; ++} ++ ++#endif // !defined(GGML_USE_MUSA) ++ + static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) { + GGML_UNUSED(reg); + if (strcmp(name, "ggml_backend_comm_init") == 0) { +@@ -5687,6 +5963,11 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con + if (strcmp(name, "ggml_backend_get_features") == 0) { + return (void *)ggml_backend_cuda_get_features; + } ++#if !defined(GGML_USE_MUSA) ++ if (strcmp(name, "ggml_backend_cachegen_decode_f16") == 0) { ++ return (void *)ggml_cuda_cachegen_decode_f16; ++ } ++#endif // !defined(GGML_USE_MUSA) + return nullptr; + } + +-- +2.54.0 (Apple Git-157) + diff --git a/third_party/llama.cpp/patches/0039-ggml-metal-align-staged-CacheGen-tiles.patch b/third_party/llama.cpp/patches/0039-ggml-metal-align-staged-CacheGen-tiles.patch new file mode 100644 index 0000000000..697b5c3c13 --- /dev/null +++ b/third_party/llama.cpp/patches/0039-ggml-metal-align-staged-CacheGen-tiles.patch @@ -0,0 +1,117 @@ +From 717bca3cc3cc9ef8272cb37c63c7c700ab567009 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 11 Sep 2026 17:43:11 +1000 +Subject: [PATCH] ggml-metal: align staged CacheGen tiles + +Assisted-by: scama +--- + ggml/src/ggml-metal/ggml-metal-device.m | 15 ++++++++++++-- + tests/test-skippy-cachegen-metal.cpp | 27 ++++++++++++++++++++++++- + 2 files changed, 39 insertions(+), 3 deletions(-) + +diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m +index c6e22a863..d05b7c32c 100644 +--- a/ggml/src/ggml-metal/ggml-metal-device.m ++++ b/ggml/src/ggml-metal/ggml-metal-device.m +@@ -2721,7 +2721,7 @@ bool ggml_metal_cachegen_decode_f16( + const struct ggml_backend_cachegen_tile * tile = &job->tiles[tile_index]; + if (tile->payload == NULL || tile->payload_bytes < 16 || tile->token_count == 0 || + tile->token_offset > job->cell_count || tile->token_count > job->cell_count - tile->token_offset || +- payload_data.length > UINT32_MAX || tile->payload_bytes > UINT32_MAX - payload_data.length || ++ payload_data.length > UINT32_MAX - 3 || + prefix_data.length / sizeof(uint32_t) > UINT32_MAX - ((size_t) job->channels + 1)) { + [payload_data release]; + [tile_data release]; +@@ -2730,6 +2730,17 @@ bool ggml_metal_cachegen_decode_f16( + [temporary_buffers release]; + return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen tile exceeds bounded geometry"); + } ++ const size_t aligned_payload_size = (payload_data.length + 3) & ~(size_t) 3; ++ if (tile->payload_bytes > UINT32_MAX - aligned_payload_size) { ++ [payload_data release]; ++ [tile_data release]; ++ [prefix_data release]; ++ [encoder endEncoding]; ++ [temporary_buffers release]; ++ return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen payload offsets overflow"); ++ } ++ const uint8_t padding[3] = { 0 }; ++ [payload_data appendBytes:padding length:aligned_payload_size - payload_data.length]; + const uint8_t * payload = (const uint8_t *) tile->payload; + const size_t lengths_offset = 16 + (size_t) tile->token_count * sizeof(float) + + (size_t) job->channels * 33 * sizeof(uint16_t); +@@ -2743,7 +2754,7 @@ bool ggml_metal_cachegen_decode_f16( + return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen tile metadata is truncated"); + } + struct ggml_metal_cachegen_tile encoded_tile = { +- (uint32_t) payload_data.length, ++ (uint32_t) aligned_payload_size, + (uint32_t) (prefix_data.length / sizeof(uint32_t)), + tile->token_offset, + tile->token_count, +diff --git a/tests/test-skippy-cachegen-metal.cpp b/tests/test-skippy-cachegen-metal.cpp +index b64ba373c..f300b18b0 100644 +--- a/tests/test-skippy-cachegen-metal.cpp ++++ b/tests/test-skippy-cachegen-metal.cpp +@@ -148,6 +148,8 @@ int main() { + ggml_tensor * transposed16 = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, capacity, channels); + ggml_tensor * row_major32 = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, channels, capacity); + ggml_tensor * transposed32 = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, capacity, channels); ++ constexpr uint32_t multi_capacity = rows * 2 + 5; ++ ggml_tensor * multi_tile = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, channels, multi_capacity); + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (buffer == nullptr) { + ggml_free(ctx); +@@ -167,11 +169,20 @@ int main() { + } + const ggml_backend_cachegen_tile tile16 = { encoded16, sizeof(encoded16), 0, rows }; + const ggml_backend_cachegen_tile tile32 = { encoded32, sizeof(encoded32), 0, rows }; ++ const ggml_backend_cachegen_tile multi_tiles[] = { ++ { encoded16, sizeof(encoded16), 0, rows }, ++ { encoded32, sizeof(encoded32), rows, rows }, ++ }; ++ std::vector multi_cells(rows * 2); ++ for (uint32_t row = 0; row < rows * 2; ++row) { ++ multi_cells[row] = multi_capacity - 1 - row; ++ } + const ggml_backend_cachegen_job jobs[] = { + { row_major16, &tile16, 1, cells.data(), cells.size(), channels, channels * 2, 2 }, + { transposed16, &tile16, 1, cells.data(), cells.size(), channels, 2, capacity * 2 }, + { row_major32, &tile32, 1, cells.data(), cells.size(), channels, channels * 2, 2 }, + { transposed32, &tile32, 1, cells.data(), cells.size(), channels, 2, capacity * 2 }, ++ { multi_tile, multi_tiles, 2, multi_cells.data(), multi_cells.size(), channels, channels * 2, 2 }, + }; + ggml_backend_dev_t device = ggml_backend_get_device(backend); + auto decode = reinterpret_cast( +@@ -181,7 +192,7 @@ int main() { + return finish(4); + } + char error[256] = {}; +- if (!decode(jobs, 4, error, sizeof(error))) { ++ if (!decode(jobs, 5, error, sizeof(error))) { + return finish(5); + } + ggml_backend_cachegen_job invalid_job = jobs[0]; +@@ -220,5 +231,19 @@ int main() { + if (const int status = check_fixture(row_major32, transposed32, expected32, 7)) { + return finish(status); + } ++ std::vector multi_bytes(ggml_nbytes(multi_tile)); ++ ggml_backend_tensor_get(multi_tile, multi_bytes.data(), 0, multi_bytes.size()); ++ const uint8_t * expected_tiles[] = { expected16, expected32 }; ++ for (uint32_t tile = 0; tile < 2; ++tile) { ++ for (uint32_t row = 0; row < rows; ++row) { ++ const uint32_t source_row = tile * rows + row; ++ if (std::memcmp( ++ multi_bytes.data() + static_cast(multi_cells[source_row]) * channels * 2, ++ expected_tiles[tile] + static_cast(row) * channels * 2, ++ channels * 2) != 0) { ++ return finish(9); ++ } ++ } ++ } + return finish(0); + } +-- +2.54.0 (Apple Git-157) + diff --git a/third_party/llama.cpp/patches/0040-ggml-optimize-CacheGen-arithmetic-decode.patch b/third_party/llama.cpp/patches/0040-ggml-optimize-CacheGen-arithmetic-decode.patch new file mode 100644 index 0000000000..76ae875671 --- /dev/null +++ b/third_party/llama.cpp/patches/0040-ggml-optimize-CacheGen-arithmetic-decode.patch @@ -0,0 +1,126 @@ +From c18e4c75324d89f866236ca2fa631c23801a0830 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 11 Sep 2026 17:51:30 +1000 +Subject: [PATCH] ggml: optimize CacheGen arithmetic decode + +Assisted-by: scama +--- + ggml/src/ggml-cuda/cachegen.cu | 39 +++++++++++++++++++--- + ggml/src/ggml-metal/kernels/cachegen.metal | 37 +++++++++++++++++--- + 2 files changed, 66 insertions(+), 10 deletions(-) + +diff --git a/ggml/src/ggml-cuda/cachegen.cu b/ggml/src/ggml-cuda/cachegen.cu +index 5d91aa0ca..e8460cc99 100644 +--- a/ggml/src/ggml-cuda/cachegen.cu ++++ b/ggml/src/ggml-cuda/cachegen.cu +@@ -11,6 +11,38 @@ static __device__ __forceinline__ uint32_t cachegen_read_bit( + return result; + } + ++static __device__ __forceinline__ uint16_t cachegen_scaled_count( ++ uint32_t value, uint32_t low, uint64_t span) { ++ const uint64_t delta = uint64_t(value) - uint64_t(low) + 1; ++ const uint64_t numerator = delta * 65536ull - 1; ++ uint32_t count = min(uint32_t(float(delta) * (65536.0f / float(span))), 65535u); ++ uint64_t product = uint64_t(count) * span; ++ while (product > numerator) { ++ --count; ++ product -= span; ++ } ++ while (count < 65535u && product + span <= numerator) { ++ ++count; ++ product += span; ++ } ++ return uint16_t(count); ++} ++ ++static __device__ __forceinline__ uint32_t cachegen_find_symbol( ++ const uint16_t * cdf, uint16_t count) { ++ uint32_t first = 0; ++ uint32_t last = 33; ++ while (first < last) { ++ const uint32_t middle = (first + last) >> 1; ++ if (cdf[middle] <= count) { ++ first = middle + 1; ++ } else { ++ last = middle; ++ } ++ } ++ return first - 1; ++} ++ + static __global__ void cachegen_decode_f16( + const uint8_t * payload, + const ggml_cuda_cachegen_tile * tiles, +@@ -50,11 +82,8 @@ static __global__ void cachegen_decode_f16( + const float center = float(bins / 2 - 1); + for (uint32_t row = 0; row < tile.rows; ++row) { + const uint64_t span = uint64_t(high) - uint64_t(low) + 1; +- const uint16_t count = uint16_t((((uint64_t(value) - uint64_t(low) + 1) * 65536ull - 1) / span)); +- uint32_t symbol = 0; +- while (symbol + 1 < 33 && cdf[symbol + 1] <= count) { +- ++symbol; +- } ++ const uint16_t count = cachegen_scaled_count(value, low, span); ++ const uint32_t symbol = cachegen_find_symbol(cdf, count); + if (symbol >= 32) { + return; + } +diff --git a/ggml/src/ggml-metal/kernels/cachegen.metal b/ggml/src/ggml-metal/kernels/cachegen.metal +index 82684a71e..22592c070 100644 +--- a/ggml/src/ggml-metal/kernels/cachegen.metal ++++ b/ggml/src/ggml-metal/kernels/cachegen.metal +@@ -22,6 +22,36 @@ static inline uint cachegen_read_bit(const device uchar * stream, uint stream_by + return result; + } + ++static inline ushort cachegen_scaled_count(uint value, uint low, ulong span) { ++ const ulong delta = ulong(value) - ulong(low) + 1; ++ const ulong numerator = delta * 65536ul - 1; ++ uint count = min(uint(float(delta) * (65536.0f / float(span))), 65535u); ++ ulong product = ulong(count) * span; ++ while (product > numerator) { ++ --count; ++ product -= span; ++ } ++ while (count < 65535u && product + span <= numerator) { ++ ++count; ++ product += span; ++ } ++ return ushort(count); ++} ++ ++static inline uint cachegen_find_symbol(const device ushort * cdf, ushort count) { ++ uint first = 0; ++ uint last = 33; ++ while (first < last) { ++ const uint middle = (first + last) >> 1; ++ if (cdf[middle] <= count) { ++ first = middle + 1; ++ } else { ++ last = middle; ++ } ++ } ++ return first - 1; ++} ++ + kernel void kernel_cachegen_decode_f16(const device uchar * payload [[buffer(0)]], + const device cachegen_tile * tiles [[buffer(1)]], + const device uint * prefixes [[buffer(2)]], +@@ -58,11 +88,8 @@ kernel void kernel_cachegen_decode_f16(const device uchar * payload [[bu + const float center = float(bins / 2 - 1); + for (uint row = 0; row < tile.rows; ++row) { + const ulong span = ulong(high) - ulong(low) + 1; +- const ushort count = ushort((((ulong(value) - ulong(low) + 1) * 65536ul - 1) / span)); +- uint symbol = 0; +- while (symbol + 1 < 33 && cdf[symbol + 1] <= count) { +- ++symbol; +- } ++ const ushort count = cachegen_scaled_count(value, low, span); ++ const uint symbol = cachegen_find_symbol(cdf, count); + if (symbol >= 32) { + return; + } +-- +2.54.0 (Apple Git-157) + diff --git a/third_party/llama.cpp/patches/0041-ggml-decode-CacheGen-into-F32-KV-tensors.patch b/third_party/llama.cpp/patches/0041-ggml-decode-CacheGen-into-F32-KV-tensors.patch new file mode 100644 index 0000000000..a732fb8818 --- /dev/null +++ b/third_party/llama.cpp/patches/0041-ggml-decode-CacheGen-into-F32-KV-tensors.patch @@ -0,0 +1,559 @@ +From 980d2df45910b4d07138a863811911c1a2d8f6d0 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 11 Sep 2026 18:22:56 +1000 +Subject: [PATCH] ggml: decode CacheGen into F32 KV tensors + +Replace the F16-specific backend hook with a typed decoder and write the entropy-decoded half values directly into F16 or F32 resident layouts. + +Assisted-by: scama +--- + ggml/include/ggml-backend.h | 6 +-- + ggml/src/ggml-cuda/cachegen.cu | 18 +++++--- + ggml/src/ggml-cuda/cachegen.cuh | 3 +- + ggml/src/ggml-cuda/ggml-cuda.cu | 28 +++++++------ + ggml/src/ggml-metal/ggml-metal-device.h | 2 +- + ggml/src/ggml-metal/ggml-metal-device.m | 18 +++++--- + ggml/src/ggml-metal/ggml-metal.cpp | 4 +- + ggml/src/ggml-metal/kernels/cachegen.metal | 21 ++++++---- + include/skippy/state.h | 2 + + src/llama-kv-cache.cpp | 49 +++++++++++++--------- + tests/test-skippy-cachegen-metal.cpp | 32 ++++++++++---- + 11 files changed, 117 insertions(+), 66 deletions(-) + +diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h +index e2460fcda..2256fecbc 100644 +--- a/ggml/include/ggml-backend.h ++++ b/ggml/include/ggml-backend.h +@@ -224,9 +224,9 @@ extern "C" { + }; + typedef struct ggml_backend_feature * (*ggml_backend_get_features_t)(ggml_backend_reg_t reg); + +- // Optional backend entry point for decoding validated CacheGen F16 tiles ++ // Optional backend entry point for decoding validated CacheGen tiles + // directly into resident tensors. Callers resolve +- // `ggml_backend_cachegen_decode_f16` through get_proc_address. ++ // `ggml_backend_cachegen_decode` through get_proc_address. + struct ggml_backend_cachegen_tile { + const void * payload; + size_t payload_bytes; +@@ -245,7 +245,7 @@ extern "C" { + uint64_t channel_stride; + }; + +- typedef bool (*ggml_backend_cachegen_decode_f16_t)( ++ typedef bool (*ggml_backend_cachegen_decode_t)( + const struct ggml_backend_cachegen_job * jobs, + size_t job_count, + char * error, +diff --git a/ggml/src/ggml-cuda/cachegen.cu b/ggml/src/ggml-cuda/cachegen.cu +index e8460cc99..cc51a634c 100644 +--- a/ggml/src/ggml-cuda/cachegen.cu ++++ b/ggml/src/ggml-cuda/cachegen.cu +@@ -43,7 +43,7 @@ static __device__ __forceinline__ uint32_t cachegen_find_symbol( + return first - 1; + } + +-static __global__ void cachegen_decode_f16( ++static __global__ void cachegen_decode( + const uint8_t * payload, + const ggml_cuda_cachegen_tile * tiles, + const uint32_t * prefixes, +@@ -52,7 +52,8 @@ static __global__ void cachegen_decode_f16( + uint32_t channels, + uint32_t tile_count, + uint64_t token_stride, +- uint64_t channel_stride) { ++ uint64_t channel_stride, ++ uint32_t element_bytes) { + const uint32_t channel = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t tile_index = blockIdx.y; + if (channel >= channels || tile_index >= tile_count) { +@@ -92,7 +93,11 @@ static __global__ void cachegen_decode_f16( + const half decoded = __float2half(normalized * maxes[row]); + const uint64_t destination = + uint64_t(cells[tile.token_offset + row]) * token_stride + uint64_t(channel) * channel_stride; +- *reinterpret_cast(dst + destination) = decoded; ++ if (element_bytes == sizeof(float)) { ++ *reinterpret_cast(dst + destination) = __half2float(decoded); ++ } else { ++ *reinterpret_cast(dst + destination) = decoded; ++ } + + if (row + 1 == tile.rows) { + break; +@@ -118,7 +123,7 @@ static __global__ void cachegen_decode_f16( + } + } + +-cudaError_t ggml_cuda_cachegen_decode_f16_launch( ++cudaError_t ggml_cuda_cachegen_decode_launch( + const uint8_t * payload, + const ggml_cuda_cachegen_tile * tiles, + const uint32_t * prefixes, +@@ -128,12 +133,13 @@ cudaError_t ggml_cuda_cachegen_decode_f16_launch( + uint32_t tile_count, + uint64_t token_stride, + uint64_t channel_stride, ++ uint32_t element_bytes, + cudaStream_t stream) { + constexpr uint32_t block_size = 64; + const dim3 block(block_size, 1, 1); + const dim3 grid((channels + block_size - 1) / block_size, tile_count, 1); +- cachegen_decode_f16<<>>( +- payload, tiles, prefixes, cells, dst, channels, tile_count, token_stride, channel_stride); ++ cachegen_decode<<>>( ++ payload, tiles, prefixes, cells, dst, channels, tile_count, token_stride, channel_stride, element_bytes); + return cudaGetLastError(); + } + +diff --git a/ggml/src/ggml-cuda/cachegen.cuh b/ggml/src/ggml-cuda/cachegen.cuh +index d6d216675..7118c9aa9 100644 +--- a/ggml/src/ggml-cuda/cachegen.cuh ++++ b/ggml/src/ggml-cuda/cachegen.cuh +@@ -11,7 +11,7 @@ struct ggml_cuda_cachegen_tile { + uint32_t rows; + }; + +-cudaError_t ggml_cuda_cachegen_decode_f16_launch( ++cudaError_t ggml_cuda_cachegen_decode_launch( + const uint8_t * payload, + const ggml_cuda_cachegen_tile * tiles, + const uint32_t * prefixes, +@@ -21,6 +21,7 @@ cudaError_t ggml_cuda_cachegen_decode_f16_launch( + uint32_t tile_count, + uint64_t token_stride, + uint64_t channel_stride, ++ uint32_t element_bytes, + cudaStream_t stream); + + #endif // !defined(GGML_USE_MUSA) +diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu +index e955c4394..3288653f5 100644 +--- a/ggml/src/ggml-cuda/ggml-cuda.cu ++++ b/ggml/src/ggml-cuda/ggml-cuda.cu +@@ -5675,6 +5675,7 @@ struct ggml_cuda_cachegen_staged_job { + int physical_device = -1; + uint8_t * dst = nullptr; + uint32_t channels = 0; ++ uint32_t element_bytes = 0; + uint64_t token_stride = 0; + uint64_t channel_stride = 0; + std::vector payload; +@@ -5749,7 +5750,7 @@ static cudaError_t ggml_cuda_cachegen_release(std::vectorcells == nullptr || job->cell_count == 0 || job->cell_count > UINT32_MAX || + job->cell_count > SIZE_MAX / sizeof(uint32_t) || job->channels == 0 || + job->channels == UINT32_MAX || +- job->dst->type != GGML_TYPE_F16 || job->token_stride == 0 || job->channel_stride == 0 || +- job->token_stride % sizeof(uint16_t) != 0 || job->channel_stride % sizeof(uint16_t) != 0) { ++ (job->dst->type != GGML_TYPE_F16 && job->dst->type != GGML_TYPE_F32) || ++ job->token_stride == 0 || job->channel_stride == 0 || ++ job->token_stride % ggml_type_size(job->dst->type) != 0 || ++ job->channel_stride % ggml_type_size(job->dst->type) != 0) { + return ggml_cuda_cachegen_error(error, error_capacity, "invalid CUDA/HIP CacheGen job"); + } + +@@ -5780,6 +5783,7 @@ static bool ggml_cuda_cachegen_decode_f16( + staged.physical_device = ggml_cuda_get_physical_device(buffer_context->device); + staged.dst = static_cast(job->dst->data); + staged.channels = job->channels; ++ staged.element_bytes = ggml_type_size(job->dst->type); + staged.token_stride = job->token_stride; + staged.channel_stride = job->channel_stride; + staged.cells.assign(job->cells, job->cells + job->cell_count); +@@ -5789,10 +5793,10 @@ static bool ggml_cuda_cachegen_decode_f16( + max_cell = std::max(max_cell, static_cast(cell)); + } + const uint64_t max_channel = static_cast(job->channels) - 1; +- if (max_cell > (UINT64_MAX - sizeof(uint16_t)) / job->token_stride || +- max_channel > (UINT64_MAX - max_cell * job->token_stride - sizeof(uint16_t)) / ++ if (max_cell > (UINT64_MAX - staged.element_bytes) / job->token_stride || ++ max_channel > (UINT64_MAX - max_cell * job->token_stride - staged.element_bytes) / + job->channel_stride || +- max_cell * job->token_stride + max_channel * job->channel_stride + sizeof(uint16_t) > ++ max_cell * job->token_stride + max_channel * job->channel_stride + staged.element_bytes > + ggml_nbytes(job->dst)) { + return ggml_cuda_cachegen_error( + error, error_capacity, "CUDA/HIP CacheGen destination geometry is out of bounds"); +@@ -5917,10 +5921,10 @@ static bool ggml_cuda_cachegen_decode_f16( + cudaMemcpyHostToDevice, cudaStreamPerThread), + "CUDA/HIP CacheGen cell upload failed"); + GGML_CACHEGEN_CUDA_CALL( +- ggml_cuda_cachegen_decode_f16_launch(job.payload_device, job.tiles_device, job.prefixes_device, +- job.cells_device, job.dst, job.channels, +- static_cast(job.tiles.size()), job.token_stride, +- job.channel_stride, cudaStreamPerThread), ++ ggml_cuda_cachegen_decode_launch(job.payload_device, job.tiles_device, job.prefixes_device, ++ job.cells_device, job.dst, job.channels, ++ static_cast(job.tiles.size()), job.token_stride, ++ job.channel_stride, job.element_bytes, cudaStreamPerThread), + "CUDA/HIP CacheGen kernel launch failed"); + #undef GGML_CACHEGEN_CUDA_CALL + } +@@ -5964,8 +5968,8 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con + return (void *)ggml_backend_cuda_get_features; + } + #if !defined(GGML_USE_MUSA) +- if (strcmp(name, "ggml_backend_cachegen_decode_f16") == 0) { +- return (void *)ggml_cuda_cachegen_decode_f16; ++ if (strcmp(name, "ggml_backend_cachegen_decode") == 0) { ++ return (void *)ggml_cuda_cachegen_decode; + } + #endif // !defined(GGML_USE_MUSA) + return nullptr; +diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h +index 22e4fb436..0a2cf000c 100644 +--- a/ggml/src/ggml-metal/ggml-metal-device.h ++++ b/ggml/src/ggml-metal/ggml-metal-device.h +@@ -351,7 +351,7 @@ void ggml_metal_buffer_clear (ggml_metal_buffer_t buf, uint8_t value); + // + struct ggml_metal_buffer_id ggml_metal_buffer_get_id(ggml_metal_buffer_t buf, const struct ggml_tensor * t); + +-bool ggml_metal_cachegen_decode_f16( ++bool ggml_metal_cachegen_decode( + const struct ggml_backend_cachegen_job * jobs, + size_t job_count, + char * error, +diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m +index d05b7c32c..d4ae1f475 100644 +--- a/ggml/src/ggml-metal/ggml-metal-device.m ++++ b/ggml/src/ggml-metal/ggml-metal-device.m +@@ -2654,6 +2654,7 @@ struct ggml_metal_buffer_id ggml_metal_buffer_get_id(ggml_metal_buffer_t buf, co + struct ggml_metal_cachegen_params { + uint32_t channels; + uint32_t tile_count; ++ uint32_t element_bytes; + uint64_t token_stride; + uint64_t channel_stride; + }; +@@ -2665,7 +2666,7 @@ static bool ggml_metal_cachegen_error(char * error, size_t capacity, const char + return false; + } + +-bool ggml_metal_cachegen_decode_f16( ++bool ggml_metal_cachegen_decode( + const struct ggml_backend_cachegen_job * jobs, + size_t job_count, + char * error, +@@ -2683,7 +2684,10 @@ bool ggml_metal_cachegen_decode_f16( + const struct ggml_backend_cachegen_job * job = &jobs[job_index]; + if (job->dst == NULL || job->dst->buffer == NULL || job->tiles == NULL || job->tile_count == 0 || + job->tile_count > UINT32_MAX || job->cells == NULL || job->cell_count == 0 || job->channels == 0 || +- job->dst->type != GGML_TYPE_F16 || job->token_stride == 0 || job->channel_stride == 0 || ++ (job->dst->type != GGML_TYPE_F16 && job->dst->type != GGML_TYPE_F32) || ++ job->token_stride == 0 || job->channel_stride == 0 || ++ job->token_stride % ggml_type_size(job->dst->type) != 0 || ++ job->channel_stride % ggml_type_size(job->dst->type) != 0 || + job->cell_count > SIZE_MAX / sizeof(uint32_t)) { + [temporary_buffers release]; + return ggml_metal_cachegen_error(error, error_capacity, "invalid Metal CacheGen job"); +@@ -2693,9 +2697,10 @@ bool ggml_metal_cachegen_decode_f16( + max_cell = MAX(max_cell, job->cells[cell_index]); + } + const uint64_t max_channel = (uint64_t) job->channels - 1; +- if (max_cell > (UINT64_MAX - sizeof(uint16_t)) / job->token_stride || +- max_channel > (UINT64_MAX - max_cell * job->token_stride - sizeof(uint16_t)) / job->channel_stride || +- max_cell * job->token_stride + max_channel * job->channel_stride + sizeof(uint16_t) > ggml_nbytes(job->dst)) { ++ const uint64_t element_bytes = ggml_type_size(job->dst->type); ++ if (max_cell > (UINT64_MAX - element_bytes) / job->token_stride || ++ max_channel > (UINT64_MAX - max_cell * job->token_stride - element_bytes) / job->channel_stride || ++ max_cell * job->token_stride + max_channel * job->channel_stride + element_bytes > ggml_nbytes(job->dst)) { + [temporary_buffers release]; + return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen destination geometry is out of bounds"); + } +@@ -2820,7 +2825,7 @@ bool ggml_metal_cachegen_decode_f16( + + const struct ggml_metal_pipeline_with_params pipeline = + ggml_metal_library_compile_pipeline(device->library, +- "kernel_cachegen_decode_f16", "kernel_cachegen_decode_f16", NULL); ++ "kernel_cachegen_decode", "kernel_cachegen_decode", NULL); + if (pipeline.pipeline == NULL) { + [encoder endEncoding]; + [temporary_buffers release]; +@@ -2835,6 +2840,7 @@ bool ggml_metal_cachegen_decode_f16( + const struct ggml_metal_cachegen_params params = { + job->channels, + (uint32_t) job->tile_count, ++ (uint32_t) element_bytes, + job->token_stride, + job->channel_stride, + }; +diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp +index 6bec83efd..a1c190a82 100644 +--- a/ggml/src/ggml-metal/ggml-metal.cpp ++++ b/ggml/src/ggml-metal/ggml-metal.cpp +@@ -910,8 +910,8 @@ static void * ggml_backend_metal_get_proc_address(ggml_backend_reg_t reg, const + if (strcmp(name, "ggml_backend_get_features") == 0) { + return (void *)ggml_backend_metal_get_features; + } +- if (strcmp(name, "ggml_backend_cachegen_decode_f16") == 0) { +- return (void *)ggml_metal_cachegen_decode_f16; ++ if (strcmp(name, "ggml_backend_cachegen_decode") == 0) { ++ return (void *)ggml_metal_cachegen_decode; + } + if (strcmp(name, "ggml_backend_metal_tuning_set_fa_vec_override") == 0) { + return (void *)ggml_backend_metal_tuning_set_fa_vec_override; +diff --git a/ggml/src/ggml-metal/kernels/cachegen.metal b/ggml/src/ggml-metal/kernels/cachegen.metal +index 22592c070..e426f7993 100644 +--- a/ggml/src/ggml-metal/kernels/cachegen.metal ++++ b/ggml/src/ggml-metal/kernels/cachegen.metal +@@ -10,6 +10,7 @@ struct cachegen_tile { + struct cachegen_params { + uint channels; + uint tile_count; ++ uint element_bytes; + ulong token_stride; + ulong channel_stride; + }; +@@ -52,13 +53,13 @@ static inline uint cachegen_find_symbol(const device ushort * cdf, ushort count) + return first - 1; + } + +-kernel void kernel_cachegen_decode_f16(const device uchar * payload [[buffer(0)]], +- const device cachegen_tile * tiles [[buffer(1)]], +- const device uint * prefixes [[buffer(2)]], +- const device uint * cells [[buffer(3)]], +- device uchar * dst [[buffer(4)]], +- constant cachegen_params & params [[buffer(5)]], +- uint2 gid [[thread_position_in_grid]]) { ++kernel void kernel_cachegen_decode(const device uchar * payload [[buffer(0)]], ++ const device cachegen_tile * tiles [[buffer(1)]], ++ const device uint * prefixes [[buffer(2)]], ++ const device uint * cells [[buffer(3)]], ++ device uchar * dst [[buffer(4)]], ++ constant cachegen_params & params [[buffer(5)]], ++ uint2 gid [[thread_position_in_grid]]) { + const uint channel = gid.x; + const uint tile_index = gid.y; + if (channel >= params.channels || tile_index >= params.tile_count) { +@@ -98,7 +99,11 @@ kernel void kernel_cachegen_decode_f16(const device uchar * payload [[bu + const half decoded = half(normalized * maxes[row]); + const ulong destination = + ulong(cells[tile.token_offset + row]) * params.token_stride + ulong(channel) * params.channel_stride; +- *reinterpret_cast(dst + destination) = decoded; ++ if (params.element_bytes == sizeof(float)) { ++ *reinterpret_cast(dst + destination) = float(decoded); ++ } else { ++ *reinterpret_cast(dst + destination) = decoded; ++ } + + if (row + 1 == tile.rows) { + break; +diff --git a/include/skippy/state.h b/include/skippy/state.h +index 7504e573c..9c78b3698 100644 +--- a/include/skippy/state.h ++++ b/include/skippy/state.h +@@ -33,6 +33,8 @@ enum skippy_cachegen_record_kind { + SKIPPY_CACHEGEN_RECORD_F16 = 0, + SKIPPY_CACHEGEN_RECORD_EXACT = 1, + SKIPPY_CACHEGEN_RECORD_F16_TRANSPOSED = 2, ++ SKIPPY_CACHEGEN_RECORD_F32 = 3, ++ SKIPPY_CACHEGEN_RECORD_F32_TRANSPOSED = 4, + }; + + /** @brief Describes one validated CacheGen record and its logical page destination. */ +diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp +index dfd00139b..0671fb296 100644 +--- a/src/llama-kv-cache.cpp ++++ b/src/llama-kv-cache.cpp +@@ -1970,17 +1970,18 @@ static uint32_t skippy_cachegen_read_u32(const uint8_t * bytes) { + static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & record, + size_t expected_rows, + size_t expected_channels, ++ size_t output_element_bytes, + std::string & error) { + constexpr size_t header_bytes = 16; + constexpr size_t cdf_entries = 33; + constexpr size_t max_rows = 256; + const auto * payload = static_cast(record.payload); + if (record.abi_version != SKIPPY_CACHEGEN_RECORD_V1_ABI_VERSION || record.reserved0 != 0 || +- record.element_bytes != 2 || payload == nullptr || record.payload_bytes < header_bytes || ++ record.element_bytes != output_element_bytes || payload == nullptr || record.payload_bytes < header_bytes || + record.decoded_bytes == 0 || record.token_count != expected_rows || expected_rows == 0 || + expected_rows > max_rows || expected_channels == 0 || + expected_channels > std::numeric_limits::max()) { +- error = "invalid CacheGen F16 record descriptor"; ++ error = "invalid CacheGen record descriptor"; + return false; + } + if (std::memcmp(payload, "LCG1", 4) != 0 || payload[5] != 0 || (payload[4] != 16 && payload[4] != 32) || +@@ -1990,8 +1991,8 @@ static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & r + return false; + } + if (expected_rows > std::numeric_limits::max() / expected_channels || +- expected_rows * expected_channels > std::numeric_limits::max() / 2 || +- record.decoded_bytes != expected_rows * expected_channels * 2) { ++ expected_rows * expected_channels > std::numeric_limits::max() / output_element_bytes || ++ record.decoded_bytes != expected_rows * expected_channels * output_element_bytes) { + error = "CacheGen record decoded geometry is inconsistent"; + return false; + } +@@ -2086,8 +2087,12 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id + if (!stage_import_kv_page(seq_id, desc, &validation_sentinel, desc.payload_bytes, error, true)) { + return false; + } +- if (records == nullptr || record_count == 0 || desc.k_type != GGML_TYPE_F16 || desc.v_type != GGML_TYPE_F16) { +- error = "CacheGen import requires non-empty F16 K/V records"; ++ const auto cachegen_type_supported = [](uint32_t type) { ++ return type == GGML_TYPE_F16 || type == GGML_TYPE_F32; ++ }; ++ if (records == nullptr || record_count == 0 || !cachegen_type_supported(desc.k_type) || ++ !cachegen_type_supported(desc.v_type)) { ++ error = "CacheGen import requires non-empty F16 or F32 K/V records"; + return false; + } + +@@ -2104,12 +2109,17 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id + size_t logical_offset = output_base; + std::vector storage; + storage.reserve(selected.size() * 2); +- auto consume_tiles = [&](ggml_tensor * dst, size_t row_bytes, bool transposed) -> bool { ++ auto consume_tiles = [&](ggml_tensor * dst, size_t row_bytes, size_t element_bytes, bool transposed) -> bool { ++ if (element_bytes == 0 || row_bytes % element_bytes != 0 || ++ (element_bytes != sizeof(uint16_t) && element_bytes != sizeof(float))) { ++ error = "CacheGen destination element geometry is invalid"; ++ return false; ++ } + skippy_cachegen_job_storage entry; + entry.dst = dst; +- entry.token_stride = transposed ? desc.v_element_bytes : row_bytes; +- entry.channel_stride = transposed ? static_cast(v_cells[strm].size()) * desc.v_element_bytes : 2; +- const size_t channels = row_bytes / 2; ++ entry.token_stride = transposed ? element_bytes : row_bytes; ++ entry.channel_stride = transposed ? static_cast(v_cells[strm].size()) * element_bytes : element_bytes; ++ const size_t channels = row_bytes / element_bytes; + entry.channels = static_cast(channels); + size_t token_offset = 0; + while (token_offset < desc.token_count) { +@@ -2119,13 +2129,14 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id + } + const auto & record = records[record_index++]; + const size_t rows = static_cast(record.token_count); +- const uint32_t expected_kind = +- transposed ? SKIPPY_CACHEGEN_RECORD_F16_TRANSPOSED : SKIPPY_CACHEGEN_RECORD_F16; ++ const uint32_t expected_kind = element_bytes == sizeof(float) ? ++ (transposed ? SKIPPY_CACHEGEN_RECORD_F32_TRANSPOSED : SKIPPY_CACHEGEN_RECORD_F32) : ++ (transposed ? SKIPPY_CACHEGEN_RECORD_F16_TRANSPOSED : SKIPPY_CACHEGEN_RECORD_F16); + if (record.kind != expected_kind || rows == 0 || rows > desc.token_count - token_offset || + record.output_offset != logical_offset + (transposed ? 0 : token_offset * row_bytes) || + record.token_start != (transposed ? token_offset : 0) || + record.total_tokens != (transposed ? desc.token_count : 0) || +- !skippy_cachegen_validate_segment(record, rows, channels, error)) { ++ !skippy_cachegen_validate_segment(record, rows, channels, element_bytes, error)) { + if (error.empty()) { + error = "CacheGen record order or destination geometry is invalid"; + } +@@ -2141,14 +2152,14 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id + }; + + for (const auto * layer : selected) { +- if (!consume_tiles(layer->k_stream[strm], desc.k_row_bytes, false)) { ++ if (!consume_tiles(layer->k_stream[strm], desc.k_row_bytes, ggml_type_size(layer->k_stream[strm]->type), false)) { + return false; + } + } + for (const auto * layer : selected) { + const size_t v_row_bytes = + v_trans ? static_cast(hparams.n_embd_v_gqa(layer->il)) * desc.v_element_bytes : desc.v_row_bytes; +- if (!consume_tiles(layer->v_stream[strm], v_row_bytes, v_trans)) { ++ if (!consume_tiles(layer->v_stream[strm], v_row_bytes, ggml_type_size(layer->v_stream[strm]->type), v_trans)) { + return false; + } + } +@@ -2182,7 +2193,7 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id + return false; + } + +- ggml_backend_cachegen_decode_f16_t decode = nullptr; ++ ggml_backend_cachegen_decode_t decode = nullptr; + for (const auto & entry : storage) { + if (entry.dst == nullptr || entry.dst->buffer == nullptr) { + error = "CacheGen destination tensor has no backend buffer"; +@@ -2193,12 +2204,12 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id + const auto reg = device == nullptr ? nullptr : ggml_backend_dev_backend_reg(device); + auto candidate = reg == nullptr ? + nullptr : +- reinterpret_cast( +- ggml_backend_reg_get_proc_address(reg, "ggml_backend_cachegen_decode_f16")); ++ reinterpret_cast( ++ ggml_backend_reg_get_proc_address(reg, "ggml_backend_cachegen_decode")); + if (candidate == nullptr || (decode != nullptr && candidate != decode)) { + invalid_argument = false; + unsupported = true; +- error = "resident KV backend does not provide CacheGen F16 decode"; ++ error = "resident KV backend does not provide typed CacheGen decode"; + return false; + } + decode = candidate; +diff --git a/tests/test-skippy-cachegen-metal.cpp b/tests/test-skippy-cachegen-metal.cpp +index f300b18b0..dac796d25 100644 +--- a/tests/test-skippy-cachegen-metal.cpp ++++ b/tests/test-skippy-cachegen-metal.cpp +@@ -146,8 +146,8 @@ int main() { + } + ggml_tensor * row_major16 = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, channels, capacity); + ggml_tensor * transposed16 = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, capacity, channels); +- ggml_tensor * row_major32 = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, channels, capacity); +- ggml_tensor * transposed32 = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, capacity, channels); ++ ggml_tensor * row_major32 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, channels, capacity); ++ ggml_tensor * transposed32 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, capacity, channels); + constexpr uint32_t multi_capacity = rows * 2 + 5; + ggml_tensor * multi_tile = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, channels, multi_capacity); + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); +@@ -180,14 +180,14 @@ int main() { + const ggml_backend_cachegen_job jobs[] = { + { row_major16, &tile16, 1, cells.data(), cells.size(), channels, channels * 2, 2 }, + { transposed16, &tile16, 1, cells.data(), cells.size(), channels, 2, capacity * 2 }, +- { row_major32, &tile32, 1, cells.data(), cells.size(), channels, channels * 2, 2 }, +- { transposed32, &tile32, 1, cells.data(), cells.size(), channels, 2, capacity * 2 }, ++ { row_major32, &tile32, 1, cells.data(), cells.size(), channels, channels * 4, 4 }, ++ { transposed32, &tile32, 1, cells.data(), cells.size(), channels, 4, capacity * 4 }, + { multi_tile, multi_tiles, 2, multi_cells.data(), multi_cells.size(), channels, channels * 2, 2 }, + }; + ggml_backend_dev_t device = ggml_backend_get_device(backend); +- auto decode = reinterpret_cast( ++ auto decode = reinterpret_cast( + ggml_backend_reg_get_proc_address( +- ggml_backend_dev_backend_reg(device), "ggml_backend_cachegen_decode_f16")); ++ ggml_backend_dev_backend_reg(device), "ggml_backend_cachegen_decode")); + if (decode == nullptr) { + return finish(4); + } +@@ -228,8 +228,24 @@ int main() { + if (const int status = check_fixture(row_major16, transposed16, expected16, 6)) { + return finish(status); + } +- if (const int status = check_fixture(row_major32, transposed32, expected32, 7)) { +- return finish(status); ++ std::vector expected32_f32(rows * channels); ++ for (size_t i = 0; i < expected32_f32.size(); ++i) { ++ uint16_t bits; ++ std::memcpy(&bits, expected32 + i * sizeof(bits), sizeof(bits)); ++ expected32_f32[i] = ggml_fp16_to_fp32(bits); ++ } ++ std::vector row32(ggml_nelements(row_major32)); ++ std::vector transposed32_values(ggml_nelements(transposed32)); ++ ggml_backend_tensor_get(row_major32, row32.data(), 0, ggml_nbytes(row_major32)); ++ ggml_backend_tensor_get(transposed32, transposed32_values.data(), 0, ggml_nbytes(transposed32)); ++ for (uint32_t row = 0; row < rows; ++row) { ++ for (uint32_t channel = 0; channel < channels; ++channel) { ++ const float expected = expected32_f32[static_cast(row) * channels + channel]; ++ if (row32[static_cast(cells[row]) * channels + channel] != expected || ++ transposed32_values[cells[row] + static_cast(channel) * capacity] != expected) { ++ return finish(7); ++ } ++ } + } + std::vector multi_bytes(ggml_nbytes(multi_tile)); + ggml_backend_tensor_get(multi_tile, multi_bytes.data(), 0, multi_bytes.size()); +-- +2.54.0 (Apple Git-157) + diff --git a/third_party/llama.cpp/patches/0042-ggml-restore-quantized-CacheGen-pages-on-device.patch b/third_party/llama.cpp/patches/0042-ggml-restore-quantized-CacheGen-pages-on-device.patch new file mode 100644 index 0000000000..15fb3fad3a --- /dev/null +++ b/third_party/llama.cpp/patches/0042-ggml-restore-quantized-CacheGen-pages-on-device.patch @@ -0,0 +1,501 @@ +From 81b1c8f52510ff07b4171020a5892f0a1e5bcead Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 11 Sep 2026 18:49:55 +1000 +Subject: [PATCH] feat(skippy): restore quantized CacheGen pages on device + +Assisted-by: scama +--- + ggml/src/ggml-cuda/cachegen.cu | 47 +++++++++- + ggml/src/ggml-cuda/ggml-cuda.cu | 9 +- + ggml/src/ggml-metal/ggml-metal-device.m | 14 ++- + ggml/src/ggml-metal/kernels/cachegen.metal | 44 ++++++++- + include/skippy/state.h | 2 + + src/llama-kv-cache.cpp | 51 ++++++++--- + tests/test-skippy-cachegen-metal.cpp | 102 ++++++++++++++++++++- + 7 files changed, 240 insertions(+), 29 deletions(-) + +diff --git a/ggml/src/ggml-cuda/cachegen.cu b/ggml/src/ggml-cuda/cachegen.cu +index cc51a634c..fa5683788 100644 +--- a/ggml/src/ggml-cuda/cachegen.cu ++++ b/ggml/src/ggml-cuda/cachegen.cu +@@ -91,11 +91,52 @@ static __global__ void cachegen_decode( + + const float normalized = (float(symbol) - center) / center; + const half decoded = __float2half(normalized * maxes[row]); +- const uint64_t destination = +- uint64_t(cells[tile.token_offset + row]) * token_stride + uint64_t(channel) * channel_stride; +- if (element_bytes == sizeof(float)) { ++ const float decoded_f = __half2float(decoded); ++ const uint64_t row_offset = uint64_t(cells[tile.token_offset + row]) * token_stride; ++ if (element_bytes == 34) { ++ const uint32_t lane = channel & 31; ++ const uint64_t destination = row_offset + uint64_t(channel >> 5) * channel_stride; ++ float amax = fabsf(decoded_f); ++#pragma unroll ++ for (uint32_t mask = 16; mask > 0; mask >>= 1) { ++ amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, mask, 32)); ++ } ++ const float scale = amax / 127.0f; ++ if (lane == 0) { ++ *reinterpret_cast(dst + destination) = __float2half(scale); ++ } ++ const int quantized = int(roundf(decoded_f * (scale == 0.0f ? 0.0f : 1.0f / scale))); ++ *(reinterpret_cast(dst + destination + 2) + lane) = int8_t(max(-127, min(127, quantized))); ++ } else if (element_bytes == 18) { ++ const uint32_t lane = channel & 31; ++ const uint64_t destination = row_offset + uint64_t(channel >> 5) * channel_stride; ++ float amax = fabsf(decoded_f); ++#pragma unroll ++ for (uint32_t mask = 16; mask > 0; mask >>= 1) { ++ amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, mask, 32)); ++ } ++ uint32_t winner = fabsf(decoded_f) == amax ? lane : 32; ++#pragma unroll ++ for (uint32_t mask = 16; mask > 0; mask >>= 1) { ++ winner = min(winner, __shfl_xor_sync(0xFFFFFFFF, winner, mask, 32)); ++ } ++ const float signed_max = __shfl_sync(0xFFFFFFFF, decoded_f, winner, 32); ++ const float scale = signed_max / -8.0f; ++ if (lane == 0) { ++ *reinterpret_cast(dst + destination) = __float2half(scale); ++ } ++ const float high_value = __shfl_down_sync(0xFFFFFFFF, decoded_f, 16, 32); ++ if (lane < 16) { ++ const float inverse = scale == 0.0f ? 0.0f : 1.0f / scale; ++ const int low = max(0, min(15, int(decoded_f * inverse + 8.5f))); ++ const int high = max(0, min(15, int(high_value * inverse + 8.5f))); ++ *(dst + destination + 2 + lane) = uint8_t(low | (high << 4)); ++ } ++ } else if (element_bytes == sizeof(float)) { ++ const uint64_t destination = row_offset + uint64_t(channel) * channel_stride; + *reinterpret_cast(dst + destination) = __half2float(decoded); + } else { ++ const uint64_t destination = row_offset + uint64_t(channel) * channel_stride; + *reinterpret_cast(dst + destination) = decoded; + } + +diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu +index 3288653f5..b68649ec4 100644 +--- a/ggml/src/ggml-cuda/ggml-cuda.cu ++++ b/ggml/src/ggml-cuda/ggml-cuda.cu +@@ -5764,13 +5764,17 @@ static bool ggml_cuda_cachegen_decode( + staged_jobs.reserve(job_count); + for (size_t job_index = 0; job_index < job_count; ++job_index) { + const struct ggml_backend_cachegen_job * job = &jobs[job_index]; ++ const bool quantized = job->dst != nullptr && ++ (job->dst->type == GGML_TYPE_Q8_0 || job->dst->type == GGML_TYPE_Q4_0); + if (job->dst == nullptr || job->dst->data == nullptr || job->dst->buffer == nullptr || + !ggml_backend_buffer_is_cuda(job->dst->buffer) || job->dst->buffer->context == nullptr || + job->tiles == nullptr || job->tile_count == 0 || job->tile_count > UINT32_MAX || + job->cells == nullptr || job->cell_count == 0 || job->cell_count > UINT32_MAX || + job->cell_count > SIZE_MAX / sizeof(uint32_t) || job->channels == 0 || + job->channels == UINT32_MAX || +- (job->dst->type != GGML_TYPE_F16 && job->dst->type != GGML_TYPE_F32) || ++ (job->dst->type != GGML_TYPE_F16 && job->dst->type != GGML_TYPE_F32 && ++ job->dst->type != GGML_TYPE_Q8_0 && job->dst->type != GGML_TYPE_Q4_0) || ++ (quantized && job->channels % 32 != 0) || + job->token_stride == 0 || job->channel_stride == 0 || + job->token_stride % ggml_type_size(job->dst->type) != 0 || + job->channel_stride % ggml_type_size(job->dst->type) != 0) { +@@ -5792,7 +5796,8 @@ static bool ggml_cuda_cachegen_decode( + for (uint32_t cell : staged.cells) { + max_cell = std::max(max_cell, static_cast(cell)); + } +- const uint64_t max_channel = static_cast(job->channels) - 1; ++ const uint64_t output_units = quantized ? job->channels / 32 : job->channels; ++ const uint64_t max_channel = output_units - 1; + if (max_cell > (UINT64_MAX - staged.element_bytes) / job->token_stride || + max_channel > (UINT64_MAX - max_cell * job->token_stride - staged.element_bytes) / + job->channel_stride || +diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m +index d4ae1f475..0092c558f 100644 +--- a/ggml/src/ggml-metal/ggml-metal-device.m ++++ b/ggml/src/ggml-metal/ggml-metal-device.m +@@ -2682,9 +2682,13 @@ bool ggml_metal_cachegen_decode( + NSMutableArray * temporary_buffers = [[NSMutableArray alloc] init]; + for (size_t job_index = 0; job_index < job_count; ++job_index) { + const struct ggml_backend_cachegen_job * job = &jobs[job_index]; ++ const bool quantized = job->dst != NULL && ++ (job->dst->type == GGML_TYPE_Q8_0 || job->dst->type == GGML_TYPE_Q4_0); + if (job->dst == NULL || job->dst->buffer == NULL || job->tiles == NULL || job->tile_count == 0 || + job->tile_count > UINT32_MAX || job->cells == NULL || job->cell_count == 0 || job->channels == 0 || +- (job->dst->type != GGML_TYPE_F16 && job->dst->type != GGML_TYPE_F32) || ++ (job->dst->type != GGML_TYPE_F16 && job->dst->type != GGML_TYPE_F32 && ++ job->dst->type != GGML_TYPE_Q8_0 && job->dst->type != GGML_TYPE_Q4_0) || ++ (quantized && job->channels % 32 != 0) || + job->token_stride == 0 || job->channel_stride == 0 || + job->token_stride % ggml_type_size(job->dst->type) != 0 || + job->channel_stride % ggml_type_size(job->dst->type) != 0 || +@@ -2696,7 +2700,8 @@ bool ggml_metal_cachegen_decode( + for (size_t cell_index = 0; cell_index < job->cell_count; ++cell_index) { + max_cell = MAX(max_cell, job->cells[cell_index]); + } +- const uint64_t max_channel = (uint64_t) job->channels - 1; ++ const uint64_t output_units = quantized ? job->channels / 32 : job->channels; ++ const uint64_t max_channel = output_units - 1; + const uint64_t element_bytes = ggml_type_size(job->dst->type); + if (max_cell > (UINT64_MAX - element_bytes) / job->token_stride || + max_channel > (UINT64_MAX - max_cell * job->token_stride - element_bytes) / job->channel_stride || +@@ -2851,8 +2856,9 @@ bool ggml_metal_cachegen_decode( + [encoder setBuffer:temporary_buffers[temporary_buffers.count - 1] offset:0 atIndex:3]; + [encoder setBuffer:dst.metal offset:dst.offs atIndex:4]; + [encoder setBytes:¶ms length:sizeof(params) atIndex:5]; +- [encoder dispatchThreadgroups:MTLSizeMake((job->channels - 1) / 64 + 1, job->tile_count, 1) +- threadsPerThreadgroup:MTLSizeMake(64, 1, 1)]; ++ const uint32_t threads_per_group = quantized ? 32 : 64; ++ [encoder dispatchThreadgroups:MTLSizeMake((job->channels - 1) / threads_per_group + 1, job->tile_count, 1) ++ threadsPerThreadgroup:MTLSizeMake(threads_per_group, 1, 1)]; + } + [encoder endEncoding]; + [cmd_buf commit]; +diff --git a/ggml/src/ggml-metal/kernels/cachegen.metal b/ggml/src/ggml-metal/kernels/cachegen.metal +index e426f7993..18c0a0fb4 100644 +--- a/ggml/src/ggml-metal/kernels/cachegen.metal ++++ b/ggml/src/ggml-metal/kernels/cachegen.metal +@@ -59,7 +59,9 @@ kernel void kernel_cachegen_decode(const device uchar * payload [[buffer + const device uint * cells [[buffer(3)]], + device uchar * dst [[buffer(4)]], + constant cachegen_params & params [[buffer(5)]], +- uint2 gid [[thread_position_in_grid]]) { ++ uint2 gid [[thread_position_in_grid]], ++ uint tid [[thread_index_in_threadgroup]]) { ++ threadgroup float quant_values[32]; + const uint channel = gid.x; + const uint tile_index = gid.y; + if (channel >= params.channels || tile_index >= params.tile_count) { +@@ -97,11 +99,45 @@ kernel void kernel_cachegen_decode(const device uchar * payload [[buffer + + const float normalized = (float(symbol) - center) / center; + const half decoded = half(normalized * maxes[row]); +- const ulong destination = +- ulong(cells[tile.token_offset + row]) * params.token_stride + ulong(channel) * params.channel_stride; +- if (params.element_bytes == sizeof(float)) { ++ const float decoded_f = float(decoded); ++ const ulong row_offset = ulong(cells[tile.token_offset + row]) * params.token_stride; ++ if (params.element_bytes == 34) { ++ const uint lane = channel & 31; ++ const ulong destination = row_offset + ulong(channel >> 5) * params.channel_stride; ++ const float amax = simd_max(abs(decoded_f)); ++ const float scale = amax / 127.0f; ++ if (lane == 0) { ++ *reinterpret_cast(dst + destination) = half(scale); ++ } ++ const int quantized = int(round(decoded_f * (scale == 0.0f ? 0.0f : 1.0f / scale))); ++ *(reinterpret_cast(dst + destination + 2) + lane) = ++ char(clamp(quantized, -127, 127)); ++ } else if (params.element_bytes == 18) { ++ const uint lane = channel & 31; ++ const ulong destination = row_offset + ulong(channel >> 5) * params.channel_stride; ++ const float amax = simd_max(abs(decoded_f)); ++ const uint winner = simd_min(select(32u, lane, abs(decoded_f) == amax)); ++ const float signed_max = simd_sum(select(0.0f, decoded_f, lane == winner)); ++ const float scale = signed_max / -8.0f; ++ quant_values[tid] = decoded_f; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (lane == 0) { ++ *reinterpret_cast(dst + destination) = half(scale); ++ const float inverse = scale == 0.0f ? 0.0f : 1.0f / scale; ++ for (uint index = 0; index < 16; ++index) { ++ const int low = clamp( ++ int(trunc(quant_values[index] * inverse + 8.5f)), 0, 15); ++ const int high = clamp( ++ int(trunc(quant_values[index + 16] * inverse + 8.5f)), 0, 15); ++ *(dst + destination + 2 + index) = uchar(low | (high << 4)); ++ } ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } else if (params.element_bytes == sizeof(float)) { ++ const ulong destination = row_offset + ulong(channel) * params.channel_stride; + *reinterpret_cast(dst + destination) = float(decoded); + } else { ++ const ulong destination = row_offset + ulong(channel) * params.channel_stride; + *reinterpret_cast(dst + destination) = decoded; + } + +diff --git a/include/skippy/state.h b/include/skippy/state.h +index 9c78b3698..981da6b5d 100644 +--- a/include/skippy/state.h ++++ b/include/skippy/state.h +@@ -35,6 +35,8 @@ enum skippy_cachegen_record_kind { + SKIPPY_CACHEGEN_RECORD_F16_TRANSPOSED = 2, + SKIPPY_CACHEGEN_RECORD_F32 = 3, + SKIPPY_CACHEGEN_RECORD_F32_TRANSPOSED = 4, ++ SKIPPY_CACHEGEN_RECORD_Q8_0 = 5, ++ SKIPPY_CACHEGEN_RECORD_Q4_0 = 6, + }; + + /** @brief Describes one validated CacheGen record and its logical page destination. */ +diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp +index 0671fb296..215076924 100644 +--- a/src/llama-kv-cache.cpp ++++ b/src/llama-kv-cache.cpp +@@ -1971,6 +1971,7 @@ static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & r + size_t expected_rows, + size_t expected_channels, + size_t output_element_bytes, ++ size_t output_row_bytes, + std::string & error) { + constexpr size_t header_bytes = 16; + constexpr size_t cdf_entries = 33; +@@ -1991,8 +1992,8 @@ static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & r + return false; + } + if (expected_rows > std::numeric_limits::max() / expected_channels || +- expected_rows * expected_channels > std::numeric_limits::max() / output_element_bytes || +- record.decoded_bytes != expected_rows * expected_channels * output_element_bytes) { ++ expected_rows > std::numeric_limits::max() / output_row_bytes || ++ record.decoded_bytes != expected_rows * output_row_bytes) { + error = "CacheGen record decoded geometry is inconsistent"; + return false; + } +@@ -2088,11 +2089,11 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id + return false; + } + const auto cachegen_type_supported = [](uint32_t type) { +- return type == GGML_TYPE_F16 || type == GGML_TYPE_F32; ++ return type == GGML_TYPE_F16 || type == GGML_TYPE_F32 || type == GGML_TYPE_Q8_0 || type == GGML_TYPE_Q4_0; + }; + if (records == nullptr || record_count == 0 || !cachegen_type_supported(desc.k_type) || + !cachegen_type_supported(desc.v_type)) { +- error = "CacheGen import requires non-empty F16 or F32 K/V records"; ++ error = "CacheGen import requires non-empty supported K/V records"; + return false; + } + +@@ -2109,17 +2110,25 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id + size_t logical_offset = output_base; + std::vector storage; + storage.reserve(selected.size() * 2); +- auto consume_tiles = [&](ggml_tensor * dst, size_t row_bytes, size_t element_bytes, bool transposed) -> bool { +- if (element_bytes == 0 || row_bytes % element_bytes != 0 || +- (element_bytes != sizeof(uint16_t) && element_bytes != sizeof(float))) { ++ auto consume_tiles = [&](ggml_tensor * dst, size_t row_bytes, bool transposed) -> bool { ++ const size_t element_bytes = ggml_type_size(dst->type); ++ const size_t block_values = ggml_blck_size(dst->type); ++ const bool quantized = dst->type == GGML_TYPE_Q8_0 || dst->type == GGML_TYPE_Q4_0; ++ if (element_bytes == 0 || block_values == 0 || row_bytes == 0 || row_bytes % element_bytes != 0 || ++ (!quantized && element_bytes != sizeof(uint16_t) && element_bytes != sizeof(float)) || ++ (quantized && (transposed || block_values != 32))) { + error = "CacheGen destination element geometry is invalid"; + return false; + } ++ const size_t channels = row_bytes / element_bytes * block_values; ++ if (channels == 0 || channels > UINT32_MAX) { ++ error = "CacheGen destination channel count is invalid"; ++ return false; ++ } + skippy_cachegen_job_storage entry; + entry.dst = dst; + entry.token_stride = transposed ? element_bytes : row_bytes; + entry.channel_stride = transposed ? static_cast(v_cells[strm].size()) * element_bytes : element_bytes; +- const size_t channels = row_bytes / element_bytes; + entry.channels = static_cast(channels); + size_t token_offset = 0; + while (token_offset < desc.token_count) { +@@ -2129,14 +2138,28 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id + } + const auto & record = records[record_index++]; + const size_t rows = static_cast(record.token_count); +- const uint32_t expected_kind = element_bytes == sizeof(float) ? +- (transposed ? SKIPPY_CACHEGEN_RECORD_F32_TRANSPOSED : SKIPPY_CACHEGEN_RECORD_F32) : +- (transposed ? SKIPPY_CACHEGEN_RECORD_F16_TRANSPOSED : SKIPPY_CACHEGEN_RECORD_F16); ++ uint32_t expected_kind = SKIPPY_CACHEGEN_RECORD_F16; ++ switch (dst->type) { ++ case GGML_TYPE_F16: ++ expected_kind = transposed ? SKIPPY_CACHEGEN_RECORD_F16_TRANSPOSED : SKIPPY_CACHEGEN_RECORD_F16; ++ break; ++ case GGML_TYPE_F32: ++ expected_kind = transposed ? SKIPPY_CACHEGEN_RECORD_F32_TRANSPOSED : SKIPPY_CACHEGEN_RECORD_F32; ++ break; ++ case GGML_TYPE_Q8_0: ++ expected_kind = SKIPPY_CACHEGEN_RECORD_Q8_0; ++ break; ++ case GGML_TYPE_Q4_0: ++ expected_kind = SKIPPY_CACHEGEN_RECORD_Q4_0; ++ break; ++ default: ++ GGML_ABORT("unreachable CacheGen destination type"); ++ } + if (record.kind != expected_kind || rows == 0 || rows > desc.token_count - token_offset || + record.output_offset != logical_offset + (transposed ? 0 : token_offset * row_bytes) || + record.token_start != (transposed ? token_offset : 0) || + record.total_tokens != (transposed ? desc.token_count : 0) || +- !skippy_cachegen_validate_segment(record, rows, channels, element_bytes, error)) { ++ !skippy_cachegen_validate_segment(record, rows, channels, element_bytes, row_bytes, error)) { + if (error.empty()) { + error = "CacheGen record order or destination geometry is invalid"; + } +@@ -2152,14 +2175,14 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id + }; + + for (const auto * layer : selected) { +- if (!consume_tiles(layer->k_stream[strm], desc.k_row_bytes, ggml_type_size(layer->k_stream[strm]->type), false)) { ++ if (!consume_tiles(layer->k_stream[strm], desc.k_row_bytes, false)) { + return false; + } + } + for (const auto * layer : selected) { + const size_t v_row_bytes = + v_trans ? static_cast(hparams.n_embd_v_gqa(layer->il)) * desc.v_element_bytes : desc.v_row_bytes; +- if (!consume_tiles(layer->v_stream[strm], v_row_bytes, ggml_type_size(layer->v_stream[strm]->type), v_trans)) { ++ if (!consume_tiles(layer->v_stream[strm], v_row_bytes, v_trans)) { + return false; + } + } +diff --git a/tests/test-skippy-cachegen-metal.cpp b/tests/test-skippy-cachegen-metal.cpp +index dac796d25..b7caf2adc 100644 +--- a/tests/test-skippy-cachegen-metal.cpp ++++ b/tests/test-skippy-cachegen-metal.cpp +@@ -6,6 +6,53 @@ + #include + #include + ++static std::vector repeat_cachegen_channels( ++ const uint8_t * payload, size_t payload_bytes, uint32_t repeats) { ++ const uint32_t rows = (uint32_t) payload[6] | (uint32_t) payload[7] << 8; ++ const uint32_t channels = (uint32_t) payload[8] | (uint32_t) payload[9] << 8 | ++ (uint32_t) payload[10] << 16 | (uint32_t) payload[11] << 24; ++ const size_t cdf_offset = 16 + (size_t) rows * sizeof(float); ++ const size_t cdf_bytes = (size_t) channels * 33 * sizeof(uint16_t); ++ const size_t lengths_offset = cdf_offset + cdf_bytes; ++ const size_t lengths_bytes = (size_t) channels * sizeof(uint16_t); ++ const size_t streams_offset = lengths_offset + lengths_bytes; ++ std::vector stream_offsets(channels + 1); ++ for (uint32_t channel = 0; channel < channels; ++channel) { ++ const size_t offset = lengths_offset + (size_t) channel * sizeof(uint16_t); ++ stream_offsets[channel + 1] = stream_offsets[channel] + ++ ((uint32_t) payload[offset] | (uint32_t) payload[offset + 1] << 8); ++ } ++ std::vector result(payload, payload + cdf_offset); ++ for (uint32_t repeat = 0; repeat < repeats; ++repeat) { ++ for (uint32_t channel = 0; channel < channels; ++channel) { ++ const uint32_t source = (channel + repeat) % channels; ++ result.insert(result.end(), payload + cdf_offset + (size_t) source * 33 * sizeof(uint16_t), ++ payload + cdf_offset + (size_t) (source + 1) * 33 * sizeof(uint16_t)); ++ } ++ } ++ for (uint32_t repeat = 0; repeat < repeats; ++repeat) { ++ for (uint32_t channel = 0; channel < channels; ++channel) { ++ const uint32_t source = (channel + repeat) % channels; ++ result.insert(result.end(), payload + lengths_offset + (size_t) source * sizeof(uint16_t), ++ payload + lengths_offset + (size_t) (source + 1) * sizeof(uint16_t)); ++ } ++ } ++ for (uint32_t repeat = 0; repeat < repeats; ++repeat) { ++ for (uint32_t channel = 0; channel < channels; ++channel) { ++ const uint32_t source = (channel + repeat) % channels; ++ result.insert(result.end(), payload + streams_offset + stream_offsets[source], ++ payload + streams_offset + stream_offsets[source + 1]); ++ } ++ } ++ const uint32_t output_channels = channels * repeats; ++ const uint32_t output_stream_bytes = (uint32_t) (payload_bytes - streams_offset) * repeats; ++ for (uint32_t byte = 0; byte < 4; ++byte) { ++ result[8 + byte] = uint8_t(output_channels >> (byte * 8)); ++ result[12 + byte] = uint8_t(output_stream_bytes >> (byte * 8)); ++ } ++ return result; ++} ++ + static const uint8_t encoded16[] = { + 0x4c, 0x43, 0x47, 0x31, 0x10, 0x00, 0x11, 0x00, 0x08, 0x00, 0x00, 0x00, 0x2b, 0x00, 0x00, 0x00, 0x00, 0xa0, 0x66, + 0x3e, 0x00, 0x80, 0xcd, 0x3e, 0x00, 0x60, 0x0e, 0x3f, 0x00, 0x40, 0x2e, 0x3f, 0x00, 0x40, 0x44, 0x3f, 0x00, 0x00, +@@ -138,7 +185,7 @@ int main() { + if (backend == nullptr) { + return 1; + } +- ggml_init_params params = { ggml_tensor_overhead() * 5, nullptr, true }; ++ ggml_init_params params = { ggml_tensor_overhead() * 8, nullptr, true }; + ggml_context * ctx = ggml_init(params); + if (ctx == nullptr) { + ggml_backend_free(backend); +@@ -148,6 +195,10 @@ int main() { + ggml_tensor * transposed16 = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, capacity, channels); + ggml_tensor * row_major32 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, channels, capacity); + ggml_tensor * transposed32 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, capacity, channels); ++ constexpr uint32_t quant_channels = 32; ++ ggml_tensor * row_major_quant_f16 = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, quant_channels, capacity); ++ ggml_tensor * row_major_q8 = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, quant_channels, capacity); ++ ggml_tensor * row_major_q4 = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_0, quant_channels, capacity); + constexpr uint32_t multi_capacity = rows * 2 + 5; + ggml_tensor * multi_tile = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, channels, multi_capacity); + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); +@@ -169,6 +220,8 @@ int main() { + } + const ggml_backend_cachegen_tile tile16 = { encoded16, sizeof(encoded16), 0, rows }; + const ggml_backend_cachegen_tile tile32 = { encoded32, sizeof(encoded32), 0, rows }; ++ const std::vector encoded_quant = repeat_cachegen_channels(encoded16, sizeof(encoded16), 4); ++ const ggml_backend_cachegen_tile tile_quant = { encoded_quant.data(), encoded_quant.size(), 0, rows }; + const ggml_backend_cachegen_tile multi_tiles[] = { + { encoded16, sizeof(encoded16), 0, rows }, + { encoded32, sizeof(encoded32), rows, rows }, +@@ -183,6 +236,9 @@ int main() { + { row_major32, &tile32, 1, cells.data(), cells.size(), channels, channels * 4, 4 }, + { transposed32, &tile32, 1, cells.data(), cells.size(), channels, 4, capacity * 4 }, + { multi_tile, multi_tiles, 2, multi_cells.data(), multi_cells.size(), channels, channels * 2, 2 }, ++ { row_major_quant_f16, &tile_quant, 1, cells.data(), cells.size(), quant_channels, quant_channels * 2, 2 }, ++ { row_major_q8, &tile_quant, 1, cells.data(), cells.size(), quant_channels, 34, 34 }, ++ { row_major_q4, &tile_quant, 1, cells.data(), cells.size(), quant_channels, 18, 18 }, + }; + ggml_backend_dev_t device = ggml_backend_get_device(backend); + auto decode = reinterpret_cast( +@@ -192,7 +248,7 @@ int main() { + return finish(4); + } + char error[256] = {}; +- if (!decode(jobs, 5, error, sizeof(error))) { ++ if (!decode(jobs, 8, error, sizeof(error))) { + return finish(5); + } + ggml_backend_cachegen_job invalid_job = jobs[0]; +@@ -261,5 +317,47 @@ int main() { + } + } + } ++ std::vector quant_source(rows * quant_channels); ++ for (uint32_t row = 0; row < rows; ++row) { ++ for (uint32_t channel = 0; channel < quant_channels; ++channel) { ++ uint16_t bits; ++ const uint32_t source_channel = (channel % channels + channel / channels) % channels; ++ std::memcpy(&bits, expected16 + (row * channels + source_channel) * sizeof(bits), sizeof(bits)); ++ quant_source[row * quant_channels + channel] = ggml_fp16_to_fp32(bits); ++ } ++ } ++ std::vector quant_f16_bytes(ggml_nbytes(row_major_quant_f16)); ++ ggml_backend_tensor_get(row_major_quant_f16, quant_f16_bytes.data(), 0, quant_f16_bytes.size()); ++ for (uint32_t row = 0; row < rows; ++row) { ++ for (uint32_t channel = 0; channel < quant_channels; ++channel) { ++ if (std::memcmp(quant_f16_bytes.data() + ((size_t) cells[row] * quant_channels + channel) * 2, ++ expected16 + ((size_t) row * channels + ++ (channel % channels + channel / channels) % channels) * 2, 2) != 0) { ++ return finish(12); ++ } ++ } ++ } ++ const auto check_quantized = [&](ggml_tensor * tensor, ggml_type type, size_t row_bytes, int status) { ++ std::vector expected(rows * row_bytes); ++ if (ggml_quantize_chunk(type, quant_source.data(), expected.data(), 0, rows, quant_channels, nullptr) != ++ expected.size()) { ++ return status; ++ } ++ std::vector actual(ggml_nbytes(tensor)); ++ ggml_backend_tensor_get(tensor, actual.data(), 0, actual.size()); ++ for (uint32_t row = 0; row < rows; ++row) { ++ if (std::memcmp(actual.data() + (size_t) cells[row] * row_bytes, ++ expected.data() + (size_t) row * row_bytes, row_bytes) != 0) { ++ return status; ++ } ++ } ++ return 0; ++ }; ++ if (const int status = check_quantized(row_major_q8, GGML_TYPE_Q8_0, 34, 10)) { ++ return finish(status); ++ } ++ if (const int status = check_quantized(row_major_q4, GGML_TYPE_Q4_0, 18, 11)) { ++ return finish(status); ++ } + return finish(0); + } +-- +2.54.0 (Apple Git-157) + diff --git a/third_party/llama.cpp/patches/0043-ggml-metal-stage-CacheGen-directly.patch b/third_party/llama.cpp/patches/0043-ggml-metal-stage-CacheGen-directly.patch new file mode 100644 index 0000000000..1d110dd6e4 --- /dev/null +++ b/third_party/llama.cpp/patches/0043-ggml-metal-stage-CacheGen-directly.patch @@ -0,0 +1,159 @@ +From: scama +Date: Fri, 12 Sep 2026 01:00:00 +1000 +Subject: [PATCH] ggml-metal: stage CacheGen payloads directly + +Avoid building each CacheGen job in NSMutableData and then copying the same +bytes again into shared Metal buffers. Validate and size the job first, then +fill its final MTLBuffers directly. +--- + ggml/src/ggml-metal/ggml-metal-device.m | 85 ++++++++++++------------ + 1 file changed, 43 insertions(+), 42 deletions(-) + +diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m +index 0092c558f..79d40466f 100644 +--- a/ggml/src/ggml-metal/ggml-metal-device.m ++++ b/ggml/src/ggml-metal/ggml-metal-device.m +@@ -2724,92 +2724,63 @@ bool ggml_metal_cachegen_decode( + } + } + +- NSMutableData * payload_data = [[NSMutableData alloc] init]; +- NSMutableData * tile_data = [[NSMutableData alloc] init]; +- NSMutableData * prefix_data = [[NSMutableData alloc] init]; ++ size_t staged_payload_bytes = 0; ++ size_t staged_prefix_count = 0; + for (size_t tile_index = 0; tile_index < job->tile_count; ++tile_index) { + const struct ggml_backend_cachegen_tile * tile = &job->tiles[tile_index]; + if (tile->payload == NULL || tile->payload_bytes < 16 || tile->token_count == 0 || + tile->token_offset > job->cell_count || tile->token_count > job->cell_count - tile->token_offset || +- payload_data.length > UINT32_MAX - 3 || +- prefix_data.length / sizeof(uint32_t) > UINT32_MAX - ((size_t) job->channels + 1)) { +- [payload_data release]; +- [tile_data release]; +- [prefix_data release]; ++ staged_payload_bytes > UINT32_MAX - 3 || ++ staged_prefix_count > UINT32_MAX - ((size_t) job->channels + 1)) { + [encoder endEncoding]; + [temporary_buffers release]; + return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen tile exceeds bounded geometry"); + } +- const size_t aligned_payload_size = (payload_data.length + 3) & ~(size_t) 3; ++ const size_t aligned_payload_size = (staged_payload_bytes + 3) & ~(size_t) 3; + if (tile->payload_bytes > UINT32_MAX - aligned_payload_size) { +- [payload_data release]; +- [tile_data release]; +- [prefix_data release]; + [encoder endEncoding]; + [temporary_buffers release]; + return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen payload offsets overflow"); + } +- const uint8_t padding[3] = { 0 }; +- [payload_data appendBytes:padding length:aligned_payload_size - payload_data.length]; + const uint8_t * payload = (const uint8_t *) tile->payload; + const size_t lengths_offset = 16 + (size_t) tile->token_count * sizeof(float) + + (size_t) job->channels * 33 * sizeof(uint16_t); + const size_t streams_offset = lengths_offset + (size_t) job->channels * sizeof(uint16_t); + if (streams_offset > tile->payload_bytes) { +- [payload_data release]; +- [tile_data release]; +- [prefix_data release]; + [encoder endEncoding]; + [temporary_buffers release]; + return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen tile metadata is truncated"); + } +- struct ggml_metal_cachegen_tile encoded_tile = { +- (uint32_t) aligned_payload_size, +- (uint32_t) (prefix_data.length / sizeof(uint32_t)), +- tile->token_offset, +- tile->token_count, +- }; + uint32_t prefix = 0; +- [prefix_data appendBytes:&prefix length:sizeof(prefix)]; + for (uint32_t channel = 0; channel < job->channels; ++channel) { + const size_t offset = lengths_offset + (size_t) channel * sizeof(uint16_t); + const uint32_t length = (uint32_t) payload[offset] | (uint32_t) payload[offset + 1] << 8; + if (prefix > UINT32_MAX - length) { +- [payload_data release]; +- [tile_data release]; +- [prefix_data release]; + [encoder endEncoding]; + [temporary_buffers release]; + return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen stream offsets overflow"); + } + prefix += length; +- [prefix_data appendBytes:&prefix length:sizeof(prefix)]; + } + const uint32_t declared_stream_bytes = (uint32_t) payload[12] | + (uint32_t) payload[13] << 8 | (uint32_t) payload[14] << 16 | (uint32_t) payload[15] << 24; + if (prefix != declared_stream_bytes || prefix != tile->payload_bytes - streams_offset) { +- [payload_data release]; +- [tile_data release]; +- [prefix_data release]; + [encoder endEncoding]; + [temporary_buffers release]; + return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen stream length is inconsistent"); + } +- [tile_data appendBytes:&encoded_tile length:sizeof(encoded_tile)]; +- [payload_data appendBytes:tile->payload length:tile->payload_bytes]; ++ staged_payload_bytes = aligned_payload_size + tile->payload_bytes; ++ staged_prefix_count += (size_t) job->channels + 1; + } + +- id payload_buffer = [device->mtl_device newBufferWithBytes:payload_data.bytes +- length:payload_data.length options:MTLResourceStorageModeShared]; +- id tile_buffer = [device->mtl_device newBufferWithBytes:tile_data.bytes +- length:tile_data.length options:MTLResourceStorageModeShared]; +- id prefix_buffer = [device->mtl_device newBufferWithBytes:prefix_data.bytes +- length:prefix_data.length options:MTLResourceStorageModeShared]; ++ id payload_buffer = [device->mtl_device newBufferWithLength:staged_payload_bytes ++ options:MTLResourceStorageModeShared]; ++ id tile_buffer = [device->mtl_device newBufferWithLength:job->tile_count * sizeof(struct ggml_metal_cachegen_tile) ++ options:MTLResourceStorageModeShared]; ++ id prefix_buffer = [device->mtl_device newBufferWithLength:staged_prefix_count * sizeof(uint32_t) ++ options:MTLResourceStorageModeShared]; + id cell_buffer = [device->mtl_device newBufferWithBytes:job->cells + length:job->cell_count * sizeof(uint32_t) options:MTLResourceStorageModeShared]; +- [payload_data release]; +- [tile_data release]; +- [prefix_data release]; + if (payload_buffer == nil || tile_buffer == nil || prefix_buffer == nil || cell_buffer == nil) { + [payload_buffer release]; + [tile_buffer release]; +@@ -2819,6 +2790,35 @@ bool ggml_metal_cachegen_decode( + [temporary_buffers release]; + return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen staging allocation failed"); + } ++ uint8_t * payload_contents = (uint8_t *) payload_buffer.contents; ++ struct ggml_metal_cachegen_tile * tile_contents = ++ (struct ggml_metal_cachegen_tile *) tile_buffer.contents; ++ uint32_t * prefix_contents = (uint32_t *) prefix_buffer.contents; ++ size_t payload_offset = 0; ++ size_t prefix_offset = 0; ++ for (size_t tile_index = 0; tile_index < job->tile_count; ++tile_index) { ++ const struct ggml_backend_cachegen_tile * tile = &job->tiles[tile_index]; ++ const size_t aligned_payload_offset = (payload_offset + 3) & ~(size_t) 3; ++ memset(payload_contents + payload_offset, 0, aligned_payload_offset - payload_offset); ++ memcpy(payload_contents + aligned_payload_offset, tile->payload, tile->payload_bytes); ++ tile_contents[tile_index] = (struct ggml_metal_cachegen_tile) { ++ (uint32_t) aligned_payload_offset, ++ (uint32_t) prefix_offset, ++ tile->token_offset, ++ tile->token_count, ++ }; ++ const uint8_t * payload = (const uint8_t *) tile->payload; ++ const size_t lengths_offset = 16 + (size_t) tile->token_count * sizeof(float) + ++ (size_t) job->channels * 33 * sizeof(uint16_t); ++ uint32_t prefix = 0; ++ prefix_contents[prefix_offset++] = prefix; ++ for (uint32_t channel = 0; channel < job->channels; ++channel) { ++ const size_t offset = lengths_offset + (size_t) channel * sizeof(uint16_t); ++ prefix += (uint32_t) payload[offset] | (uint32_t) payload[offset + 1] << 8; ++ prefix_contents[prefix_offset++] = prefix; ++ } ++ payload_offset = aligned_payload_offset + tile->payload_bytes; ++ } + [temporary_buffers addObject:payload_buffer]; + [temporary_buffers addObject:tile_buffer]; + [temporary_buffers addObject:prefix_buffer]; +-- +2.52.0 diff --git a/third_party/llama.cpp/patches/0044-ggml-decode-packed-CacheGen-symbols-on-device.patch b/third_party/llama.cpp/patches/0044-ggml-decode-packed-CacheGen-symbols-on-device.patch new file mode 100644 index 0000000000..299468d419 --- /dev/null +++ b/third_party/llama.cpp/patches/0044-ggml-decode-packed-CacheGen-symbols-on-device.patch @@ -0,0 +1,581 @@ +From 9253c686456c6f90182cd49094c96ec1cb9291c3 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Sat, 12 Sep 2026 05:42:09 +1000 +Subject: [PATCH] ggml: decode packed CacheGen symbols on device + +--- + ggml/src/ggml-cuda/cachegen.cu | 52 ++++++++++---- + ggml/src/ggml-cuda/ggml-cuda.cu | 76 +++++++++++++------- + ggml/src/ggml-metal/ggml-metal-device.m | 82 +++++++++++++++------- + ggml/src/ggml-metal/kernels/cachegen.metal | 55 +++++++++++---- + src/llama-kv-cache.cpp | 45 +++++++++--- + tests/test-skippy-cachegen-metal.cpp | 40 ++++++++++- + 6 files changed, 261 insertions(+), 89 deletions(-) + +diff --git a/ggml/src/ggml-cuda/cachegen.cu b/ggml/src/ggml-cuda/cachegen.cu +index fa5683788..685dc586a 100644 +--- a/ggml/src/ggml-cuda/cachegen.cu ++++ b/ggml/src/ggml-cuda/cachegen.cu +@@ -63,28 +63,53 @@ static __global__ void cachegen_decode( + const ggml_cuda_cachegen_tile tile = tiles[tile_index]; + const uint8_t * segment = payload + tile.payload_offset; + const uint32_t bins = segment[4]; ++ const bool packed = segment[3] == '2'; + const float * maxes = reinterpret_cast(segment + 16); +- const uint16_t * cdfs = reinterpret_cast(segment + 16 + tile.rows * sizeof(float)); +- const uint16_t * cdf = cdfs + channel * 33; +- const uint32_t prefix_index = tile.prefix_offset + channel; +- const uint32_t stream_start = prefixes[prefix_index]; +- const uint32_t stream_bytes = prefixes[prefix_index + 1] - stream_start; +- const uint32_t lengths_offset = 16 + tile.rows * sizeof(float) + channels * 33 * sizeof(uint16_t); +- const uint32_t streams_offset = lengths_offset + channels * sizeof(uint16_t); +- const uint8_t * stream = segment + streams_offset + stream_start; ++ const uint16_t * cdf = nullptr; ++ const uint8_t * stream = nullptr; ++ uint32_t stream_bytes = 0; ++ if (packed) { ++ stream = segment + 16 + tile.rows * sizeof(float); ++ stream_bytes = uint32_t(segment[12]) | uint32_t(segment[13]) << 8 | ++ uint32_t(segment[14]) << 16 | uint32_t(segment[15]) << 24; ++ } else { ++ const uint16_t * cdfs = reinterpret_cast(segment + 16 + tile.rows * sizeof(float)); ++ cdf = cdfs + channel * 33; ++ const uint32_t prefix_index = tile.prefix_offset + channel; ++ const uint32_t stream_start = prefixes[prefix_index]; ++ stream_bytes = prefixes[prefix_index + 1] - stream_start; ++ const uint32_t lengths_offset = 16 + tile.rows * sizeof(float) + channels * 33 * sizeof(uint16_t); ++ const uint32_t streams_offset = lengths_offset + channels * sizeof(uint16_t); ++ stream = segment + streams_offset + stream_start; ++ } + + uint32_t bit = 0; + uint32_t value = 0; +- for (uint32_t i = 0; i < 32; ++i) { +- value = (value << 1) | cachegen_read_bit(stream, stream_bytes, bit); ++ if (!packed) { ++ for (uint32_t i = 0; i < 32; ++i) { ++ value = (value << 1) | cachegen_read_bit(stream, stream_bytes, bit); ++ } + } + uint32_t low = 0; + uint32_t high = 0xffffffffu; + const float center = float(bins / 2 - 1); + for (uint32_t row = 0; row < tile.rows; ++row) { + const uint64_t span = uint64_t(high) - uint64_t(low) + 1; +- const uint16_t count = cachegen_scaled_count(value, low, span); +- const uint32_t symbol = cachegen_find_symbol(cdf, count); ++ uint32_t symbol; ++ if (packed) { ++ const uint32_t bits_per_symbol = segment[5]; ++ const uint64_t symbol_bit = (uint64_t(row) * channels + channel) * bits_per_symbol; ++ const uint64_t byte_index = symbol_bit >> 3; ++ const uint32_t shift = symbol_bit & 7; ++ uint32_t word = stream[byte_index]; ++ if (shift + bits_per_symbol > 8) { ++ word |= uint32_t(stream[byte_index + 1]) << 8; ++ } ++ symbol = (word >> shift) & ((1u << bits_per_symbol) - 1); ++ } else { ++ const uint16_t count = cachegen_scaled_count(value, low, span); ++ symbol = cachegen_find_symbol(cdf, count); ++ } + if (symbol >= 32) { + return; + } +@@ -143,6 +168,9 @@ static __global__ void cachegen_decode( + if (row + 1 == tile.rows) { + break; + } ++ if (packed) { ++ continue; ++ } + const uint64_t cdf_low = cdf[symbol]; + const uint64_t cdf_high = symbol == 31 ? 65536ull : cdf[symbol + 1]; + high = low - 1 + uint32_t((span * cdf_high) >> 16); +diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu +index b68649ec4..ccfda8a9d 100644 +--- a/ggml/src/ggml-cuda/ggml-cuda.cu ++++ b/ggml/src/ggml-cuda/ggml-cuda.cu +@@ -5814,7 +5814,7 @@ static bool ggml_cuda_cachegen_decode( + for (size_t tile_index = 0; tile_index < job->tile_count; ++tile_index) { + const struct ggml_backend_cachegen_tile * tile = &job->tiles[tile_index]; + if (tile->payload == nullptr || tile->payload_bytes < 16 || tile->payload_bytes > UINT32_MAX || +- tile->token_count == 0 || tile->token_offset > job->cell_count || ++ tile->token_count == 0 || tile->token_count > 256 || tile->token_offset > job->cell_count || + tile->token_count > job->cell_count - tile->token_offset || + staged.payload.size() > UINT32_MAX - 3 || + staged.prefixes.size() > UINT32_MAX - (static_cast(job->channels) + 1)) { +@@ -5830,41 +5830,67 @@ static bool ggml_cuda_cachegen_decode( + staged.payload.resize(aligned_payload_size, 0); + + const uint8_t * payload = static_cast(tile->payload); +- const uint64_t streams_offset_u64 = 16ull + static_cast(tile->token_count) * sizeof(float) + +- static_cast(job->channels) * 33 * sizeof(uint16_t) + +- static_cast(job->channels) * sizeof(uint16_t); +- if (streams_offset_u64 > tile->payload_bytes || streams_offset_u64 > UINT32_MAX) { ++ const bool packed = memcmp(payload, "LCG2", 4) == 0; ++ const uint32_t bins = payload[4]; ++ const uint32_t expected_bits = bins == 16 ? 4 : 5; ++ const uint32_t header_rows = static_cast(payload[6]) | ++ static_cast(payload[7]) << 8; ++ const uint32_t header_channels = static_cast(payload[8]) | ++ static_cast(payload[9]) << 8 | static_cast(payload[10]) << 16 | ++ static_cast(payload[11]) << 24; ++ if ((!packed && memcmp(payload, "LCG1", 4) != 0) || (bins != 16 && bins != 32) || ++ (packed ? payload[5] != expected_bits : payload[5] != 0) || ++ header_rows != tile->token_count || header_channels != job->channels) { + return ggml_cuda_cachegen_error( +- error, error_capacity, "CUDA/HIP CacheGen tile metadata is truncated"); ++ error, error_capacity, "CUDA/HIP CacheGen tile header is inconsistent"); + } +- const size_t lengths_offset = static_cast(streams_offset_u64) - +- static_cast(job->channels) * sizeof(uint16_t); +- const size_t streams_offset = static_cast(streams_offset_u64); + const ggml_cuda_cachegen_tile encoded_tile = { + static_cast(staged.payload.size()), + static_cast(staged.prefixes.size()), + tile->token_offset, + tile->token_count, + }; +- uint32_t prefix = 0; +- staged.prefixes.push_back(prefix); +- for (uint32_t channel = 0; channel < job->channels; ++channel) { +- const size_t offset = lengths_offset + static_cast(channel) * sizeof(uint16_t); +- const uint32_t length = static_cast(payload[offset]) | +- static_cast(payload[offset + 1]) << 8; +- if (prefix > UINT32_MAX - length) { +- return ggml_cuda_cachegen_error( +- error, error_capacity, "CUDA/HIP CacheGen stream offsets overflow"); +- } +- prefix += length; +- staged.prefixes.push_back(prefix); +- } + const uint32_t declared_stream_bytes = static_cast(payload[12]) | + static_cast(payload[13]) << 8 | static_cast(payload[14]) << 16 | + static_cast(payload[15]) << 24; +- if (prefix != declared_stream_bytes || prefix != tile->payload_bytes - streams_offset) { +- return ggml_cuda_cachegen_error( +- error, error_capacity, "CUDA/HIP CacheGen stream length is inconsistent"); ++ staged.prefixes.push_back(0); ++ if (packed) { ++ const uint64_t values = static_cast(tile->token_count) * job->channels; ++ const uint64_t expected_stream_bytes = (values * expected_bits + 7) / 8; ++ const size_t streams_offset = 16 + static_cast(tile->token_count) * sizeof(float); ++ if (expected_stream_bytes > UINT32_MAX || streams_offset > tile->payload_bytes || ++ declared_stream_bytes != expected_stream_bytes || ++ declared_stream_bytes != tile->payload_bytes - streams_offset) { ++ return ggml_cuda_cachegen_error( ++ error, error_capacity, "CUDA/HIP CacheGen packed length is inconsistent"); ++ } ++ } else { ++ const uint64_t streams_offset_u64 = 16ull + static_cast(tile->token_count) * sizeof(float) + ++ static_cast(job->channels) * 33 * sizeof(uint16_t) + ++ static_cast(job->channels) * sizeof(uint16_t); ++ if (streams_offset_u64 > tile->payload_bytes || streams_offset_u64 > UINT32_MAX) { ++ return ggml_cuda_cachegen_error( ++ error, error_capacity, "CUDA/HIP CacheGen tile metadata is truncated"); ++ } ++ const size_t lengths_offset = static_cast(streams_offset_u64) - ++ static_cast(job->channels) * sizeof(uint16_t); ++ const size_t streams_offset = static_cast(streams_offset_u64); ++ uint32_t prefix = 0; ++ for (uint32_t channel = 0; channel < job->channels; ++channel) { ++ const size_t offset = lengths_offset + static_cast(channel) * sizeof(uint16_t); ++ const uint32_t length = static_cast(payload[offset]) | ++ static_cast(payload[offset + 1]) << 8; ++ if (prefix > UINT32_MAX - length) { ++ return ggml_cuda_cachegen_error( ++ error, error_capacity, "CUDA/HIP CacheGen stream offsets overflow"); ++ } ++ prefix += length; ++ staged.prefixes.push_back(prefix); ++ } ++ if (prefix != declared_stream_bytes || prefix != tile->payload_bytes - streams_offset) { ++ return ggml_cuda_cachegen_error( ++ error, error_capacity, "CUDA/HIP CacheGen stream length is inconsistent"); ++ } + } + staged.tiles.push_back(encoded_tile); + staged.payload.insert(staged.payload.end(), payload, payload + tile->payload_bytes); +diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m +index 79d40466f..2a38e0d4b 100644 +--- a/ggml/src/ggml-metal/ggml-metal-device.m ++++ b/ggml/src/ggml-metal/ggml-metal-device.m +@@ -2686,6 +2686,7 @@ bool ggml_metal_cachegen_decode( + (job->dst->type == GGML_TYPE_Q8_0 || job->dst->type == GGML_TYPE_Q4_0); + if (job->dst == NULL || job->dst->buffer == NULL || job->tiles == NULL || job->tile_count == 0 || + job->tile_count > UINT32_MAX || job->cells == NULL || job->cell_count == 0 || job->channels == 0 || ++ job->channels == UINT32_MAX || + (job->dst->type != GGML_TYPE_F16 && job->dst->type != GGML_TYPE_F32 && + job->dst->type != GGML_TYPE_Q8_0 && job->dst->type != GGML_TYPE_Q4_0) || + (quantized && job->channels % 32 != 0) || +@@ -2729,6 +2730,7 @@ bool ggml_metal_cachegen_decode( + for (size_t tile_index = 0; tile_index < job->tile_count; ++tile_index) { + const struct ggml_backend_cachegen_tile * tile = &job->tiles[tile_index]; + if (tile->payload == NULL || tile->payload_bytes < 16 || tile->token_count == 0 || ++ tile->token_count > 256 || + tile->token_offset > job->cell_count || tile->token_count > job->cell_count - tile->token_offset || + staged_payload_bytes > UINT32_MAX - 3 || + staged_prefix_count > UINT32_MAX - ((size_t) job->channels + 1)) { +@@ -2743,34 +2745,60 @@ bool ggml_metal_cachegen_decode( + return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen payload offsets overflow"); + } + const uint8_t * payload = (const uint8_t *) tile->payload; +- const size_t lengths_offset = 16 + (size_t) tile->token_count * sizeof(float) + +- (size_t) job->channels * 33 * sizeof(uint16_t); +- const size_t streams_offset = lengths_offset + (size_t) job->channels * sizeof(uint16_t); +- if (streams_offset > tile->payload_bytes) { ++ const bool packed = memcmp(payload, "LCG2", 4) == 0; ++ const uint32_t bins = payload[4]; ++ const uint32_t expected_bits = bins == 16 ? 4 : 5; ++ const uint32_t header_rows = (uint32_t) payload[6] | (uint32_t) payload[7] << 8; ++ const uint32_t header_channels = (uint32_t) payload[8] | (uint32_t) payload[9] << 8 | ++ (uint32_t) payload[10] << 16 | (uint32_t) payload[11] << 24; ++ if ((!packed && memcmp(payload, "LCG1", 4) != 0) || (bins != 16 && bins != 32) || ++ (packed ? payload[5] != expected_bits : payload[5] != 0) || ++ header_rows != tile->token_count || header_channels != job->channels) { + [encoder endEncoding]; + [temporary_buffers release]; +- return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen tile metadata is truncated"); ++ return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen tile header is inconsistent"); + } +- uint32_t prefix = 0; +- for (uint32_t channel = 0; channel < job->channels; ++channel) { +- const size_t offset = lengths_offset + (size_t) channel * sizeof(uint16_t); +- const uint32_t length = (uint32_t) payload[offset] | (uint32_t) payload[offset + 1] << 8; +- if (prefix > UINT32_MAX - length) { ++ const uint32_t declared_stream_bytes = (uint32_t) payload[12] | ++ (uint32_t) payload[13] << 8 | (uint32_t) payload[14] << 16 | (uint32_t) payload[15] << 24; ++ if (packed) { ++ const uint64_t values = (uint64_t) tile->token_count * job->channels; ++ const uint64_t expected_stream_bytes = (values * expected_bits + 7) / 8; ++ const size_t streams_offset = 16 + (size_t) tile->token_count * sizeof(float); ++ if (expected_stream_bytes > UINT32_MAX || streams_offset > tile->payload_bytes || ++ declared_stream_bytes != expected_stream_bytes || ++ declared_stream_bytes != tile->payload_bytes - streams_offset) { + [encoder endEncoding]; + [temporary_buffers release]; +- return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen stream offsets overflow"); ++ return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen packed length is inconsistent"); ++ } ++ } else { ++ const size_t lengths_offset = 16 + (size_t) tile->token_count * sizeof(float) + ++ (size_t) job->channels * 33 * sizeof(uint16_t); ++ const size_t streams_offset = lengths_offset + (size_t) job->channels * sizeof(uint16_t); ++ if (streams_offset > tile->payload_bytes) { ++ [encoder endEncoding]; ++ [temporary_buffers release]; ++ return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen tile metadata is truncated"); ++ } ++ uint32_t prefix = 0; ++ for (uint32_t channel = 0; channel < job->channels; ++channel) { ++ const size_t offset = lengths_offset + (size_t) channel * sizeof(uint16_t); ++ const uint32_t length = (uint32_t) payload[offset] | (uint32_t) payload[offset + 1] << 8; ++ if (prefix > UINT32_MAX - length) { ++ [encoder endEncoding]; ++ [temporary_buffers release]; ++ return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen stream offsets overflow"); ++ } ++ prefix += length; ++ } ++ if (prefix != declared_stream_bytes || prefix != tile->payload_bytes - streams_offset) { ++ [encoder endEncoding]; ++ [temporary_buffers release]; ++ return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen stream length is inconsistent"); + } +- prefix += length; +- } +- const uint32_t declared_stream_bytes = (uint32_t) payload[12] | +- (uint32_t) payload[13] << 8 | (uint32_t) payload[14] << 16 | (uint32_t) payload[15] << 24; +- if (prefix != declared_stream_bytes || prefix != tile->payload_bytes - streams_offset) { +- [encoder endEncoding]; +- [temporary_buffers release]; +- return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen stream length is inconsistent"); + } + staged_payload_bytes = aligned_payload_size + tile->payload_bytes; +- staged_prefix_count += (size_t) job->channels + 1; ++ staged_prefix_count += packed ? 1 : (size_t) job->channels + 1; + } + + id payload_buffer = [device->mtl_device newBufferWithLength:staged_payload_bytes +@@ -2808,14 +2836,16 @@ bool ggml_metal_cachegen_decode( + tile->token_count, + }; + const uint8_t * payload = (const uint8_t *) tile->payload; +- const size_t lengths_offset = 16 + (size_t) tile->token_count * sizeof(float) + +- (size_t) job->channels * 33 * sizeof(uint16_t); + uint32_t prefix = 0; + prefix_contents[prefix_offset++] = prefix; +- for (uint32_t channel = 0; channel < job->channels; ++channel) { +- const size_t offset = lengths_offset + (size_t) channel * sizeof(uint16_t); +- prefix += (uint32_t) payload[offset] | (uint32_t) payload[offset + 1] << 8; +- prefix_contents[prefix_offset++] = prefix; ++ if (memcmp(payload, "LCG2", 4) != 0) { ++ const size_t lengths_offset = 16 + (size_t) tile->token_count * sizeof(float) + ++ (size_t) job->channels * 33 * sizeof(uint16_t); ++ for (uint32_t channel = 0; channel < job->channels; ++channel) { ++ const size_t offset = lengths_offset + (size_t) channel * sizeof(uint16_t); ++ prefix += (uint32_t) payload[offset] | (uint32_t) payload[offset + 1] << 8; ++ prefix_contents[prefix_offset++] = prefix; ++ } + } + payload_offset = aligned_payload_offset + tile->payload_bytes; + } +diff --git a/ggml/src/ggml-metal/kernels/cachegen.metal b/ggml/src/ggml-metal/kernels/cachegen.metal +index 18c0a0fb4..728899e4b 100644 +--- a/ggml/src/ggml-metal/kernels/cachegen.metal ++++ b/ggml/src/ggml-metal/kernels/cachegen.metal +@@ -71,28 +71,54 @@ kernel void kernel_cachegen_decode(const device uchar * payload [[buffer + const cachegen_tile tile = tiles[tile_index]; + const device uchar * segment = payload + tile.payload_offset; + const uint bins = segment[4]; ++ const bool packed = segment[3] == '2'; + const device float * maxes = reinterpret_cast(segment + 16); +- const device ushort * cdfs = reinterpret_cast(segment + 16 + tile.rows * sizeof(float)); +- const device ushort * cdf = cdfs + channel * 33; +- const uint prefix_index = tile.prefix_offset + channel; +- const uint stream_start = prefixes[prefix_index]; +- const uint stream_bytes = prefixes[prefix_index + 1] - stream_start; +- const uint lengths_offset = 16 + tile.rows * sizeof(float) + params.channels * 33 * sizeof(ushort); +- const uint streams_offset = lengths_offset + params.channels * sizeof(ushort); +- const device uchar * stream = segment + streams_offset + stream_start; ++ const device ushort * cdf = nullptr; ++ const device uchar * stream = nullptr; ++ uint stream_bytes = 0; ++ if (packed) { ++ stream = segment + 16 + tile.rows * sizeof(float); ++ stream_bytes = uint(segment[12]) | uint(segment[13]) << 8 | ++ uint(segment[14]) << 16 | uint(segment[15]) << 24; ++ } else { ++ const device ushort * cdfs = ++ reinterpret_cast(segment + 16 + tile.rows * sizeof(float)); ++ cdf = cdfs + channel * 33; ++ const uint prefix_index = tile.prefix_offset + channel; ++ const uint stream_start = prefixes[prefix_index]; ++ stream_bytes = prefixes[prefix_index + 1] - stream_start; ++ const uint lengths_offset = 16 + tile.rows * sizeof(float) + params.channels * 33 * sizeof(ushort); ++ const uint streams_offset = lengths_offset + params.channels * sizeof(ushort); ++ stream = segment + streams_offset + stream_start; ++ } + + uint bit = 0; + uint value = 0; +- for (uint i = 0; i < 32; ++i) { +- value = (value << 1) | cachegen_read_bit(stream, stream_bytes, bit); ++ if (!packed) { ++ for (uint i = 0; i < 32; ++i) { ++ value = (value << 1) | cachegen_read_bit(stream, stream_bytes, bit); ++ } + } + uint low = 0; + uint high = 0xffffffffu; + const float center = float(bins / 2 - 1); + for (uint row = 0; row < tile.rows; ++row) { +- const ulong span = ulong(high) - ulong(low) + 1; +- const ushort count = cachegen_scaled_count(value, low, span); +- const uint symbol = cachegen_find_symbol(cdf, count); ++ const ulong span = ulong(high) - ulong(low) + 1; ++ uint symbol; ++ if (packed) { ++ const uint bits_per_symbol = segment[5]; ++ const ulong symbol_bit = (ulong(row) * params.channels + channel) * bits_per_symbol; ++ const ulong byte_index = symbol_bit >> 3; ++ const uint shift = symbol_bit & 7; ++ uint word = stream[byte_index]; ++ if (shift + bits_per_symbol > 8) { ++ word |= uint(stream[byte_index + 1]) << 8; ++ } ++ symbol = (word >> shift) & ((1u << bits_per_symbol) - 1); ++ } else { ++ const ushort count = cachegen_scaled_count(value, low, span); ++ symbol = cachegen_find_symbol(cdf, count); ++ } + if (symbol >= 32) { + return; + } +@@ -144,6 +170,9 @@ kernel void kernel_cachegen_decode(const device uchar * payload [[buffer + if (row + 1 == tile.rows) { + break; + } ++ if (packed) { ++ continue; ++ } + const ulong cdf_low = cdf[symbol]; + const ulong cdf_high = symbol == 31 ? 65536ul : cdf[symbol + 1]; + high = low - 1 + uint((span * cdf_high) >> 16); +diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp +index 215076924..dad99f62a 100644 +--- a/src/llama-kv-cache.cpp ++++ b/src/llama-kv-cache.cpp +@@ -1985,7 +1985,11 @@ static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & r + error = "invalid CacheGen record descriptor"; + return false; + } +- if (std::memcmp(payload, "LCG1", 4) != 0 || payload[5] != 0 || (payload[4] != 16 && payload[4] != 32) || ++ const bool packed = std::memcmp(payload, "LCG2", 4) == 0; ++ const uint8_t expected_bits = payload[4] == 16 ? 4 : 5; ++ if ((!packed && std::memcmp(payload, "LCG1", 4) != 0) || ++ (packed ? payload[5] != expected_bits : payload[5] != 0) || ++ (payload[4] != 16 && payload[4] != 32) || + skippy_cachegen_read_u16(payload + 6) != expected_rows || + skippy_cachegen_read_u32(payload + 8) != expected_channels) { + error = "CacheGen record disagrees with its segment header"; +@@ -1998,6 +2002,35 @@ static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & r + return false; + } + const size_t max_bytes = expected_rows * sizeof(float); ++ if (header_bytes > std::numeric_limits::max() - max_bytes || ++ header_bytes + max_bytes > record.payload_bytes) { ++ error = "CacheGen segment maximum metadata is truncated"; ++ return false; ++ } ++ for (size_t row = 0; row < expected_rows; ++row) { ++ const uint32_t bits = skippy_cachegen_read_u32(payload + header_bytes + row * sizeof(float)); ++ float value; ++ std::memcpy(&value, &bits, sizeof(value)); ++ if (!std::isfinite(value) || value < 0.0f) { ++ error = "CacheGen segment contains an invalid row maximum"; ++ return false; ++ } ++ } ++ const size_t stream_bytes = skippy_cachegen_read_u32(payload + 12); ++ if (packed) { ++ const size_t values = expected_rows * expected_channels; ++ if (values > (std::numeric_limits::max() - 7) / expected_bits) { ++ error = "CacheGen packed segment length overflows"; ++ return false; ++ } ++ const size_t expected_stream_bytes = (values * expected_bits + 7) / 8; ++ if (stream_bytes != expected_stream_bytes || ++ stream_bytes != record.payload_bytes - (header_bytes + max_bytes)) { ++ error = "CacheGen packed segment length is inconsistent"; ++ return false; ++ } ++ return true; ++ } + if (expected_channels > std::numeric_limits::max() / (cdf_entries * sizeof(uint16_t)) || + expected_channels > std::numeric_limits::max() / sizeof(uint16_t)) { + error = "CacheGen segment metadata overflows"; +@@ -2014,20 +2047,10 @@ static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & r + const size_t cdf_offset = header_bytes + max_bytes; + const size_t length_offset = cdf_offset + cdf_bytes; + const size_t stream_offset = length_offset + length_bytes; +- const size_t stream_bytes = skippy_cachegen_read_u32(payload + 12); + if (stream_offset > record.payload_bytes || stream_bytes != record.payload_bytes - stream_offset) { + error = "CacheGen segment stream length is inconsistent"; + return false; + } +- for (size_t row = 0; row < expected_rows; ++row) { +- const uint32_t bits = skippy_cachegen_read_u32(payload + header_bytes + row * sizeof(float)); +- float value; +- std::memcpy(&value, &bits, sizeof(value)); +- if (!std::isfinite(value) || value < 0.0f) { +- error = "CacheGen segment contains an invalid row maximum"; +- return false; +- } +- } + for (size_t channel = 0; channel < expected_channels; ++channel) { + const uint8_t * cdf = payload + cdf_offset + channel * cdf_entries * sizeof(uint16_t); + uint16_t previous = skippy_cachegen_read_u16(cdf); +diff --git a/tests/test-skippy-cachegen-metal.cpp b/tests/test-skippy-cachegen-metal.cpp +index b7caf2adc..b15b5e53a 100644 +--- a/tests/test-skippy-cachegen-metal.cpp ++++ b/tests/test-skippy-cachegen-metal.cpp +@@ -92,6 +92,19 @@ static const uint8_t encoded16[] = { + 0xf8, 0x88, 0xfe, 0xfc, 0x5d, 0xf0, + }; + ++static const uint8_t encoded16_packed[] = { ++ 0x4c, 0x43, 0x47, 0x32, 0x10, 0x04, 0x11, 0x00, 0x08, 0x00, 0x00, 0x00, 0x44, 0x00, 0x00, 0x00, ++ 0x00, 0xa0, 0x66, 0x3e, 0x00, 0x80, 0xcd, 0x3e, 0x00, 0x60, 0x0e, 0x3f, 0x00, 0x40, 0x2e, 0x3f, ++ 0x00, 0x40, 0x44, 0x3f, 0x00, 0x00, 0x4f, 0x3f, 0x00, 0xe0, 0x4d, 0x3f, 0x00, 0x00, 0x41, 0x3f, ++ 0x00, 0x20, 0x29, 0x3f, 0x00, 0xa0, 0x09, 0x3f, 0x00, 0xe0, 0xcd, 0x3e, 0x00, 0x20, 0x75, 0x3e, ++ 0x00, 0x80, 0x71, 0x3d, 0x00, 0xc0, 0x36, 0x3e, 0x00, 0x80, 0xb1, 0x3e, 0x00, 0xa0, 0xfa, 0x3e, ++ 0x00, 0x40, 0x1d, 0x3f, 0x76, 0x98, 0xca, 0xed, 0xa9, 0xbb, 0xdc, 0xed, 0xbb, 0xcc, 0xdd, 0xee, ++ 0xcc, 0xdc, 0xdd, 0xee, 0xdc, 0xdd, 0xed, 0xee, 0xdd, 0xdd, 0xee, 0xee, 0xdd, 0xed, 0xee, 0xee, ++ 0xed, 0xee, 0xee, 0xee, 0xee, 0xee, 0xee, 0xee, 0xee, 0xee, 0xee, 0xee, 0xee, 0xee, 0xee, 0xde, ++ 0xee, 0xde, 0xdd, 0xdd, 0xde, 0xbc, 0x9a, 0x88, 0x22, 0x11, 0x11, 0x00, 0x11, 0x00, 0x00, 0x00, ++ 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, ++}; ++ + static const uint8_t expected16[] = { + 0x1e, 0xa8, 0x00, 0x00, 0x1e, 0x28, 0x1e, 0x2c, 0x2d, 0x2e, 0x26, 0x31, 0x2d, 0x32, 0x35, 0x33, 0x57, 0x2f, 0x81, + 0x31, 0x57, 0x33, 0x57, 0x33, 0x96, 0x34, 0x81, 0x35, 0x81, 0x35, 0x6c, 0x36, 0x16, 0x35, 0x16, 0x35, 0x5b, 0x36, +@@ -185,13 +198,14 @@ int main() { + if (backend == nullptr) { + return 1; + } +- ggml_init_params params = { ggml_tensor_overhead() * 8, nullptr, true }; ++ ggml_init_params params = { ggml_tensor_overhead() * 9, nullptr, true }; + ggml_context * ctx = ggml_init(params); + if (ctx == nullptr) { + ggml_backend_free(backend); + return 2; + } + ggml_tensor * row_major16 = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, channels, capacity); ++ ggml_tensor * row_major16_packed = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, channels, capacity); + ggml_tensor * transposed16 = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, capacity, channels); + ggml_tensor * row_major32 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, channels, capacity); + ggml_tensor * transposed32 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, capacity, channels); +@@ -219,6 +233,7 @@ int main() { + cells[row] = capacity - 1 - row; + } + const ggml_backend_cachegen_tile tile16 = { encoded16, sizeof(encoded16), 0, rows }; ++ const ggml_backend_cachegen_tile tile16_packed = { encoded16_packed, sizeof(encoded16_packed), 0, rows }; + const ggml_backend_cachegen_tile tile32 = { encoded32, sizeof(encoded32), 0, rows }; + const std::vector encoded_quant = repeat_cachegen_channels(encoded16, sizeof(encoded16), 4); + const ggml_backend_cachegen_tile tile_quant = { encoded_quant.data(), encoded_quant.size(), 0, rows }; +@@ -232,6 +247,7 @@ int main() { + } + const ggml_backend_cachegen_job jobs[] = { + { row_major16, &tile16, 1, cells.data(), cells.size(), channels, channels * 2, 2 }, ++ { row_major16_packed, &tile16_packed, 1, cells.data(), cells.size(), channels, channels * 2, 2 }, + { transposed16, &tile16, 1, cells.data(), cells.size(), channels, 2, capacity * 2 }, + { row_major32, &tile32, 1, cells.data(), cells.size(), channels, channels * 4, 4 }, + { transposed32, &tile32, 1, cells.data(), cells.size(), channels, 4, capacity * 4 }, +@@ -248,7 +264,7 @@ int main() { + return finish(4); + } + char error[256] = {}; +- if (!decode(jobs, 8, error, sizeof(error))) { ++ if (!decode(jobs, 9, error, sizeof(error))) { + return finish(5); + } + ggml_backend_cachegen_job invalid_job = jobs[0]; +@@ -256,6 +272,16 @@ int main() { + if (decode(&invalid_job, 1, error, sizeof(error))) { + return finish(8); + } ++ std::vector invalid_packed(encoded16_packed, encoded16_packed + sizeof(encoded16_packed)); ++ invalid_packed[5] = 5; ++ const ggml_backend_cachegen_tile invalid_packed_tile = { ++ invalid_packed.data(), invalid_packed.size(), 0, rows, ++ }; ++ ggml_backend_cachegen_job invalid_packed_job = jobs[1]; ++ invalid_packed_job.tiles = &invalid_packed_tile; ++ if (decode(&invalid_packed_job, 1, error, sizeof(error))) { ++ return finish(14); ++ } + + const auto check_fixture = [&](ggml_tensor * row_major, ggml_tensor * transposed, + const uint8_t * expected, int status) { +@@ -284,6 +310,16 @@ int main() { + if (const int status = check_fixture(row_major16, transposed16, expected16, 6)) { + return finish(status); + } ++ std::vector packed16_bytes(ggml_nbytes(row_major16_packed)); ++ ggml_backend_tensor_get(row_major16_packed, packed16_bytes.data(), 0, packed16_bytes.size()); ++ for (uint32_t row = 0; row < rows; ++row) { ++ if (std::memcmp( ++ packed16_bytes.data() + static_cast(cells[row]) * channels * 2, ++ expected16 + static_cast(row) * channels * 2, ++ channels * 2) != 0) { ++ return finish(13); ++ } ++ } + std::vector expected32_f32(rows * channels); + for (size_t i = 0; i < expected32_f32.size(); ++i) { + uint16_t bits; +-- +2.54.0 (Apple Git-157) + diff --git a/third_party/llama.cpp/patches/0045-ggml-cuda-stage-CacheGen-payloads-directly.patch b/third_party/llama.cpp/patches/0045-ggml-cuda-stage-CacheGen-payloads-directly.patch new file mode 100644 index 0000000000..86da779ccd --- /dev/null +++ b/third_party/llama.cpp/patches/0045-ggml-cuda-stage-CacheGen-payloads-directly.patch @@ -0,0 +1,115 @@ +From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 +From: scama +Date: Sat, 12 Sep 2026 08:45:00 +1000 +Subject: [PATCH] ggml-cuda: stage CacheGen payloads directly + +--- + ggml/src/ggml-cuda/ggml-cuda.cu | 36 +++++++++++++++++++++++++----------- + 1 file changed, 25 insertions(+), 11 deletions(-) + +diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu +index ccfda8a9d..d32175261 100644 +--- a/ggml/src/ggml-cuda/ggml-cuda.cu ++++ b/ggml/src/ggml-cuda/ggml-cuda.cu +@@ -5671,6 +5671,12 @@ static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t + + #if !defined(GGML_USE_MUSA) + ++struct ggml_cuda_cachegen_payload_copy { ++ const uint8_t * source = nullptr; ++ uint32_t destination_offset = 0; ++ uint32_t bytes = 0; ++}; ++ + struct ggml_cuda_cachegen_staged_job { + int physical_device = -1; + uint8_t * dst = nullptr; +@@ -5678,7 +5684,8 @@ struct ggml_cuda_cachegen_staged_job { + uint32_t element_bytes = 0; + uint64_t token_stride = 0; + uint64_t channel_stride = 0; +- std::vector payload; ++ size_t payload_bytes = 0; ++ std::vector payload_copies; + std::vector tiles; + std::vector prefixes; + std::vector cells; +@@ -5811,24 +5818,23 @@ static bool ggml_cuda_cachegen_decode( + return ggml_cuda_cachegen_error(error, error_capacity, "CUDA/HIP CacheGen tile geometry overflows"); + } + staged.tiles.reserve(job->tile_count); ++ staged.payload_copies.reserve(job->tile_count); + for (size_t tile_index = 0; tile_index < job->tile_count; ++tile_index) { + const struct ggml_backend_cachegen_tile * tile = &job->tiles[tile_index]; + if (tile->payload == nullptr || tile->payload_bytes < 16 || tile->payload_bytes > UINT32_MAX || + tile->token_count == 0 || tile->token_count > 256 || tile->token_offset > job->cell_count || + tile->token_count > job->cell_count - tile->token_offset || +- staged.payload.size() > UINT32_MAX - 3 || ++ staged.payload_bytes > UINT32_MAX - 3 || + staged.prefixes.size() > UINT32_MAX - (static_cast(job->channels) + 1)) { + return ggml_cuda_cachegen_error( + error, error_capacity, "CUDA/HIP CacheGen tile exceeds bounded geometry"); + } + +- const size_t aligned_payload_size = (staged.payload.size() + 3) & ~size_t(3); ++ const size_t aligned_payload_size = (staged.payload_bytes + 3) & ~size_t(3); + if (tile->payload_bytes > UINT32_MAX - aligned_payload_size) { + return ggml_cuda_cachegen_error( + error, error_capacity, "CUDA/HIP CacheGen payload offsets overflow"); + } +- staged.payload.resize(aligned_payload_size, 0); +- + const uint8_t * payload = static_cast(tile->payload); + const bool packed = memcmp(payload, "LCG2", 4) == 0; + const uint32_t bins = payload[4]; +@@ -5845,7 +5851,7 @@ static bool ggml_cuda_cachegen_decode( + error, error_capacity, "CUDA/HIP CacheGen tile header is inconsistent"); + } + const ggml_cuda_cachegen_tile encoded_tile = { +- static_cast(staged.payload.size()), ++ static_cast(aligned_payload_size), + static_cast(staged.prefixes.size()), + tile->token_offset, + tile->token_count, +@@ -5893,7 +5899,12 @@ static bool ggml_cuda_cachegen_decode( + } + } + staged.tiles.push_back(encoded_tile); +- staged.payload.insert(staged.payload.end(), payload, payload + tile->payload_bytes); ++ staged.payload_copies.push_back({ ++ payload, ++ static_cast(aligned_payload_size), ++ static_cast(tile->payload_bytes), ++ }); ++ staged.payload_bytes = aligned_payload_size + tile->payload_bytes; + } + } + } catch (const std::bad_alloc &) { +@@ -5925,7 +5936,7 @@ static bool ggml_cuda_cachegen_decode( + return fail((operation), status); \ + } \ + } while (0) +- GGML_CACHEGEN_CUDA_CALL(cudaMalloc(reinterpret_cast(&job.payload_device), job.payload.size()), ++ GGML_CACHEGEN_CUDA_CALL(cudaMalloc(reinterpret_cast(&job.payload_device), job.payload_bytes), + "CUDA/HIP CacheGen payload allocation failed"); + GGML_CACHEGEN_CUDA_CALL(cudaMalloc(reinterpret_cast(&job.tiles_device), + job.tiles.size() * sizeof(ggml_cuda_cachegen_tile)), +@@ -5936,9 +5947,12 @@ static bool ggml_cuda_cachegen_decode( + GGML_CACHEGEN_CUDA_CALL(cudaMalloc(reinterpret_cast(&job.cells_device), + job.cells.size() * sizeof(uint32_t)), + "CUDA/HIP CacheGen cell allocation failed"); +- GGML_CACHEGEN_CUDA_CALL(cudaMemcpyAsync(job.payload_device, job.payload.data(), job.payload.size(), +- cudaMemcpyHostToDevice, cudaStreamPerThread), +- "CUDA/HIP CacheGen payload upload failed"); ++ for (const ggml_cuda_cachegen_payload_copy & copy : job.payload_copies) { ++ GGML_CACHEGEN_CUDA_CALL( ++ cudaMemcpyAsync(job.payload_device + copy.destination_offset, copy.source, copy.bytes, ++ cudaMemcpyHostToDevice, cudaStreamPerThread), ++ "CUDA/HIP CacheGen payload upload failed"); ++ } + GGML_CACHEGEN_CUDA_CALL(cudaMemcpyAsync(job.tiles_device, job.tiles.data(), + job.tiles.size() * sizeof(ggml_cuda_cachegen_tile), + cudaMemcpyHostToDevice, cudaStreamPerThread), +-- +2.51.0 + diff --git a/third_party/llama.cpp/patches/0046-ggml-cuda-use-native-CacheGen-shuffle-masks.patch b/third_party/llama.cpp/patches/0046-ggml-cuda-use-native-CacheGen-shuffle-masks.patch new file mode 100644 index 0000000000..477fb4744f --- /dev/null +++ b/third_party/llama.cpp/patches/0046-ggml-cuda-use-native-CacheGen-shuffle-masks.patch @@ -0,0 +1,49 @@ +From ca99c017af46be9b6cd07411ac0be33d9799d5ca Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Sat, 12 Sep 2026 11:49:37 +1000 +Subject: [PATCH] ggml-cuda: use native CacheGen shuffle masks + +--- + ggml/src/ggml-cuda/cachegen.cu | 10 +++++----- + 1 file changed, 5 insertions(+), 5 deletions(-) + +diff --git a/ggml/src/ggml-cuda/cachegen.cu b/ggml/src/ggml-cuda/cachegen.cu +index 685dc586a..ab959eae3 100644 +--- a/ggml/src/ggml-cuda/cachegen.cu ++++ b/ggml/src/ggml-cuda/cachegen.cu +@@ -124,7 +124,7 @@ static __global__ void cachegen_decode( + float amax = fabsf(decoded_f); + #pragma unroll + for (uint32_t mask = 16; mask > 0; mask >>= 1) { +- amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, mask, 32)); ++ amax = fmaxf(amax, __shfl_xor_sync(__activemask(), amax, mask, 32)); + } + const float scale = amax / 127.0f; + if (lane == 0) { +@@ -138,19 +138,19 @@ static __global__ void cachegen_decode( + float amax = fabsf(decoded_f); + #pragma unroll + for (uint32_t mask = 16; mask > 0; mask >>= 1) { +- amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, mask, 32)); ++ amax = fmaxf(amax, __shfl_xor_sync(__activemask(), amax, mask, 32)); + } + uint32_t winner = fabsf(decoded_f) == amax ? lane : 32; + #pragma unroll + for (uint32_t mask = 16; mask > 0; mask >>= 1) { +- winner = min(winner, __shfl_xor_sync(0xFFFFFFFF, winner, mask, 32)); ++ winner = min(winner, __shfl_xor_sync(__activemask(), winner, mask, 32)); + } +- const float signed_max = __shfl_sync(0xFFFFFFFF, decoded_f, winner, 32); ++ const float signed_max = __shfl_sync(__activemask(), decoded_f, winner, 32); + const float scale = signed_max / -8.0f; + if (lane == 0) { + *reinterpret_cast(dst + destination) = __float2half(scale); + } +- const float high_value = __shfl_down_sync(0xFFFFFFFF, decoded_f, 16, 32); ++ const float high_value = __shfl_down_sync(__activemask(), decoded_f, 16, 32); + if (lane < 16) { + const float inverse = scale == 0.0f ? 0.0f : 1.0f / scale; + const int low = max(0, min(15, int(decoded_f * inverse + 8.5f))); +-- +2.54.0 (Apple Git-157) + diff --git a/tools/xtask/data/console_print_allowlist.json b/tools/xtask/data/console_print_allowlist.json index ec3680a76b..2ab476eec6 100644 --- a/tools/xtask/data/console_print_allowlist.json +++ b/tools/xtask/data/console_print_allowlist.json @@ -733,6 +733,44 @@ "macro_name": "eprintln!" } ], + "crates/mesh-llm-commands/src/kv_cache.rs": [ + { + "line": 151, + "macro_name": "eprint!" + }, + { + "line": 163, + "macro_name": "println!" + }, + { + "line": 173, + "macro_name": "println!" + }, + { + "line": 181, + "macro_name": "println!" + }, + { + "line": 189, + "macro_name": "println!" + }, + { + "line": 195, + "macro_name": "println!" + }, + { + "line": 200, + "macro_name": "println!" + }, + { + "line": 205, + "macro_name": "println!" + }, + { + "line": 210, + "macro_name": "println!" + } + ], "crates/mesh-llm-commands/src/model_package.rs": [ { "line": 117, @@ -3637,6 +3675,12 @@ "macro_name": "println!" } ], + "crates/skippy-bench/src/l2_tier.rs": [ + { + "line": 254, + "macro_name": "println!" + } + ], "crates/skippy-bench/src/local_single.rs": [ { "line": 196, @@ -3663,11 +3707,11 @@ ], "crates/skippy-bench/src/main.rs": [ { - "line": 34, + "line": 35, "macro_name": "eprintln!" }, { - "line": 42, + "line": 43, "macro_name": "eprintln!" } ], @@ -3683,6 +3727,50 @@ "macro_name": "println!" } ], + "crates/skippy-cache/examples/cachegen_cubecl_spike.rs": [ + { + "line": 273, + "macro_name": "println!" + }, + { + "line": 274, + "macro_name": "println!" + }, + { + "line": 281, + "macro_name": "println!" + }, + { + "line": 285, + "macro_name": "println!" + }, + { + "line": 291, + "macro_name": "println!" + }, + { + "line": 292, + "macro_name": "println!" + }, + { + "line": 350, + "macro_name": "println!" + }, + { + "line": 380, + "macro_name": "println!" + }, + { + "line": 383, + "macro_name": "eprintln!" + } + ], + "crates/skippy-cache/src/l3/tests.rs": [ + { + "line": 603, + "macro_name": "println!" + } + ], "crates/skippy-correctness/src/glm_dsa_trace.rs": [ { "line": 1434, @@ -3691,11 +3779,11 @@ ], "crates/skippy-correctness/src/main.rs": [ { - "line": 22, + "line": 25, "macro_name": "eprintln!" }, { - "line": 30, + "line": 33, "macro_name": "eprintln!" } ], @@ -3705,12 +3793,56 @@ "macro_name": "println!" } ], + "crates/skippy-correctness/src/runner/kv_page_growth.rs": [ + { + "line": 285, + "macro_name": "println!" + }, + { + "line": 340, + "macro_name": "println!" + } + ], "crates/skippy-correctness/src/runner/native_mtp.rs": [ { "line": 250, "macro_name": "println!" } ], + "crates/skippy-correctness/src/runner/remote_handoff.rs": [ + { + "line": 188, + "macro_name": "eprintln!" + }, + { + "line": 224, + "macro_name": "eprintln!" + }, + { + "line": 933, + "macro_name": "eprintln!" + }, + { + "line": 942, + "macro_name": "eprintln!" + }, + { + "line": 954, + "macro_name": "eprintln!" + }, + { + "line": 957, + "macro_name": "eprintln!" + }, + { + "line": 961, + "macro_name": "eprintln!" + }, + { + "line": 1309, + "macro_name": "eprintln!" + } + ], "crates/skippy-correctness/src/runner/stage_fa_parity.rs": [ { "line": 51, @@ -4317,27 +4449,27 @@ ], "crates/skippy-server/src/binary_transport/binary_messaging.rs": [ { - "line": 424, + "line": 427, "macro_name": "eprintln!" }, { - "line": 434, + "line": 437, "macro_name": "println!" }, { - "line": 469, + "line": 472, "macro_name": "eprintln!" }, { - "line": 490, + "line": 493, "macro_name": "eprintln!" }, { - "line": 498, + "line": 501, "macro_name": "eprintln!" }, { - "line": 567, + "line": 570, "macro_name": "eprintln!" } ], @@ -4429,7 +4561,7 @@ "macro_name": "println!" }, { - "line": 347, + "line": 350, "macro_name": "println!" } ], diff --git a/website/src/_data/docs.js b/website/src/_data/docs.js index 7a8e478f83..40d5505097 100644 --- a/website/src/_data/docs.js +++ b/website/src/_data/docs.js @@ -40,6 +40,7 @@ export default [ links: [ ["OpenAI-compatible API", "/docs/pages/openai-compatible-api/"], ["Automatic routing", "/docs/pages/automatic-routing/"], + ["KV caching", "/docs/pages/kv-caching/"], ["Streaming", "/docs/pages/openai-compatible-api/#streaming"], ["Tool calling", "/docs/pages/openai-compatible-api/#tool-calling"], ["Structured outputs", "/docs/pages/openai-compatible-api/#structured-outputs"] diff --git a/website/src/docs/pages/config-defaults.md b/website/src/docs/pages/config-defaults.md index d69d7ec76f..97f6b63306 100644 --- a/website/src/docs/pages/config-defaults.md +++ b/website/src/docs/pages/config-defaults.md @@ -14,8 +14,9 @@ Shared default settings applied to every model. Individual model entries can ove ctx_size = 0 # Context window (0 = auto) batch = 0 # Batch size (0 = auto) ubatch = 0 # Micro-batch size (0 = auto) -cache_type_k = "f16" # Key cache dtype -cache_type_v = "f16" # Value cache dtype +cache_type_k = "auto" # Key cache dtype (resolved by policy) +cache_type_v = "auto" # Value cache dtype (resolved by policy) +kv_cache_policy = "balanced" # "balanced", "quality", or "saver" kv_offload = "auto" # KV-cache offload policy prompt_cache = "auto" # Prompt-cache policy flash_attention = "auto" # Flash-attention policy diff --git a/website/src/docs/pages/config-reference.md b/website/src/docs/pages/config-reference.md index dc65d8f669..cd5cdd9c56 100644 --- a/website/src/docs/pages/config-reference.md +++ b/website/src/docs/pages/config-reference.md @@ -117,6 +117,10 @@ produces a clear startup error rather than a partial start. | `runtime.reconcile_model_targets` | boolean | `false` | node-level | process restart | wired | none | | `runtime.reconcile_model_target_demand_upgrades` | boolean | `false` | node-level | process restart | wired | none | | `runtime.native_runtime.mesh_version`
`runtime.native_runtime.skippy_abi`
`runtime.native_runtime.selection` | string | unset (auto-selected) | node-level | process restart | wired | none | +| `runtime.kv_cache.disk.mode` | enum | `off` (default), `auto`, `fixed` | node-level | process restart | wired | `--kv-cache-disk` | +| `runtime.kv_cache.disk.directory` | absolute path | `$MESH_LLM_HOME/kv-cache` | node-level | process restart | wired | `--kv-cache-disk-dir` | +| `runtime.kv_cache.disk.budget_mib` | integer | required and > 0 only for `fixed` | node-level | applies dynamically | wired | fixed size passed to `--kv-cache-disk` | +| `runtime.kv_cache.disk.minimum_free_mib` | integer | `16384`; minimum `1024` | node-level | applies dynamically | wired | `--kv-cache-min-free` | | `runtime.model_target_demand_upgrade_min_requests` | integer | `2` | node-level | process restart | wired | none | | `runtime.model_target_demand_upgrade_max_age_secs` | integer | `3600` | node-level | process restart | wired | none | | `advanced.server.alias` | string | unset; per-model alias overrides the default | both | model reload | wired; becomes the served identity used by `/v1/models` and routing | none | @@ -135,14 +139,14 @@ for the activity policy and privacy boundary. | `model_fit.ctx_size` | integer | `0` = auto | both | model reload | wired | `--ctx-size` on the ad-hoc single-model path | | `model_fit.batch` | integer | `0` = auto (`n_batch`) | both | model reload | wired | none | | `model_fit.ubatch` | integer | `0` = auto (`n_ubatch`); should not exceed `batch` | both | model reload | wired | none | -| `model_fit.cache_type_k`
`model_fit.cache_type_v` | enum (dtype) | `auto`, `f32`, `f16` (default), `bf16`, `q8_0`, `q4_0`, `q4_1`, `iq4_nl`, `q5_0`, `q5_1`; explicit value overrides `kv_cache_policy` | both | model reload | wired | none | -| `model_fit.kv_cache_policy` | enum | `auto`, `quality`, `balanced`, `saver`; expands into cache dtypes | both | model reload | wired | none | +| `model_fit.cache_type_k`
`model_fit.cache_type_v` | enum (dtype) | `auto` (default), `f16`, `q8_0`, `q4_0`; schema validation also accepts `f32`, `bf16`, `q4_1`, `iq4_nl`, `q5_0`, `q5_1`, but the normal serving runtime rejects them at model load; explicit value overrides `kv_cache_policy` | both | model reload | wired | none | +| `model_fit.kv_cache_policy` | enum | `balanced` (default), `auto`, `quality`, `saver`; expands into cache dtypes | both | model reload | wired | none | | `model_fit.kv_offload` | bool-or-`auto` | `auto` | both | model reload | wired | none | | `model_fit.kv_unified` | bool-or-`auto` | `auto` | both | model reload | wired (recurrent/hybrid architectures still force this true natively) | none | | `model_fit.cache_ram_mib` | integer | unset (no cap) | both | model reload | unwired (any positive value fails at model load) | none | -| `model_fit.cache_idle_slots` | integer | unset (unbounded) | both | model reload | wired | none | +| `model_fit.cache_idle_slots` | integer | unset uses the runtime lane count; `0` drops every reset lane, positive values cap retained idle sessions | both | model reload | wired | none | | `model_fit.prompt_cache` | bool-or-`auto` | `auto` | both | model reload | wired | none | -| `model_fit.prefix_cache.enabled` | boolean | unset (disabled) | both | model reload | wired | none | +| `model_fit.prefix_cache.enabled` | boolean | unset uses family defaults; `false` disables | both | model reload | wired | none | | `model_fit.prefix_cache.max_entries` | integer | runtime default | both | model reload | wired | none | | `model_fit.prefix_cache.max_bytes` | integer | `0`/unset = no cap | both | model reload | wired | none | | `model_fit.prefix_cache.min_tokens` | integer | runtime default | both | model reload | wired | none | @@ -160,6 +164,9 @@ Missing TOML for this group: GGUF metadata `kv_overrides`. There is no schema key for it yet; do not expect an override path until a later PR adds one. +See [KV Caching](/docs/pages/kv-caching/) for the default behavior and common +configuration recipes. + ## Group 4: device selection, GPU offload, multi-GPU, CPU MoE, and loading behavior | Key path | Type | Allowed values / default (`auto`) | `[defaults]` / `[[models]]` | Restart | Status | CLI equivalent | diff --git a/website/src/docs/pages/kv-caching.md b/website/src/docs/pages/kv-caching.md new file mode 100644 index 0000000000..df79ac2005 --- /dev/null +++ b/website/src/docs/pages/kv-caching.md @@ -0,0 +1,256 @@ +--- +title: KV Caching +description: Reuse prompt work in memory or on disk and tune KV memory use +--- + +# KV Caching + +KV caching lets Mesh reuse the model state created while reading a prompt. +When a later request starts with the same tokens, Mesh can restore that state +and skip some or all of the repeated prefill work. This lowers time to first +token for repeated system prompts, long documents, and multi-turn chats. + +Applications do not need a cache-specific API. Keep using the +[OpenAI-compatible API](/docs/pages/openai-compatible-api/) and send the full +conversation or prompt on every request. Mesh identifies matching prefixes and +reports reused prompt tokens in `usage.prompt_tokens_details.cached_tokens`. +Callers that already use OpenAI's `prompt_cache_key` may keep sending it; Mesh +trims the value and uses it as a cache and routing namespace. Keep it stable for +the requests that should share cached prefixes. It does not force a cache hit, +and requests without it share the default namespace. Do not put secrets in the +key because Mesh records it verbatim in request telemetry. + +For the Responses API, an explicit `prompt_cache_key` wins. When it is absent, +Mesh uses `previous_response_id`, then the conversation ID, as the cache key. + +## Default experience + +With no cache settings: + +- Mesh enables the in-memory prefix cache with family-aware limits. Prefixes of + at least 256 tokens are eligible, and Mesh selects the stored state format + from the model architecture. +- With no explicit K/V dtype, the resolver selects Q8_0 for both caches when + the model is smaller than 50 GiB and Q4_0 for both caches at 50 GiB or above. + It checks GGUF architecture and head dimensions before applying that default + and falls back to F16 when metadata proves the quantized layout is invalid. + An explicit incompatible dtype fails instead of changing silently. +- KV offload and unified-cache behavior remain automatic. +- The durable disk cache is off, so a process restart starts with an empty + prompt cache and Mesh writes no prompt state to disk. + +These defaults apply across the supported CPU, Metal, CUDA, and ROCm runtimes. + +## The controls at a glance + +| Goal | Setting | Default | Where to set it | +|---|---|---|---| +| Pin key/value cache formats | `model_fit.cache_type_k`, `model_fit.cache_type_v` | `auto` | config file | +| Control KV device offload | `model_fit.kv_offload` | `auto` | config file | +| Control unified KV allocation | `model_fit.kv_unified` | `auto` | config file | +| Control the attention kernel required by quantized V | `model_fit.flash_attention` | derived from V dtype | config file | +| Cap retained idle native sessions | `model_fit.cache_idle_slots` | lane count | config file | +| Disable all prompt-prefix reuse | `model_fit.prompt_cache` | `auto` | config file | +| Tune or disable in-memory prefix reuse | `model_fit.prefix_cache.*` | family defaults | config file | +| Persist prompt state across restarts | `runtime.kv_cache.disk.*` | `off` | config, environment, or `serve` flags | +| Inspect or remove disk entries | `mesh-llm kv-cache ...` | n/a | CLI | + +Model-level cache controls do not currently have CLI equivalents. Use +`~/.mesh-llm/config.toml`, or pass a different file with `mesh-llm serve +--config PATH`. The disk tier has CLI overrides for one-off runs. + +## Choose the KV representation + +Configure the representation directly when you need deterministic behavior: + +```toml +[defaults.model_fit] +cache_type_k = "q8_0" +cache_type_v = "q8_0" +kv_offload = "auto" +kv_unified = "auto" +flash_attention = "enabled" +``` + +The relevant controls are: + +| Setting | Loadable embedded-runtime values | Runtime effect | +|---|---|---| +| `cache_type_k` | `auto`, `f16`, `q8_0`, `q4_0` | Storage and compute dtype for attention keys | +| `cache_type_v` | `auto`, `f16`, `q8_0`, `q4_0` | Storage and compute dtype for attention values | +| `kv_offload` | `auto`, `true`, `false` | Whether KV tensors may reside on the selected accelerator rather than host memory | +| `kv_unified` | `auto`, `true`, `false` | Whether runtime slots use the backend unified KV allocation | +| `flash_attention` | `auto`, `enabled`, `disabled` | Selects the fused attention path; a quantized V cache requires the enabled path | + +Q8_0 and Q4_0 encode values in 32-element blocks. A model whose KV head +dimension cannot satisfy that block layout cannot use the corresponding +quantized cache. Automatic selection can detect that from GGUF metadata and +fall back to F16. Explicit dtype selection bypasses that fallback so invalid +combinations fail during model load. A backend Flash Attention capability +failure is only known when the runtime loads, so metadata validation alone +cannot prove that every quantized combination will start. + +K and V may use different dtypes. For these technical fields, per-model values +override global values, which override family defaults and finally the built-in +size rule. + +The config validator currently recognizes additional GGML dtype labels that +the pinned embedded runtime does not load. The table above lists the values +accepted by `skippy_runtime::parse_cache_type`; use those values for a serving +configuration. `auto` is consumed by the resolver and does not reach that +parser. + +You can override one model without changing the others: + +```toml +[[models]] +model = "org/model-GGUF" + +[models.model_fit] +cache_type_k = "f16" +cache_type_v = "f16" +kv_offload = false +``` + +## Tune in-memory prefix reuse + +The automatic prefix cache is usually the right choice. To disable it for a +model, set: + +```toml +[defaults.model_fit] +prompt_cache = false +``` + +To keep it enabled but set explicit bounds: + +```toml +[defaults.model_fit.prefix_cache] +enabled = true +payload_mode = "auto" +min_tokens = 512 +max_entries = 256 +max_bytes = 8589934592 +shared_stride_tokens = 128 +shared_record_limit = 4 +``` + +`payload_mode = "auto"` stores resident KV for known dense models and KV plus +recurrent state for known recurrent or hybrid models. Unknown architectures do +not cache automatically. The size fields are byte counts; `max_bytes = 0` +means no explicit byte cap. Setting `prompt_cache = false` disables prefix +caching and conflicts with an explicitly enabled `prefix_cache` block. + +Per-model `[[models]]` values override `[defaults]`; explicit cache dtypes +override policy-derived dtypes; unresolved values fall back through family +policy and built-ins. Model cache changes apply when the model reloads. + +The OpenAI request field `prompt_cache_retention` accepts `in_memory` and +`24h`. Mesh records it as telemetry, but neither value currently enforces a +cache lifetime. Use the runtime limits above and the disk maintenance commands +below to control retention. + +`cache_idle_slots` limits how many reset native sessions remain available for +reuse. Unset means the runtime lane count is the bound, `0` drops every reset +lane, and a positive value adds a lower cap. `cache_ram_mib` is reserved; any +positive value currently fails model loading, so there is no configurable +host-RAM L2 tier. + +Cache matches require the exact token prefix and exact runtime identity. Mesh +does not use fuzzy or semantic prompt matching. + +## Enable durable disk caching + +Disk caching preserves reusable prompt state across process restarts. It is +node-local and disabled until you opt in. + +For an automatically sized cache: + +```toml +[runtime.kv_cache.disk] +mode = "auto" +directory = "/var/lib/mesh-llm/kv-cache" +minimum_free_mib = 16384 +``` + +For a fixed 32 GiB cap: + +```toml +[runtime.kv_cache.disk] +mode = "fixed" +directory = "/var/lib/mesh-llm/kv-cache" +budget_mib = 32768 +minimum_free_mib = 16384 +``` + +The directory must be absolute. If omitted, it is +`$MESH_LLM_HOME/kv-cache`, or `~/.mesh-llm/kv-cache`. Auto mode uses at most +20% of the filesystem capacity basis, never consumes the configured free-space +reserve, and is capped at 64 GiB. + +The same settings can be supplied for one run: + +```bash +mesh-llm serve --kv-cache-disk auto + +mesh-llm serve \ + --kv-cache-disk 32GiB \ + --kv-cache-disk-dir /var/lib/mesh-llm/kv-cache \ + --kv-cache-min-free 16GiB +``` + +Environment equivalents are `MESH_LLM_KV_CACHE_DISK`, +`MESH_LLM_KV_CACHE_DISK_DIR`, and `MESH_LLM_KV_CACHE_MIN_FREE`. CLI values win +over environment values, which win over the config file, independently for +each field. + +Invalid disk settings stop startup with a configuration error. Once a valid +configuration is running, storage trouble fails open: Mesh logs the problem +and serves the request with cold prefill. + +## Inspect and maintain the disk cache + +```bash +mesh-llm kv-cache status +mesh-llm kv-cache status --json + +mesh-llm kv-cache prune --target 16GiB --yes +mesh-llm kv-cache clear --yes +``` + +Human-readable `status` shows the effective state, configured mode, root, and +used/budget bytes. `status --json` also includes the free-space reserve, entry +counts, degradation reason, activity, reconciliation, and per-model inventory. +`prune` and `clear` only remove inactive entries. Both can be limited to an +exact numeric model identity with `--model-identity ID`. + +Budget and minimum-free changes apply live. Changing the mode or directory +requires a node restart. + +## Advanced environment controls + +These environment variables are intended for incident response and controlled +experiments: + +- Setting `SKIPPY_KV_CACHE=off` or `SKIPPY_PREFIX_CACHE=off` disables worker + prefix-cache storage even when the stage plan enables it. Other + `SKIPPY_KV_CACHE_*` tuning variables only construct cache settings when the + stage plan did not supply them. +- `MESH_LLM_DISABLE_PREFIX_AFFINITY` disables routing toward a peer with known + prefix state. `MESH_LLM_DISABLE_STICKY_ROUTING` disables sticky routing, and + `MESH_LLM_PREFIX_ONLY=1` uses the request's prefix hash as the deterministic + fallback when neither cache evidence nor a session route applies. These + change peer selection; they do not disable or resize cache storage. + +## CacheGen status + +CacheGen is an experimental compressed representation for saved KV state. The +implementation targets CPU, Metal, CUDA, and ROCm, but it is still undergoing +backend qualification, including completion and verification of the ROCm path, +and is not selected by the normal serving config or CLI today. Use the stable +resident/exact-state cache paths described above for production operation; do +not add an invented `cachegen` setting to the config. + +For every field and allowed value, see the [Config Reference](/docs/pages/config-reference/). +For disk storage details, failure modes, and recovery procedures, see the +[KV-cache disk operator guide](https://github.com/Mesh-LLM/mesh-llm/blob/main/docs/skippy/KV_CACHE_DISK.md). diff --git a/website/src/docs/pages/openai-compatible-api.md b/website/src/docs/pages/openai-compatible-api.md index b1eb9e73c0..bafd2674af 100644 --- a/website/src/docs/pages/openai-compatible-api.md +++ b/website/src/docs/pages/openai-compatible-api.md @@ -48,6 +48,11 @@ curl -s http://localhost:9337/v1/chat/completions -H "Content-Type: application/ Clients that support streamed OpenAI-compatible responses can use the same base URL. +Repeated prompt prefixes are reused automatically when the selected model +supports them. Responses report reused prompt tokens in +`usage.prompt_tokens_details.cached_tokens`. See [KV Caching](/docs/pages/kv-caching/) +for defaults, memory policy, and durable disk-cache controls. + ## Tool calling Tool-calling support depends on the selected model and the agent client. Start with console chat, then test the specific agent workflow you plan to use. diff --git a/website/src/docs/pages/skippy-api.md b/website/src/docs/pages/skippy-api.md index 207c8dbd61..8cb3189890 100644 --- a/website/src/docs/pages/skippy-api.md +++ b/website/src/docs/pages/skippy-api.md @@ -9,7 +9,7 @@ description: Generated reference for the capability-oriented Skippy C ABI. This reference is generated from the patched llama.cpp public headers. It documents the native C ABI used by Skippy's Rust FFI layer and staged runtime. The ABI is experimental and versioned for lockstep native/Rust builds. -Current generated surface: **15 headers** and **99 exported functions**. +Current generated surface: **15 headers** and **100 exported functions**. ## Quick navigation @@ -156,7 +156,7 @@ Current generated surface: **15 headers** and **99 exported functions**.
- state.h14 functions + state.h15 functions
skippy_export_state skippy_import_state @@ -168,6 +168,7 @@ Current generated surface: **15 headers** and **99 exported functions**. skippy_retire_verify_checkpoint skippy_export_kv_page skippy_import_kv_page + skippy_import_cachegen_kv_page_v1 skippy_session_save_prefix skippy_session_restore_prefix skippy_session_memory_used_cells @@ -1508,6 +1509,20 @@ LLAMA_API enum skippy_status skippy_import_kv_page( struct skippy_error ** out_error); ``` + +#### `skippy_import_cachegen_kv_page_v1` + +Imports validated CacheGen records directly into resident KV storage. + +```cpp +LLAMA_API enum skippy_status skippy_import_cachegen_kv_page_v1( + struct skippy_session * session, + const struct skippy_kv_page_desc * desc, + const struct skippy_cachegen_record_v1 * records, + size_t record_count, + struct skippy_error ** out_error); +``` + #### `skippy_session_save_prefix` @@ -1674,7 +1689,7 @@ SKIPPY_COMMON_API enum skippy_status skippy_parse_chat_response_json( The headers also define the following enums, structs, opaque handles, and ABI constants: - `activation.h`: `skippy_activation_dtype`, `skippy_activation_layout`, `skippy_activation_boundary_desc`, `skippy_activation_desc`, `SKIPPY_ACTIVATION_BOUNDARY_DESC_VERSION = 1`, `SKIPPY_ACTIVATION_SIDEBAND_TOKEN_IDS = (UINT64_C(1) << 0)`, `SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST = (UINT64_C(1) << 0)`, `SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP = (UINT64_C(1) << 1)`, `SKIPPY_ACTIVATION_FLAG_INKLING_MTP_EMBD = (UINT64_C(1) << 2)`, `SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K = (UINT64_C(1) << 3)` -- `common.h`: `skippy_feature`, `skippy_status`, `skippy_error`, `skippy_abi_version`, `SKIPPY_ABI_VERSION_MAJOR = 0`, `SKIPPY_ABI_VERSION_MINOR = 1`, `SKIPPY_ABI_VERSION_PATCH = 54`, `SKIPPY_FEATURE_RUNTIME_EVENT_REPORTER = ((uint64_t)1 << 31)`, `SKIPPY_FEATURE_MODEL_LOAD_EVENTS_V2 = ((uint64_t)1 << 32)`, `SKIPPY_FEATURE_KV_EVENTS = ((uint64_t)1 << 33)`, `SKIPPY_FEATURE_DEVICE_EVENTS = ((uint64_t)1 << 34)`, `SKIPPY_FEATURE_DIAGNOSTIC_EVENTS = ((uint64_t)1 << 35)`, `SKIPPY_FEATURE_UNLOAD_EVENTS = ((uint64_t)1 << 36)` +- `common.h`: `skippy_feature`, `skippy_status`, `skippy_error`, `skippy_abi_version`, `SKIPPY_ABI_VERSION_MAJOR = 0`, `SKIPPY_ABI_VERSION_MINOR = 1`, `SKIPPY_ABI_VERSION_PATCH = 55`, `SKIPPY_FEATURE_RUNTIME_EVENT_REPORTER = ((uint64_t)1 << 31)`, `SKIPPY_FEATURE_MODEL_LOAD_EVENTS_V2 = ((uint64_t)1 << 32)`, `SKIPPY_FEATURE_KV_EVENTS = ((uint64_t)1 << 33)`, `SKIPPY_FEATURE_DEVICE_EVENTS = ((uint64_t)1 << 34)`, `SKIPPY_FEATURE_DIAGNOSTIC_EVENTS = ((uint64_t)1 << 35)`, `SKIPPY_FEATURE_UNLOAD_EVENTS = ((uint64_t)1 << 36)`, `SKIPPY_FEATURE_CACHEGEN_KV_PAGE = (UINT64_C(1) << 37)` - `devices.h`: `skippy_backend_device_type`, `skippy_backend_device_cap`, `skippy_backend_device` - `events.h`: `skippy_runtime_event_v1`, `skippy_runtime_event_reporter_v1`, `SKIPPY_RUNTIME_EVENT_V1_ABI_VERSION = 1` - `execution.h`: `skippy_iteration_request` @@ -1685,6 +1700,6 @@ The headers also define the following enums, structs, opaque handles, and ABI co - `signals.h`: `skippy_token_signal`, `skippy_generation_signal_window` - `speculative_decoding.h`: `skippy_ngram_cache`, `skippy_native_mtp_draft`, `SKIPPY_NATIVE_MTP_MAX_DRAFT_TOKENS = 8` - `stage_plan.h`: `skippy_stage_planner`, `skippy_stage_plan`, `skippy_stage_plan_string_ref_v1`, `skippy_stage_planner_tensor_v1`, `skippy_stage_planner_profile_v1`, `skippy_stage_planner_config_v1`, `skippy_stage_plan_value_kind`, `skippy_stage_plan_state_kind`, `skippy_stage_plan_state_access`, `skippy_stage_plan_desc_v1`, `skippy_stage_plan_profile_desc_v1`, `skippy_stage_plan_value_desc_v1`, `skippy_stage_plan_state_desc_v1`, `SKIPPY_STAGE_PLANNER_CONFIG_V1_ABI_VERSION = 1`, `SKIPPY_STAGE_PLANNER_TENSOR_V1_ABI_VERSION = 1`, `SKIPPY_STAGE_PLANNER_PROFILE_V1_ABI_VERSION = 1`, `SKIPPY_STAGE_PLAN_DESC_V1_ABI_VERSION = 1`, `SKIPPY_STAGE_PLAN_PROFILE_DESC_V1_ABI_VERSION = 1`, `SKIPPY_STAGE_PLAN_VALUE_DESC_V1_ABI_VERSION = 1`, `SKIPPY_STAGE_PLAN_STATE_DESC_V1_ABI_VERSION = 1`, `SKIPPY_STAGE_PLAN_MAX_DIMS = 4` -- `state.h`: `skippy_kv_page_flag`, `skippy_kv_page_codec`, `skippy_kv_page_component_role`, `skippy_kv_page_component_desc`, `skippy_kv_page_desc` +- `state.h`: `skippy_kv_page_flag`, `skippy_kv_page_codec`, `skippy_kv_page_component_role`, `skippy_cachegen_record_kind`, `skippy_cachegen_record_v1`, `skippy_kv_page_component_desc`, `skippy_kv_page_desc`, `SKIPPY_CACHEGEN_RECORD_V1_ABI_VERSION = 1` Source directory: `include/skippy/`. Regenerate this page after changing any public header or exported function. From 551e9353462d4634369cc895408aef5849ccfa5d Mon Sep 17 00:00:00 2001 From: scama Date: Sun, 13 Sep 2026 16:30:36 +1000 Subject: [PATCH 02/16] fix(ci): preserve native gate evidence --- .../tests/runtime_events_native.rs | 79 +++++++++++++------ .../tests/test_runtime_events_native_gate.py | 1 - tools/xtask/data/console_print_allowlist.json | 12 ++- 3 files changed, 61 insertions(+), 31 deletions(-) diff --git a/crates/skippy-runtime/tests/runtime_events_native.rs b/crates/skippy-runtime/tests/runtime_events_native.rs index e78c273221..d54f548924 100644 --- a/crates/skippy-runtime/tests/runtime_events_native.rs +++ b/crates/skippy-runtime/tests/runtime_events_native.rs @@ -34,46 +34,62 @@ const EVIDENCE_FILE_ENV: &str = "MESH_LLM_RUNTIME_EVENTS_EVIDENCE_FILE"; /// `MESH_LLM_RUNTIME_EVENTS_EVIDENCE_FILE` is unset. The actual file I/O is /// `skippy_runtime::write_evidence_marker`, unit tested directly in /// `crates/skippy-runtime/src/native_test_evidence.rs`. -fn write_marker(line: &str) { - let path = env::var_os(EVIDENCE_FILE_ENV).map(PathBuf::from); - skippy_runtime::write_evidence_marker(path.as_deref(), line); +fn write_marker(path: Option<&std::path::Path>, line: &str) { + skippy_runtime::write_evidence_marker(path, line); } #[test] fn runtime_events_native_gate() { + // Resolve the evidence destination before loading native libraries. The + // loader and model-open path are process-global, so the marker destination + // must remain stable for the full gate. + let evidence_path = env::var_os(EVIDENCE_FILE_ENV).map(PathBuf::from); + if env::var(GATE_ENV).ok().as_deref() != Some("1") { println!("BLOCKED: {GATE_ENV} unset"); - write_marker("blocked-when-ungated: gate unset, no native symbol was touched"); + write_marker( + evidence_path.as_deref(), + "blocked-when-ungated: gate unset, no native symbol was touched", + ); return; } #[cfg(not(feature = "dynamic-native-runtime"))] { println!("BLOCKED: dynamic-native-runtime feature not enabled"); - write_marker("blocked: dynamic-native-runtime feature is not enabled for this run"); + write_marker( + evidence_path.as_deref(), + "blocked: dynamic-native-runtime feature is not enabled for this run", + ); panic!("{GATE_ENV}=1 requires the dynamic-native-runtime feature"); } #[cfg(feature = "dynamic-native-runtime")] { - run_real_native_gate(); + run_real_native_gate(evidence_path); } } #[cfg(feature = "dynamic-native-runtime")] -fn run_real_native_gate() { +fn run_real_native_gate(evidence_path: Option) { + let evidence_path = evidence_path.unwrap_or_else(|| { + println!("BLOCKED: {EVIDENCE_FILE_ENV} unset"); + panic!("{GATE_ENV}=1 requires {EVIDENCE_FILE_ENV} to name the evidence file") + }); let bundle_dir = env::var(BUNDLE_DIR_ENV).unwrap_or_else(|_| { println!("BLOCKED: {BUNDLE_DIR_ENV} unset"); - write_marker(&format!( - "blocked: {BUNDLE_DIR_ENV} unset, required when {GATE_ENV}=1" - )); + write_marker( + Some(&evidence_path), + &format!("blocked: {BUNDLE_DIR_ENV} unset, required when {GATE_ENV}=1"), + ); panic!("{GATE_ENV}=1 requires {BUNDLE_DIR_ENV} to point at a dynamic native runtime") }); let model_path = env::var(MODEL_ENV).unwrap_or_else(|_| { println!("BLOCKED: {MODEL_ENV} unset"); - write_marker(&format!( - "blocked: {MODEL_ENV} unset, required when {GATE_ENV}=1" - )); + write_marker( + Some(&evidence_path), + &format!("blocked: {MODEL_ENV} unset, required when {GATE_ENV}=1"), + ); panic!("{GATE_ENV}=1 requires {MODEL_ENV} to name a readable model") }); @@ -197,22 +213,33 @@ fn run_real_native_gate() { // Only after reporter installation, successful model-open, structured // production callbacks, and the unload exercise have all completed do we // claim that this opt-in path actually executed. - write_marker("executed"); + write_marker(Some(&evidence_path), "executed"); write_marker( + Some(&evidence_path), "exact-abi-admission: native runtime loaded (loader enforces exact major.minor.patch)", ); - write_marker(&format!( - "capability-probe: confirmed={:#x} health_messages={}", - report.confirmed, - report.health_messages.len() - )); - write_marker("reporter-install: true"); - write_marker("model-open: single-part real model-open succeeded"); - write_marker(&format!( - "structured-production-callbacks: {structured_count}" - )); - write_marker(&format!("unload-callbacks: {unload_count}")); - write_marker("reporter-clear: returned"); + write_marker( + Some(&evidence_path), + &format!( + "capability-probe: confirmed={:#x} health_messages={}", + report.confirmed, + report.health_messages.len() + ), + ); + write_marker(Some(&evidence_path), "reporter-install: true"); + write_marker( + Some(&evidence_path), + "model-open: single-part real model-open succeeded", + ); + write_marker( + Some(&evidence_path), + &format!("structured-production-callbacks: {structured_count}"), + ); + write_marker( + Some(&evidence_path), + &format!("unload-callbacks: {unload_count}"), + ); + write_marker(Some(&evidence_path), "reporter-clear: returned"); } /// Resolves the real installed-runtime layout: `MESH_LLM_NATIVE_RUNTIME_BUNDLE_DIR` diff --git a/scripts/tests/test_runtime_events_native_gate.py b/scripts/tests/test_runtime_events_native_gate.py index 18e0acb502..6cb0274eba 100644 --- a/scripts/tests/test_runtime_events_native_gate.py +++ b/scripts/tests/test_runtime_events_native_gate.py @@ -103,7 +103,6 @@ def run_gate( evidence = root / "evidence.txt" if evidence_seed is not None: evidence.write_text(evidence_seed, encoding="utf-8") - return subprocess.run( [ "bash", diff --git a/tools/xtask/data/console_print_allowlist.json b/tools/xtask/data/console_print_allowlist.json index 2ab476eec6..94c3d891da 100644 --- a/tools/xtask/data/console_print_allowlist.json +++ b/tools/xtask/data/console_print_allowlist.json @@ -4425,19 +4425,23 @@ ], "crates/skippy-runtime/tests/runtime_events_native.rs": [ { - "line": 45, + "line": 49, "macro_name": "println!" }, { - "line": 52, + "line": 59, "macro_name": "println!" }, { - "line": 66, + "line": 76, "macro_name": "println!" }, { - "line": 73, + "line": 80, + "macro_name": "println!" + }, + { + "line": 88, "macro_name": "println!" } ], From 5a406369d892fec2e1dda72154c2f9851acbd45d Mon Sep 17 00:00:00 2001 From: scama Date: Sun, 13 Sep 2026 17:31:18 +1000 Subject: [PATCH 03/16] fix(ggml-cuda): make CacheGen masks portable --- ...da-use-native-CacheGen-shuffle-masks.patch | 41 ++++++++++++++----- 1 file changed, 30 insertions(+), 11 deletions(-) diff --git a/third_party/llama.cpp/patches/0046-ggml-cuda-use-native-CacheGen-shuffle-masks.patch b/third_party/llama.cpp/patches/0046-ggml-cuda-use-native-CacheGen-shuffle-masks.patch index 477fb4744f..96bd122998 100644 --- a/third_party/llama.cpp/patches/0046-ggml-cuda-use-native-CacheGen-shuffle-masks.patch +++ b/third_party/llama.cpp/patches/0046-ggml-cuda-use-native-CacheGen-shuffle-masks.patch @@ -4,46 +4,65 @@ Date: Sat, 12 Sep 2026 11:49:37 +1000 Subject: [PATCH] ggml-cuda: use native CacheGen shuffle masks --- - ggml/src/ggml-cuda/cachegen.cu | 10 +++++----- - 1 file changed, 5 insertions(+), 5 deletions(-) + ggml/src/ggml-cuda/cachegen.cu | 18 +++++++++++++----- + 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-cuda/cachegen.cu b/ggml/src/ggml-cuda/cachegen.cu -index 685dc586a..ab959eae3 100644 +index 685dc586a..394c1d45e 100644 --- a/ggml/src/ggml-cuda/cachegen.cu +++ b/ggml/src/ggml-cuda/cachegen.cu -@@ -124,7 +124,7 @@ static __global__ void cachegen_decode( +@@ -2,6 +2,12 @@ + + #if !defined(GGML_USE_MUSA) + ++#if defined(GGML_USE_HIP) ++#define GGML_CACHEGEN_ACTIVE_MASK() __ballot(1) ++#else ++#define GGML_CACHEGEN_ACTIVE_MASK() __activemask() ++#endif ++ + static __device__ __forceinline__ uint32_t cachegen_read_bit( + const uint8_t * stream, uint32_t stream_bytes, uint32_t & bit) { + const uint32_t byte_index = bit >> 3; +@@ -124,7 +130,7 @@ static __global__ void cachegen_decode( float amax = fabsf(decoded_f); #pragma unroll for (uint32_t mask = 16; mask > 0; mask >>= 1) { - amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, mask, 32)); -+ amax = fmaxf(amax, __shfl_xor_sync(__activemask(), amax, mask, 32)); ++ amax = fmaxf(amax, __shfl_xor_sync(GGML_CACHEGEN_ACTIVE_MASK(), amax, mask, 32)); } const float scale = amax / 127.0f; if (lane == 0) { -@@ -138,19 +138,19 @@ static __global__ void cachegen_decode( +@@ -138,19 +144,19 @@ static __global__ void cachegen_decode( float amax = fabsf(decoded_f); #pragma unroll for (uint32_t mask = 16; mask > 0; mask >>= 1) { - amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, mask, 32)); -+ amax = fmaxf(amax, __shfl_xor_sync(__activemask(), amax, mask, 32)); ++ amax = fmaxf(amax, __shfl_xor_sync(GGML_CACHEGEN_ACTIVE_MASK(), amax, mask, 32)); } uint32_t winner = fabsf(decoded_f) == amax ? lane : 32; #pragma unroll for (uint32_t mask = 16; mask > 0; mask >>= 1) { - winner = min(winner, __shfl_xor_sync(0xFFFFFFFF, winner, mask, 32)); -+ winner = min(winner, __shfl_xor_sync(__activemask(), winner, mask, 32)); ++ winner = min(winner, __shfl_xor_sync(GGML_CACHEGEN_ACTIVE_MASK(), winner, mask, 32)); } - const float signed_max = __shfl_sync(0xFFFFFFFF, decoded_f, winner, 32); -+ const float signed_max = __shfl_sync(__activemask(), decoded_f, winner, 32); ++ const float signed_max = __shfl_sync(GGML_CACHEGEN_ACTIVE_MASK(), decoded_f, winner, 32); const float scale = signed_max / -8.0f; if (lane == 0) { *reinterpret_cast(dst + destination) = __float2half(scale); } - const float high_value = __shfl_down_sync(0xFFFFFFFF, decoded_f, 16, 32); -+ const float high_value = __shfl_down_sync(__activemask(), decoded_f, 16, 32); ++ const float high_value = __shfl_down_sync(GGML_CACHEGEN_ACTIVE_MASK(), decoded_f, 16, 32); if (lane < 16) { const float inverse = scale == 0.0f ? 0.0f : 1.0f / scale; const int low = max(0, min(15, int(decoded_f * inverse + 8.5f))); +@@ -212,4 +218,6 @@ cudaError_t ggml_cuda_cachegen_decode_launch( + return cudaGetLastError(); + } + ++#undef GGML_CACHEGEN_ACTIVE_MASK ++ + #endif // !defined(GGML_USE_MUSA) -- 2.54.0 (Apple Git-157) - From a85f8c85d3f5dca6980590fb4d3992de38dc2dca Mon Sep 17 00:00:00 2001 From: scama Date: Sun, 13 Sep 2026 18:30:39 +1000 Subject: [PATCH 04/16] fix(cachegen): map shuffle-down intrinsic for HIP --- ...da-use-native-CacheGen-shuffle-masks.patch | 20 ++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/third_party/llama.cpp/patches/0046-ggml-cuda-use-native-CacheGen-shuffle-masks.patch b/third_party/llama.cpp/patches/0046-ggml-cuda-use-native-CacheGen-shuffle-masks.patch index 96bd122998..27b5efd925 100644 --- a/third_party/llama.cpp/patches/0046-ggml-cuda-use-native-CacheGen-shuffle-masks.patch +++ b/third_party/llama.cpp/patches/0046-ggml-cuda-use-native-CacheGen-shuffle-masks.patch @@ -1,11 +1,12 @@ -From ca99c017af46be9b6cd07411ac0be33d9799d5ca Mon Sep 17 00:00:00 2001 +From 2014cb675af3629ff1bb6d61ef7885bc32b12ed2 Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Sat, 12 Sep 2026 11:49:37 +1000 Subject: [PATCH] ggml-cuda: use native CacheGen shuffle masks --- - ggml/src/ggml-cuda/cachegen.cu | 18 +++++++++++++----- - 1 file changed, 13 insertions(+), 5 deletions(-) + ggml/src/ggml-cuda/cachegen.cu | 18 +++++++++++++----- + ggml/src/ggml-cuda/vendors/hip.h | 1 + + 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-cuda/cachegen.cu b/ggml/src/ggml-cuda/cachegen.cu index 685dc586a..394c1d45e 100644 @@ -64,5 +65,18 @@ index 685dc586a..394c1d45e 100644 +#undef GGML_CACHEGEN_ACTIVE_MASK + #endif // !defined(GGML_USE_MUSA) +diff --git a/ggml/src/ggml-cuda/vendors/hip.h b/ggml/src/ggml-cuda/vendors/hip.h +index 2fc0fe9fd..d842ad754 100644 +--- a/ggml/src/ggml-cuda/vendors/hip.h ++++ b/ggml/src/ggml-cuda/vendors/hip.h +@@ -30,6 +30,7 @@ + #define CU_MEM_ACCESS_FLAGS_PROT_READWRITE hipMemAccessFlagsProtReadWrite + #define CU_CHECK(fn) {hipError_t err = fn; if(err != hipSuccess) { GGML_ABORT("HipVMM Failure: %s\n", hipGetErrorString(err)); }} + #define __shfl_sync(mask, var, laneMask, width) __shfl(var, laneMask, width) ++#define __shfl_down_sync(mask, var, delta, width) __shfl_down(var, delta, width) + #define __shfl_up_sync(mask, var, laneMask, width) __shfl_up(var, laneMask, width) + #define __shfl_xor_sync(mask, var, laneMask, width) __shfl_xor(var, laneMask, width) + #define __all_sync(mask, var) __all(var) -- 2.54.0 (Apple Git-157) + From ce0741eb51913cc5b7af0bfdacc0c8a78fe66eac Mon Sep 17 00:00:00 2001 From: scama Date: Mon, 14 Sep 2026 07:25:13 +1000 Subject: [PATCH 05/16] feat(skippy): wire bounded host-RAM L2 serving --- .../src/model/built_in_schema/declarations.rs | 2 +- crates/mesh-llm-config/src/wiring_status.rs | 8 +- .../src/inference/skippy/family_policy.rs | 1 + .../inference/skippy/resolver/resolution.rs | 13 +- .../src/inference/skippy/resolver/support.rs | 3 - .../src/inference/skippy/resolver/tests.rs | 2 + .../inference/skippy/resolver/translation.rs | 21 +- .../src/inference/skippy/resolver/types.rs | 1 + .../src/runtime/local_model_only.rs | 1 + .../config_schema_defaults_ui_reference.json | 7 + crates/skippy-cache/src/l2/mod.rs | 318 ++++++++++++- .../src/prompt_cli/stage_config.rs | 2 + crates/skippy-protocol/src/config.rs | 4 + .../src/binary_transport/stage_execution.rs | 1 + .../src/frontend/local_generation/tests.rs | 1 + .../token_generation/kv_restore.rs | 2 +- .../src/frontend/prefix_cache.rs | 2 + .../src/frontend/tests/support.rs | 1 + .../src/kv_integration/activation.rs | 1 + .../src/kv_integration/cache_affinity.rs | 1 + .../src/kv_integration/config.rs | 126 ++++- .../src/kv_integration/exact_state.rs | 27 +- .../src/kv_integration/identity.rs | 1 + .../src/kv_integration/l2_serving.rs | 437 ++++++++++++++++++ .../skippy-server/src/kv_integration/mod.rs | 32 ++ .../src/kv_integration/resident_prefix.rs | 1 + docs/USAGE.md | 2 +- docs/skippy/CONFIGURATION.md | 2 +- docs/skippy/PROMPT_CACHE.md | 21 + tools/xtask/data/console_print_allowlist.json | 8 +- website/src/docs/pages/config-reference.md | 2 +- 31 files changed, 1018 insertions(+), 33 deletions(-) create mode 100644 crates/skippy-server/src/kv_integration/l2_serving.rs diff --git a/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs b/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs index 741bd946c2..625c3edcce 100644 --- a/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs +++ b/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs @@ -342,7 +342,7 @@ fn model_fit_settings( ), basic_setting(&format!("{prefix}.kv_offload"), bool_or_auto_schema()), basic_setting(&format!("{prefix}.kv_unified"), bool_or_auto_schema()), - unwired_setting( + basic_setting( &format!("{prefix}.cache_ram_mib"), ConfigValueSchema::Integer, ), diff --git a/crates/mesh-llm-config/src/wiring_status.rs b/crates/mesh-llm-config/src/wiring_status.rs index caa7c63640..f757699335 100644 --- a/crates/mesh-llm-config/src/wiring_status.rs +++ b/crates/mesh-llm-config/src/wiring_status.rs @@ -622,10 +622,10 @@ pub const WIRING_MANIFEST: &[WiringEntry] = &[ }, WiringEntry { path: "model_fit.cache_ram_mib", - status: WiringStatus::Unwired, - owner: "PR2", - reason: "Any positive value fails at model load", - behavior: WiringBehavior::BailsDownstream, + status: WiringStatus::Wired, + owner: "n/a", + reason: "", + behavior: WiringBehavior::None, }, WiringEntry { path: "model_fit.cache_idle_slots", diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs index 7441df5e4f..ca91b85e3e 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs @@ -54,6 +54,7 @@ impl FamilyPolicy { payload: StageKvCachePayload::Auto, max_entries: bounded_entries, max_bytes, + l2_max_bytes: 0, min_tokens, shared_prefix_stride_tokens: 128, shared_prefix_record_limit: derive_shared_prefix_record_limit(bounded_entries), diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs index 834df81d14..e737503d8f 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs @@ -17,7 +17,7 @@ use super::types::{ BUILTIN_BATCH, BUILTIN_CTX_SIZE, BUILTIN_PARALLEL, BUILTIN_PREFILL_CHUNK_SIZE, BUILTIN_SAFETY_MARGIN_GB, BUILTIN_UBATCH, ResolvedHardwareConfig, ResolvedModelFitConfig, ResolvedMultimodalConfig, ResolvedSkippyConfig, ResolvedSkippyExecutionConfig, - ResolvedThroughputConfig, SkippyConfigResolveRequest, + ResolvedStageKvCache, ResolvedThroughputConfig, SkippyConfigResolveRequest, }; use crate::plugin::{ BoolOrAuto, ModelConfigDefaults, ModelConfigEntry, ModelFitConfig, ThroughputConfig, @@ -266,6 +266,16 @@ fn resolve_model_fit_config( .or(context.global_model_fit.and_then(|fit| fit.flash_attention)) .unwrap_or_else(|| effective_flash_attention(&cache_type_v)); let prefix_cache = resolve_prefix_cache(context.model_fit, context.global_model_fit)?; + let l2_max_bytes = pick_owned( + context.model_fit.and_then(|fit| fit.cache_ram_mib), + context.global_model_fit.and_then(|fit| fit.cache_ram_mib), + ) + .unwrap_or(0) + .checked_mul(1024 * 1024) + .ok_or_else(|| anyhow::anyhow!("model_fit.cache_ram_mib exceeds the byte range"))?; + if l2_max_bytes > 0 && matches!(prefix_cache, ResolvedStageKvCache::Disabled) { + anyhow::bail!("model_fit.cache_ram_mib requires prefix caching to be enabled"); + } Ok(ResolvedModelFitConfig { ctx_size, @@ -275,6 +285,7 @@ fn resolve_model_fit_config( cache_type_v, kv_cache_policy: kv.effective_policy, prefix_cache, + l2_max_bytes, kv_offload, kv_offload_resolved, kv_unified, diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/support.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/support.rs index 3b4cee8854..036412fdcb 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/support.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/support.rs @@ -68,9 +68,6 @@ pub(super) fn reject_unsupported_model_fit_controls( let Some(config) = config else { return Ok(()); }; - if config.cache_ram_mib.unwrap_or(0) > 0 { - bail!("skippy model_fit.cache_ram_mib is not supported by the pinned runtime"); - } if config.keep_tokens.unwrap_or(0) > 0 { bail!("skippy model_fit.keep_tokens is not supported by the pinned runtime"); } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs index 9193a99f52..2459917ade 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs @@ -2033,6 +2033,7 @@ fn staged_controls_propagate_into_stage_config_and_embedded_openai_args() { r#" [defaults.model_fit] prompt_cache = true +cache_ram_mib = 64 [defaults.model_fit.prefix_cache] enabled = true @@ -2080,6 +2081,7 @@ draft_max_tokens = 8 assert_eq!(kv_cache.shared_prefix_stride_tokens, 48); assert_eq!(kv_cache.shared_prefix_record_limit, 3); assert_eq!(kv_cache.payload, StageKvCachePayload::ResidentKv); + assert_eq!(kv_cache.l2_max_bytes, 64 * 1024 * 1024); let openai = resolved .to_embedded_openai_args(4096, true) diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs index be415bf193..8a79e5f3cd 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs @@ -472,23 +472,25 @@ impl ResolvedSkippyConfig { &self, family_default: Option, ) -> Result> { - match &self.model_fit.prefix_cache { - ResolvedStageKvCache::FamilyDefault => Ok(family_default), - ResolvedStageKvCache::Disabled => Ok(Some(StageKvCacheConfig { + let mut resolved = match &self.model_fit.prefix_cache { + ResolvedStageKvCache::FamilyDefault => family_default, + ResolvedStageKvCache::Disabled => Some(StageKvCacheConfig { mode: StageKvCacheMode::Disabled, payload: StageKvCachePayload::Auto, max_entries: 0, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 0, shared_prefix_stride_tokens: 0, shared_prefix_record_limit: 0, - })), + }), ResolvedStageKvCache::Explicit(template) => { let mut cache = family_default.unwrap_or(StageKvCacheConfig { mode: template.mode.clone(), payload: StageKvCachePayload::Auto, max_entries: 128, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 256, shared_prefix_stride_tokens: 128, shared_prefix_record_limit: 2, @@ -510,9 +512,18 @@ impl ResolvedSkippyConfig { if let Some(value) = template.shared_prefix_record_limit { cache.shared_prefix_record_limit = value as u64; } - Ok(Some(cache)) + Some(cache) } + }; + if self.model_fit.l2_max_bytes > 0 && resolved.is_none() { + anyhow::bail!( + "model_fit.cache_ram_mib requires an executable prefix-cache configuration" + ); + } + if let Some(cache) = resolved.as_mut() { + cache.l2_max_bytes = self.model_fit.l2_max_bytes; } + Ok(resolved) } } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs index f49d29e7e7..4e2217184b 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs @@ -74,6 +74,7 @@ pub(crate) struct ResolvedModelFitConfig { pub(crate) cache_type_v: String, pub(crate) kv_cache_policy: String, pub(crate) prefix_cache: ResolvedStageKvCache, + pub(crate) l2_max_bytes: u64, pub(crate) kv_offload: String, /// Parsed `kv_offload` for the native tri-state control. `None` covers /// both "auto" and any value that did not parse to a bool. diff --git a/crates/mesh-llm-host-runtime/src/runtime/local_model_only.rs b/crates/mesh-llm-host-runtime/src/runtime/local_model_only.rs index da4127a4c0..c6408f7d11 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local_model_only.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local_model_only.rs @@ -505,6 +505,7 @@ mod tests { payload: StageKvCachePayload::Auto, max_entries: 8, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 8, shared_prefix_stride_tokens: 8, shared_prefix_record_limit: 2, diff --git a/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_defaults_ui_reference.json b/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_defaults_ui_reference.json index ac43057d6d..cbab8a580b 100644 --- a/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_defaults_ui_reference.json +++ b/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_defaults_ui_reference.json @@ -112,6 +112,13 @@ "kind": "built_in" } }, + { + "canonical_path": "defaults.model_fit.cache_ram_mib", + "support": "supported", + "source": { + "kind": "built_in" + } + }, { "canonical_path": "defaults.model_fit.cache_type_k", "support": "supported", diff --git a/crates/skippy-cache/src/l2/mod.rs b/crates/skippy-cache/src/l2/mod.rs index 3591b589bf..5d72cc8b88 100644 --- a/crates/skippy-cache/src/l2/mod.rs +++ b/crates/skippy-cache/src/l2/mod.rs @@ -37,9 +37,8 @@ //! `CacheBytes` is a block-backed view over the shared segment storages, //! contiguous in the single-segment case. //! -//! This first slice is a standalone store with no wiring into the request -//! path; the benchmark harness drives it directly. L2 promotion/demotion -//! policy and server integration land in a later slice. +//! `skippy-server` wires this tier as a bounded mirror of repeated or +//! high-value L3 fills. An L2 hit rewarms L1. use std::{ collections::HashMap, ops::Range, @@ -49,10 +48,14 @@ use std::{ }, }; -use crate::payload::{CacheBytes, ExactStatePayloadKind}; -use crate::{HandoffManifest, segment_digest}; #[cfg(test)] -use crate::{HandoffSegmentRef, MANIFEST_VERSION, PayloadCodec, SegmentCodecIdentity}; +use crate::PayloadCodec; +use crate::payload::{CacheBytes, ExactStatePayload, ExactStatePayloadKind}; +use crate::{ + HandoffManifest, HandoffSegmentRef, MANIFEST_VERSION, SegmentCodecIdentity, segment_digest, +}; + +const DIRECT_SEGMENT_BYTES: usize = 1024 * 1024; /// Where an entry came from, for telemetry and promotion policy later. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -117,6 +120,9 @@ pub struct L2Layout { pub total_bytes: u64, pub kv_bytes: u64, pub recurrent_bytes: u64, + /// Opaque serialized runtime KV-page descriptor. The serving layer + /// validates this before importing the payload. + pub kv_desc_json: Option, /// `(segment digest, byte range within the assembled wire)` per /// manifest segment, in manifest order. Ranges concatenate to /// `0..total_bytes` exactly as the L3 manifest tiles them. @@ -150,6 +156,10 @@ impl ExactStatePayloadMirror { } } + pub fn kv_desc_json(&self) -> Option<&str> { + self.layout().kv_desc_json.as_deref() + } + /// Build a mirror from a captured L3 manifest. Callers must verify the /// payload wire against `manifest.payload_digest` — `admit` does this — /// before the mirror is stored. @@ -199,6 +209,7 @@ impl ExactStatePayloadMirror { total_bytes: manifest.total_bytes, kv_bytes: manifest.kv_bytes, recurrent_bytes: manifest.recurrent_bytes, + kv_desc_json: manifest.kv_desc_json.clone(), segments, }; Ok(match kind { @@ -574,6 +585,33 @@ fn validate_layout<'a>( } Ok(validated) } + +fn payload_wire(payload: &ExactStatePayload) -> Result<(Vec, u64, u64), L2InsertRefusal> { + let malformed = |error: anyhow::Error| L2InsertRefusal::MalformedManifest(error.to_string()); + match payload { + ExactStatePayload::FullState { bytes } => { + let wire = bytes.as_cow().map_err(malformed)?.into_owned(); + let kv_bytes = wire.len() as u64; + Ok((wire, kv_bytes, 0)) + } + ExactStatePayload::RecurrentOnly { recurrent } => { + let wire = recurrent.as_cow().map_err(malformed)?.into_owned(); + let recurrent_bytes = wire.len() as u64; + Ok((wire, 0, recurrent_bytes)) + } + ExactStatePayload::KvRecurrent { kv, recurrent } => { + let kv = kv.as_cow().map_err(malformed)?; + let recurrent = recurrent.as_cow().map_err(malformed)?; + let kv_bytes = kv.len() as u64; + let recurrent_bytes = recurrent.len() as u64; + let mut wire = Vec::with_capacity(kv.len().saturating_add(recurrent.len())); + wire.extend_from_slice(kv.as_ref()); + wire.extend_from_slice(recurrent.as_ref()); + Ok((wire, kv_bytes, recurrent_bytes)) + } + } +} + impl L2Tier { pub fn new(budget_bytes: u64) -> Self { Self { @@ -587,6 +625,64 @@ impl L2Tier { self.budget_bytes } + /// Admit an exact-state payload restored from the authoritative L3 tier. + /// + /// The payload is cut into stable one-MiB content chunks. This preserves + /// immutable sharing between related entries while keeping all hashing + /// and copying on the existing cache worker rather than the request path. + pub fn admit_payload( + &self, + cache_key: String, + token_count: u64, + expected_payload_digest: &str, + payload: &ExactStatePayload, + kv_desc_json: Option, + origin: L2Origin, + ) -> Result, L2InsertRefusal> { + let (wire, kv_bytes, recurrent_bytes) = payload_wire(payload)?; + let payload_digest = segment_digest(&wire); + if payload_digest != expected_payload_digest { + self.stats.admission_rejects.fetch_add(1, Ordering::Relaxed); + return Err(L2InsertRefusal::DigestMismatch { + expected: expected_payload_digest.to_string(), + actual: payload_digest, + }); + } + let mut manifest = HandoffManifest::new(String::new(), payload.kind().to_string()); + manifest.version = MANIFEST_VERSION; + manifest.total_bytes = wire.len() as u64; + manifest.payload_digest = expected_payload_digest.to_string(); + manifest.kv_bytes = kv_bytes; + manifest.recurrent_bytes = recurrent_bytes; + manifest.kv_desc_json = kv_desc_json; + manifest.token_count = token_count; + manifest.segments = wire + .chunks(DIRECT_SEGMENT_BYTES) + .enumerate() + .scan(0u64, |offset, (index, bytes)| { + let start = *offset; + *offset = offset.saturating_add(bytes.len() as u64); + Some(HandoffSegmentRef { + index: index as u32, + offset: start, + bytes: bytes.len() as u64, + digest: segment_digest(bytes), + codec_identity: Some(SegmentCodecIdentity::raw(bytes.len() as u64)), + meta_json: None, + }) + }) + .collect(); + let mirror = ExactStatePayloadMirror::from_manifest(&manifest)?; + self.admit( + cache_key, + token_count, + expected_payload_digest.to_string(), + &wire, + mirror, + origin, + ) + } + /// Admit an assembled entry. /// /// `wire` is the payload's concatenated L3 wire — the exact bytes whose @@ -980,6 +1076,62 @@ impl L2Tier { bytes } + /// Remove every entry that mirrors one durable payload digest. + pub fn remove_by_digest(&self, payload_digest: &str) -> Vec { + let mut inner = self.inner.lock().expect("L2 map lock poisoned"); + let keys = inner + .map + .iter() + .filter(|(_, entry)| entry.payload_digest == payload_digest) + .map(|(key, _)| key.clone()) + .collect::>(); + let mut evictions = Vec::with_capacity(keys.len()); + for key in keys { + let Some(removed) = inner.map.remove(&key) else { + continue; + }; + let mut retained = 0u64; + for digest in removed.payload.segment_digests() { + if inner + .map + .values() + .any(|other| other.payload.segment_digests().contains(&digest)) + { + retained = retained.saturating_add( + inner + .segments + .get(digest) + .map(|handle| handle.bytes.len() as u64) + .unwrap_or(0), + ); + } + } + Self::recompute_all_charges(&mut inner); + let before = inner.bytes; + self.release_entry_segments(&mut inner, &removed, &[]); + evictions.push(L2Eviction { + cache_key: key, + freed_bytes: before.saturating_sub(inner.bytes), + retained_bytes: retained, + }); + } + evictions + } + + /// Evict least-recently-used entries until physical usage is at or below + /// `target_bytes`. Targets above the configured budget are clamped. + pub fn shrink_to(&self, target_bytes: u64) -> Vec { + let mut inner = self.inner.lock().expect("L2 map lock poisoned"); + let mut journal = AdmitJournal::default(); + self.evict_to_limit( + &mut inner, + target_bytes.min(self.budget_bytes), + "", + &[], + &mut journal, + ) + } + pub fn len(&self) -> usize { self.inner.lock().expect("L2 map lock poisoned").map.len() } @@ -1214,6 +1366,7 @@ mod tests { total_bytes: len, kv_bytes: len, recurrent_bytes: 0, + kv_desc_json: None, segments: vec![(segment_digest(w), 0..len)], }, } @@ -1241,6 +1394,7 @@ mod tests { total_bytes: len, kv_bytes: len, recurrent_bytes: 0, + kv_desc_json: None, segments, }, } @@ -1280,6 +1434,94 @@ mod tests { assert_eq!(bytes.as_ref(), &w[..], "served bytes must equal the wire"); } + #[test] + fn admit_payload_round_trips_composite_state_and_descriptor() { + let tier = L2Tier::new(1 << 20); + let k = key("ns", &[1, 2, 3]); + let kv = vec![1, 2, 3, 4]; + let recurrent = vec![5, 6, 7]; + let descriptor = r#"{"token_start":0,"token_count":3}"#.to_string(); + let expected_digest = segment_digest(&[kv.as_slice(), recurrent.as_slice()].concat()); + + tier.admit_payload( + k.clone(), + 3, + &expected_digest, + &ExactStatePayload::kv_recurrent(kv.clone(), recurrent.clone()), + Some(descriptor.clone()), + L2Origin::Direct, + ) + .expect("direct payload admission must fit"); + + let hit = tier.get(&k).expect("admitted payload must hit"); + assert_eq!(hit.payload.kv_desc_json(), Some(descriptor.as_str())); + let payload = hit.to_payload(); + assert_eq!( + payload + .kv_bytes() + .expect("read KV bytes") + .expect("composite payload has KV") + .as_ref(), + kv.as_slice() + ); + assert_eq!( + payload + .recurrent_state_bytes() + .expect("read recurrent bytes") + .as_ref(), + recurrent.as_slice() + ); + } + + #[test] + fn admit_payload_round_trips_full_state() { + let tier = L2Tier::new(1 << 20); + let k = key("ns", &[9]); + let bytes = vec![7; 128]; + let expected_digest = segment_digest(&bytes); + + tier.admit_payload( + k.clone(), + 1, + &expected_digest, + &ExactStatePayload::full_state(bytes.clone()), + None, + L2Origin::Direct, + ) + .expect("full-state admission must fit"); + + let payload = tier + .get(&k) + .expect("admitted payload must hit") + .to_payload(); + assert_eq!( + payload + .full_state_bytes_timed() + .expect("read full state") + .0 + .as_ref(), + bytes.as_slice() + ); + } + + #[test] + fn direct_payload_admission_requires_the_durable_digest() { + let tier = L2Tier::new(1 << 20); + let payload = ExactStatePayload::full_state(vec![7; 128]); + let error = tier + .admit_payload( + key("ns", &[9]), + 1, + &segment_digest(b"different"), + &payload, + None, + L2Origin::FromL3, + ) + .expect_err("mismatched durable identity must be refused"); + assert!(matches!(error, L2InsertRefusal::DigestMismatch { .. })); + assert!(tier.is_empty()); + } + #[test] fn admission_digest_mismatch_refuses_and_stores_nothing() { let tier = L2Tier::new(1 << 20); @@ -1856,6 +2098,65 @@ mod tests { assert_eq!(tier.stats().segments, 0); } + #[test] + fn remove_by_digest_drops_only_matching_mirrors() { + let tier = L2Tier::new(1 << 20); + let (shared, shared_digest) = wire(64, 8); + let (other, other_digest) = wire(64, 9); + for tokens in [&[1][..], &[2][..]] { + tier.admit( + key("ns", tokens), + 1, + shared_digest.clone(), + &shared, + single_segment_mirror(&shared), + L2Origin::FromL3, + ) + .expect("shared mirror fits"); + } + let other_key = key("ns", &[3]); + tier.admit( + other_key.clone(), + 1, + other_digest, + &other, + single_segment_mirror(&other), + L2Origin::FromL3, + ) + .expect("other mirror fits"); + + let removed = tier.remove_by_digest(&shared_digest); + assert_eq!(removed.len(), 2); + assert_eq!(tier.len(), 1); + assert!(tier.get(&other_key).is_some()); + } + + #[test] + fn shrink_to_evicts_lru_until_the_target_is_met() { + let tier = L2Tier::new(256); + let mut keys = Vec::new(); + for i in 0..3i32 { + let (bytes, digest) = wire(64, i as u8 + 1); + let cache_key = key("ns", &[i]); + tier.admit( + cache_key.clone(), + 1, + digest, + &bytes, + single_segment_mirror(&bytes), + L2Origin::FromL3, + ) + .expect("entry fits"); + keys.push(cache_key); + } + assert!(tier.get(&keys[0]).is_some(), "first entry becomes hottest"); + + let evicted = tier.shrink_to(128); + assert_eq!(evicted.len(), 1); + assert_eq!(evicted[0].cache_key, keys[1]); + assert_eq!(tier.stats().bytes, 128); + } + #[test] fn identical_wire_same_key_readmit_keeps_its_own_segments() { // The original failure: re-admitting an identical wire at the same @@ -1947,6 +2248,7 @@ mod tests { total_bytes: total, kv_bytes: total, recurrent_bytes: 0, + kv_desc_json: None, segments: grown_segments, }, }; @@ -2031,6 +2333,7 @@ mod tests { total_bytes: 120, kv_bytes: 120, recurrent_bytes: 0, + kv_desc_json: None, segments: vec![(segment_digest(&x), 0..60), (segment_digest(&y), 60..120)], }, }; @@ -2138,6 +2441,7 @@ mod tests { total_bytes: 32, kv_bytes: 32, recurrent_bytes: 0, + kv_desc_json: None, segments: vec![ (segment_digest(&seg), 0..16), (segment_digest(&seg), 16..32), @@ -2287,6 +2591,7 @@ mod tests { total_bytes: 64, kv_bytes: 64, recurrent_bytes: 0, + kv_desc_json: None, segments: vec![(x_digest.clone(), 0..32), (x_digest, 32..64)], }, }; @@ -2346,6 +2651,7 @@ mod tests { total_bytes: 48, kv_bytes: 48, recurrent_bytes: 0, + kv_desc_json: None, segments: vec![(stolen, 0..24), (segment_digest(&w2[24..]), 24..48)], }, }; diff --git a/crates/skippy-prompt/src/prompt_cli/stage_config.rs b/crates/skippy-prompt/src/prompt_cli/stage_config.rs index 84964bbedc..954541104a 100644 --- a/crates/skippy-prompt/src/prompt_cli/stage_config.rs +++ b/crates/skippy-prompt/src/prompt_cli/stage_config.rs @@ -132,6 +132,7 @@ fn prompt_stage_kv_cache_config( payload: StageKvCachePayload::Auto, max_entries: 1, max_bytes: 0, + l2_max_bytes: 0, min_tokens: args.kv_page_size_tokens.max(1), shared_prefix_stride_tokens: 128, shared_prefix_record_limit: 2, @@ -145,6 +146,7 @@ fn prompt_stage_kv_cache_config( payload, max_entries: 128, max_bytes, + l2_max_bytes: 0, min_tokens: args.kv_page_size_tokens.max(1), shared_prefix_stride_tokens: 128, shared_prefix_record_limit: 2, diff --git a/crates/skippy-protocol/src/config.rs b/crates/skippy-protocol/src/config.rs index 25afdc68d0..61148b0827 100644 --- a/crates/skippy-protocol/src/config.rs +++ b/crates/skippy-protocol/src/config.rs @@ -339,6 +339,10 @@ pub struct StageKvCacheConfig { pub max_entries: usize, #[serde(default)] pub max_bytes: u64, + /// Hard byte budget for the opt-in host-RAM L2 exact-state tier. + /// Zero keeps L2 disabled. + #[serde(default)] + pub l2_max_bytes: u64, #[serde(default = "default_kv_cache_min_tokens")] pub min_tokens: u64, #[serde(default = "default_kv_cache_shared_stride_tokens")] diff --git a/crates/skippy-server/src/binary_transport/stage_execution.rs b/crates/skippy-server/src/binary_transport/stage_execution.rs index f61a716961..66faf9da68 100644 --- a/crates/skippy-server/src/binary_transport/stage_execution.rs +++ b/crates/skippy-server/src/binary_transport/stage_execution.rs @@ -976,6 +976,7 @@ pub(in crate::binary_transport) fn prefix_cache_test_config() -> StageConfig { payload: StageKvCachePayload::ResidentKv, max_entries: 8, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 256, shared_prefix_stride_tokens: 128, shared_prefix_record_limit: 2, diff --git a/crates/skippy-server/src/frontend/local_generation/tests.rs b/crates/skippy-server/src/frontend/local_generation/tests.rs index 72e6842ca2..3cba546ce4 100644 --- a/crates/skippy-server/src/frontend/local_generation/tests.rs +++ b/crates/skippy-server/src/frontend/local_generation/tests.rs @@ -147,6 +147,7 @@ fn recurrent_test_backend( payload: StageKvCachePayload::KvRecurrent, max_entries: 8, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 1, shared_prefix_stride_tokens: 1, shared_prefix_record_limit: 0, diff --git a/crates/skippy-server/src/frontend/local_generation/token_generation/kv_restore.rs b/crates/skippy-server/src/frontend/local_generation/token_generation/kv_restore.rs index 41d9c2cd95..7ddd6f074e 100644 --- a/crates/skippy-server/src/frontend/local_generation/token_generation/kv_restore.rs +++ b/crates/skippy-server/src/frontend/local_generation/token_generation/kv_restore.rs @@ -41,7 +41,7 @@ impl StageOpenAiBackend { "skippy.exact_cache.source".to_string(), json!(restored.source), ); - if restored.source == "l3" { + if restored.source != "radix" { attrs.insert( "skippy.exact_cache.fill_ms".to_string(), json!(restored.fill_ms), diff --git a/crates/skippy-server/src/frontend/prefix_cache.rs b/crates/skippy-server/src/frontend/prefix_cache.rs index 61e9ed46ca..d39695595e 100644 --- a/crates/skippy-server/src/frontend/prefix_cache.rs +++ b/crates/skippy-server/src/frontend/prefix_cache.rs @@ -1372,6 +1372,7 @@ mod tests { payload: skippy_protocol::StageKvCachePayload::ResidentKv, max_entries: 8, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 1, shared_prefix_stride_tokens: 1, shared_prefix_record_limit: 1, @@ -1449,6 +1450,7 @@ mod tests { payload: skippy_protocol::StageKvCachePayload::ResidentKv, max_entries: 8, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 1, shared_prefix_stride_tokens: 1, shared_prefix_record_limit: 1, diff --git a/crates/skippy-server/src/frontend/tests/support.rs b/crates/skippy-server/src/frontend/tests/support.rs index 1de05d1a68..15d99b9f48 100644 --- a/crates/skippy-server/src/frontend/tests/support.rs +++ b/crates/skippy-server/src/frontend/tests/support.rs @@ -47,6 +47,7 @@ pub(super) fn prefix_cache_test_config() -> StageConfig { payload: StageKvCachePayload::ResidentKv, max_entries: 8, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 256, shared_prefix_stride_tokens: 128, shared_prefix_record_limit: 2, diff --git a/crates/skippy-server/src/kv_integration/activation.rs b/crates/skippy-server/src/kv_integration/activation.rs index 3dcf3d66d8..b7f874738e 100644 --- a/crates/skippy-server/src/kv_integration/activation.rs +++ b/crates/skippy-server/src/kv_integration/activation.rs @@ -189,6 +189,7 @@ mod tests { payload: StageKvCachePayload::ResidentKv, max_entries: 8, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 256, shared_prefix_stride_tokens: 128, shared_prefix_record_limit: 2, diff --git a/crates/skippy-server/src/kv_integration/cache_affinity.rs b/crates/skippy-server/src/kv_integration/cache_affinity.rs index aa56428c2c..612164f229 100644 --- a/crates/skippy-server/src/kv_integration/cache_affinity.rs +++ b/crates/skippy-server/src/kv_integration/cache_affinity.rs @@ -109,6 +109,7 @@ mod tests { payload: StageKvCachePayload::ResidentKv, max_entries: 8, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 64, shared_prefix_stride_tokens: 32, shared_prefix_record_limit: 2, diff --git a/crates/skippy-server/src/kv_integration/config.rs b/crates/skippy-server/src/kv_integration/config.rs index 968a4e5231..fb574935bb 100644 --- a/crates/skippy-server/src/kv_integration/config.rs +++ b/crates/skippy-server/src/kv_integration/config.rs @@ -133,7 +133,16 @@ impl KvStageIntegration { return Ok(None); } let l3_manager = manager()?; - let durable_payload = l3_manager.as_ref().map(|_| { + if cache_config.l2_max_bytes > 0 && l3_manager.is_none() { + let _ = mesh_llm_events::emit_event(OutputEvent::Warning { + message: "Skippy L2 host-RAM cache disabled for this model stage".to_string(), + context: Some(format!( + "stage_id={} model_id={} reason=L2 requires an active L3 cache", + config.stage_id, config.model_id + )), + }); + } + let durable_payload = l3_manager.is_some().then(|| { if payload == StagePrefixCachePayload::ResidentKv && dense_without_recurrent { // Resident KV stays the in-process fast path. Dense families // export KV pages with an empty recurrent snapshot for L3; @@ -147,6 +156,17 @@ impl KvStageIntegration { .zip(durable_payload) .map(|(manager, payload)| l3_tier_for_manager(config, payload, manager)) .transpose()?; + let l2 = l3 + .as_ref() + .filter(|_| cache_config.l2_max_bytes > 0) + .and(durable_payload) + .map(|payload| { + super::l2_serving::StageL2::new( + cache_config.l2_max_bytes, + numerical_model_identity_for_stage(config), + exact_state_identity_for_stage(config, l3_payload_kind(payload)), + ) + }); // FullState is architecture-neutral: the native runtime serializes the // complete session state for both dense and recurrent model families. if matches!(model_capability, ModelKvCapability::KnownRecurrent) { @@ -178,6 +198,7 @@ impl KvStageIntegration { let worker_radix = radix.clone(); let worker_exact_blobs = exact_blobs.clone(); let worker_l3 = l3.clone(); + let worker_l2 = l2.clone(); let inflight_records: Arc>> = l3.as_ref().map_or_else( || Arc::new(Mutex::new(BTreeSet::new())), |tier| tier.manager().record_claims(tier.state_identity()), @@ -222,6 +243,7 @@ impl KvStageIntegration { &worker_exact_blobs, exact_max_entries, exact_byte_limits, + worker_l2.as_ref(), worker_l3.as_deref(), pending, ) @@ -280,6 +302,7 @@ impl KvStageIntegration { split_prefill_tokens: Arc::new(Mutex::new(BTreeMap::new())), kv_lifecycle_observer: observer, exact_state_record_queue_bytes, + l2, l3, inflight_fills, dense_without_recurrent, @@ -535,6 +558,7 @@ fn store_exact_radix_record( blobs: &Mutex, max_entries: usize, limits: ExactStateByteLimits, + l2: Option<&super::l2_serving::StageL2>, l3: Option<&L3Tier>, pending: PendingExactStateRecord, ) -> Result<()> { @@ -543,7 +567,9 @@ fn store_exact_radix_record( // or failing disk must not fail the in-memory record. The refusal reason // lands in the tier's status; one warning per process keeps a full disk // from flooding the log. - if let Some(l3) = l3 { + if pending.write_through_l3 + && let Some(l3) = l3 + { let kv_desc_json = pending .extra .kv_desc @@ -573,6 +599,15 @@ fn store_exact_radix_record( } } } + if let (Some(l2), Some(payload_digest)) = (l2, pending.l2_promotion_digest.as_deref()) { + let _ = l2.promote( + &pending.namespace, + &pending.token_ids, + payload_digest, + &pending.payload, + &pending.extra, + ); + } let logical_bytes = pending.payload.byte_len(); let (payload, _) = pending.payload.dedupe_into( &mut blobs @@ -785,6 +820,7 @@ fn effective_cache_config(config: &StageConfig) -> Option { payload, max_entries, max_bytes, + l2_max_bytes: 0, min_tokens, shared_prefix_stride_tokens, shared_prefix_record_limit, @@ -833,6 +869,8 @@ mod tests { namespace: "model".to_string(), token_ids: tokens.to_vec(), l3_fill_claim: None, + write_through_l3: true, + l2_promotion_digest: None, } } @@ -847,6 +885,7 @@ mod tests { 1, limits(0, 0), None, + None, pending("first", &[1, 2], b"aaaabbbb"), ) .unwrap(); @@ -856,6 +895,7 @@ mod tests { 1, limits(0, 0), None, + None, pending("second", &[1, 3], b"aaaacccc"), ) .unwrap(); @@ -875,6 +915,28 @@ mod tests { ); } + #[test] + fn selected_l3_fill_promotes_to_l2_on_the_record_worker_path() { + let radix = Mutex::new(UnifiedRadixCache::new()); + let blobs = Mutex::new(CacheBlobStore::new(4)); + let l2 = super::super::l2_serving::StageL2::new( + 1 << 20, + "model-identity".to_string(), + "state-identity".to_string(), + ); + let bytes = b"filled-state"; + let mut record = pending("filled", &[1, 2], bytes); + record.write_through_l3 = false; + record.l2_promotion_digest = Some(skippy_cache::segment_digest(bytes)); + + store_exact_radix_record(&radix, &blobs, 1, limits(0, 0), Some(&l2), None, record).unwrap(); + + let stats = l2.stats(); + assert_eq!(stats.entries, 1); + assert_eq!(stats.inserts, 1); + assert_eq!(stats.logical_bytes, bytes.len() as u64); + } + #[test] fn invalid_exact_radix_key_releases_deduped_payload() { let radix = Mutex::new(UnifiedRadixCache::new()); @@ -886,6 +948,7 @@ mod tests { 1, limits(0, 0), None, + None, pending("empty", &[], b"aaaabbbb"), ) .unwrap_err(); @@ -914,6 +977,7 @@ mod tests { &blobs, 1, limits(0, 0), + None, Some(&tier), pending("first", &[1, 2], b"first-exact-state"), ) @@ -923,6 +987,7 @@ mod tests { &blobs, 1, limits(0, 0), + None, Some(&tier), pending("second", &[1, 3], b"second-exact-state"), ) @@ -1076,6 +1141,7 @@ mod tests { 8, limits(4, 1024), None, + None, pending(page_id, &tokens, bytes), ) .unwrap(); @@ -1107,6 +1173,7 @@ mod tests { 2, limits(4, 1_024), None, + None, pending(page_id, &tokens, bytes), ) .unwrap(); @@ -1135,6 +1202,7 @@ mod tests { 8, limits(4, 8), None, + None, pending(page_id, &tokens, bytes), ) .unwrap(); @@ -1161,6 +1229,7 @@ mod tests { 8, limits(2, 4), None, + None, pending("checkpoint", &[1, 2], b"aaaabbbb"), ) .unwrap(); @@ -1402,6 +1471,58 @@ mod tests { assert_eq!(invalid.exact_state_payload(), None); } + #[test] + fn cache_ram_budget_enables_stage_l2_with_disk_authority() { + let mut config = enabled_auto_config("future/model"); + config + .kv_cache + .as_mut() + .expect("enabled cache config") + .l2_max_bytes = 64 * 1024 * 1024; + let root = tempfile::tempdir().unwrap(); + let manager = L3CacheManager::acquire(root.path(), StoreLimits::new(1 << 30, 0)).unwrap(); + + let kv = KvStageIntegration::from_loaded_model_with_l3_manager( + &config, + Some(ModelStateKind::Dense), + Some(manager), + None, + ) + .unwrap() + .expect("prefix cache should remain enabled"); + + assert!(kv.l2.is_some()); + assert!(kv.l3.is_some()); + let attrs = kv.attrs().into_iter().collect::>(); + assert_eq!(attrs["skippy.kv.l2.enabled"], serde_json::json!(true)); + assert_eq!( + attrs["skippy.kv.l2.budget_bytes"], + serde_json::json!(64 * 1024 * 1024u64) + ); + } + + #[test] + fn cache_ram_budget_does_not_create_an_authority_free_l2() { + let mut config = enabled_auto_config("future/model"); + config + .kv_cache + .as_mut() + .expect("enabled cache config") + .l2_max_bytes = 64 * 1024 * 1024; + + let kv = KvStageIntegration::from_loaded_model_with_l3_manager( + &config, + Some(ModelStateKind::Dense), + None, + None, + ) + .unwrap() + .expect("L1 prefix cache should remain enabled"); + + assert!(kv.l2.is_none()); + assert!(kv.l3.is_none()); + } + #[test] fn parses_cache_mode_and_payload_aliases() { assert_eq!( @@ -1483,6 +1604,7 @@ mod tests { payload: StageKvCachePayload::Auto, max_entries: 512, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 1, shared_prefix_stride_tokens: 1, shared_prefix_record_limit: 1, diff --git a/crates/skippy-server/src/kv_integration/exact_state.rs b/crates/skippy-server/src/kv_integration/exact_state.rs index 1d5047cef6..41d137755c 100644 --- a/crates/skippy-server/src/kv_integration/exact_state.rs +++ b/crates/skippy-server/src/kv_integration/exact_state.rs @@ -88,7 +88,7 @@ impl KvStageIntegration { (lookup, entries) }; let Some(lookup) = lookup else { - // Radix miss: the durable tier may still hold this prefix. + // The durable tiers may still hold this prefix. // Runs inside the restore transaction, so a failed import // rolls the lane back exactly as a radix restore would. if let Some(restored) = @@ -395,6 +395,8 @@ impl KvStageIntegration { namespace: identity.namespace.clone(), token_ids: identity.token_ids.clone(), l3_fill_claim: None, + write_through_l3: true, + l2_promotion_digest: None, }) { ExactStateRecordAdmission::Queued => { // Recording owns the radix/blob locks while it hashes a potentially @@ -462,6 +464,13 @@ impl KvStageIntegration { // recorded the reason for the status surface. Ok(None) | Err(_) => return Ok(None), }; + // L3 remains the authority for L2: locate and validate the current + // manifest identity before a host-RAM mirror may serve the request. + if let Some(restored) = + self.restore_from_l2(runtime, session_id, identity, &location, lookup_started)? + { + return Ok(Some(restored)); + } // Segment and manifest digests intentionally deduplicate bytes across // numerical states. A fill claim must not: one state's fill cannot // warm another state's radix namespace, even when their payload bytes @@ -511,8 +520,14 @@ impl KvStageIntegration { // A load failure (corrupt segment, now quarantined) is a miss, not a // request failure. Import failures below do propagate: the transaction // rolls the lane back and the caller falls back to cold prefill. - let Ok(fill) = l3.load(location) else { - return Ok(None); + let fill = match l3.load(location) { + Ok(fill) => fill, + Err(_) => { + if let Some(l2) = &self.l2 { + l2.invalidate_digest(&location.manifest_key); + } + return Ok(None); + } }; if fill.payload.byte_len() == 0 { return Ok(None); @@ -607,6 +622,10 @@ impl KvStageIntegration { // Re-warm the RAM tier off the request path. A drop is fine: the // disk copy stays authoritative. The fill claim rides along so the // worker releases it only once the entry is radix-resident. + let l2_promotion_digest = self.l2.as_ref().and_then(|l2| { + l2.consider_l3_fill(&location.manifest_key, token_count, fill.payload.byte_len()) + .then(|| location.manifest_key.clone()) + }); let admission = self.enqueue_exact_state_record(PendingExactStateRecord { page_id: identity.page_id.clone(), payload: fill.payload, @@ -614,6 +633,8 @@ impl KvStageIntegration { namespace: identity.namespace.clone(), token_ids: identity.token_ids[..token_count as usize].to_vec(), l3_fill_claim: Some(l3_fill_claim_key(l3, location)), + write_through_l3: false, + l2_promotion_digest, }); let rewarm_enqueued = matches!(admission, ExactStateRecordAdmission::Queued); Ok(Some(ExactStateRestore { diff --git a/crates/skippy-server/src/kv_integration/identity.rs b/crates/skippy-server/src/kv_integration/identity.rs index 71fa2bb0f9..16f5e045f2 100644 --- a/crates/skippy-server/src/kv_integration/identity.rs +++ b/crates/skippy-server/src/kv_integration/identity.rs @@ -225,6 +225,7 @@ mod tests { payload: StageKvCachePayload::ResidentKv, max_entries: 8, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 64, shared_prefix_stride_tokens: 32, shared_prefix_record_limit: 2, diff --git a/crates/skippy-server/src/kv_integration/l2_serving.rs b/crates/skippy-server/src/kv_integration/l2_serving.rs new file mode 100644 index 0000000000..1ae1e84796 --- /dev/null +++ b/crates/skippy-server/src/kv_integration/l2_serving.rs @@ -0,0 +1,437 @@ +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + time::{Duration, Instant}, +}; + +use anyhow::{Context, Result}; +use skippy_cache::{L2Hit, L2Origin, L2Stats, L2Tier, l2_cache_key}; +use skippy_runtime::RuntimeKvPageDesc; + +use crate::runtime_state::RuntimeState; + +use super::{ + ExactStateExtra, ExactStateRecordAdmission, ExactStateRestore, KvStageIntegration, + PendingExactStateRecord, PrefillKvIdentity, StagePrefixCachePayload, + records::add_reconstruct_stats, +}; + +const PROMOTION_MAX_BYTES: u64 = 64 * 1024 * 1024; +const RESTORE_WORTH_TOKENS: u64 = 4_096; +const SECOND_HIT_WINDOW: Duration = Duration::from_secs(10 * 60); +const MAX_PROMOTION_OBSERVATIONS: usize = 4_096; + +#[derive(Clone, Copy)] +struct PromotionObservation { + first_seen: Instant, + hits: u8, + terminal: bool, +} + +/// Stage-scoped view of the host-RAM tier. The identities are fixed when the +/// model loads, so every lookup and promotion uses the same numerical boundary +/// as L3 without exposing model or prompt fingerprints in telemetry. +#[derive(Clone)] +pub(crate) struct StageL2 { + tier: Arc, + model_identity: String, + state_identity: String, + promotions: Arc>>, +} + +impl StageL2 { + pub(crate) fn new(budget_bytes: u64, model_identity: String, state_identity: String) -> Self { + Self { + tier: Arc::new(L2Tier::new(budget_bytes)), + model_identity, + state_identity, + promotions: Arc::new(Mutex::new(HashMap::new())), + } + } + + pub(crate) fn get( + &self, + identity: &PrefillKvIdentity, + location: &skippy_cache::L3Location, + ) -> Option<(String, L2Hit)> { + let token_count = usize::try_from(location.token_count).ok()?; + let tokens = identity.token_ids.get(..token_count)?; + let key = self.cache_key(&identity.namespace, tokens); + let hit = self.tier.get(&key)?; + if hit.token_count != location.token_count || hit.payload_digest != location.manifest_key { + self.remove(&key); + return None; + } + Some((key, hit)) + } + + pub(crate) fn remove(&self, key: &str) { + let _ = self.tier.remove(key); + } + + pub(crate) fn invalidate_digest(&self, payload_digest: &str) { + let _ = self.tier.remove_by_digest(payload_digest); + } + + pub(crate) fn consider_l3_fill( + &self, + manifest_key: &str, + token_count: u64, + payload_bytes: u64, + ) -> bool { + self.should_promote(manifest_key, token_count, payload_bytes) + } + + pub(crate) fn promote( + &self, + namespace: &str, + tokens: &[i32], + expected_payload_digest: &str, + payload: &skippy_cache::ExactStatePayload, + extra: &ExactStateExtra, + ) -> bool { + let key = self.cache_key(namespace, tokens); + let kv_desc_json = extra + .kv_desc + .as_ref() + .and_then(|desc| serde_json::to_string(desc).ok()); + let admitted = self + .tier + .admit_payload( + key, + tokens.len() as u64, + expected_payload_digest, + payload, + kv_desc_json, + L2Origin::FromL3, + ) + .is_ok(); + if let Some(observation) = self + .promotions + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get_mut(expected_payload_digest) + { + observation.terminal = true; + } + admitted + } + + pub(crate) fn stats(&self) -> L2Stats { + self.tier.stats() + } + + fn cache_key(&self, namespace: &str, tokens: &[i32]) -> String { + l2_cache_key( + &self.model_identity, + &self.state_identity, + namespace, + tokens, + ) + } + + fn should_promote(&self, manifest_key: &str, token_count: u64, payload_bytes: u64) -> bool { + let now = Instant::now(); + let mut observations = self + .promotions + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + observations.retain(|_, observation| { + now.saturating_duration_since(observation.first_seen) <= SECOND_HIT_WINDOW + }); + if observations.len() >= MAX_PROMOTION_OBSERVATIONS + && !observations.contains_key(manifest_key) + && let Some(oldest) = observations + .iter() + .min_by(|left, right| { + left.1 + .first_seen + .cmp(&right.1.first_seen) + .then_with(|| left.0.cmp(right.0)) + }) + .map(|(key, _)| key.clone()) + { + observations.remove(&oldest); + } + let observation = + observations + .entry(manifest_key.to_string()) + .or_insert(PromotionObservation { + first_seen: now, + hits: 0, + terminal: false, + }); + if observation.terminal { + return false; + } + observation.hits = observation.hits.saturating_add(1); + if payload_bytes > PROMOTION_MAX_BYTES { + observation.terminal = true; + return false; + } + // Keep the observation eligible until the worker has attempted the + // promotion. If the bounded worker queue drops this fill, the next L3 + // restore may retry instead of suppressing the entry for ten minutes. + token_count >= RESTORE_WORTH_TOKENS || observation.hits >= 2 + } +} + +impl KvStageIntegration { + pub(super) fn restore_from_l2( + &self, + runtime: &mut RuntimeState, + session_id: &str, + identity: &PrefillKvIdentity, + location: &skippy_cache::L3Location, + lookup_started: Instant, + ) -> Result> { + let Some(l2) = &self.l2 else { + return Ok(None); + }; + let Some((key, hit)) = l2.get(identity, location) else { + return Ok(None); + }; + let token_count = hit.token_count; + if token_count == 0 || token_count != location.token_count { + l2.remove(&key); + return Ok(None); + } + let payload = hit.to_payload(); + let kv_desc = match hit.payload.kv_desc_json() { + Some(json) => match serde_json::from_str::(json) { + Ok(desc) => Some(desc), + Err(_) => { + l2.remove(&key); + return Ok(None); + } + }, + None => None, + }; + let fill_ms = lookup_started.elapsed().as_secs_f64() * 1000.0; + let mut reconstruct_ms = 0.0; + let mut reconstruct_bytes = 0u64; + let mut reconstruct_blocks = 0usize; + let mut kv_import_ms = 0.0; + let mut recurrent_import_ms = 0.0; + let mut deterministic_failure = false; + let restore = (|| -> Result { + match payload.kind().into() { + StagePrefixCachePayload::FullState => { + let (bytes, stats) = payload + .full_state_bytes_timed() + .context("reconstruct L2 full-state payload") + .map_err(|error| mark_failure(&mut deterministic_failure, error))?; + if bytes.is_empty() { + deterministic_failure = true; + anyhow::bail!("L2 full-state payload is empty"); + } + add_reconstruct_stats( + &mut reconstruct_ms, + &mut reconstruct_bytes, + &mut reconstruct_blocks, + stats, + ); + let started = Instant::now(); + runtime.import_full_state_for_token_count( + session_id, + bytes.as_ref(), + token_count, + )?; + kv_import_ms = started.elapsed().as_secs_f64() * 1000.0; + } + StagePrefixCachePayload::KvRecurrent => { + if let Some((kv, stats)) = payload + .kv_bytes_timed() + .context("reconstruct L2 KV payload") + .map_err(|error| mark_failure(&mut deterministic_failure, error))? + { + add_reconstruct_stats( + &mut reconstruct_ms, + &mut reconstruct_bytes, + &mut reconstruct_blocks, + stats, + ); + match kv_desc.as_ref() { + Some(desc) => { + desc.validate_payload(kv.len()).map_err(|error| { + mark_failure(&mut deterministic_failure, error) + })?; + if desc.token_start != 0 || desc.token_count != token_count { + deterministic_failure = true; + anyhow::bail!("L2 KV page token range mismatch"); + } + let started = Instant::now(); + runtime.import_kv_page(session_id, desc, kv.as_ref())?; + kv_import_ms = started.elapsed().as_secs_f64() * 1000.0; + } + None if !kv.is_empty() => { + deterministic_failure = true; + anyhow::bail!("L2 KV payload is missing its descriptor"); + } + None => {} + } + } + let (recurrent, stats) = payload + .recurrent_state_bytes_timed() + .context("reconstruct L2 recurrent payload") + .map_err(|error| mark_failure(&mut deterministic_failure, error))?; + if recurrent.is_empty() && !self.dense_without_recurrent { + deterministic_failure = true; + anyhow::bail!("L2 recurrent-state payload is empty"); + } + add_reconstruct_stats( + &mut reconstruct_ms, + &mut reconstruct_bytes, + &mut reconstruct_blocks, + stats, + ); + let started = Instant::now(); + if recurrent.is_empty() { + runtime.set_session_position(session_id, token_count)?; + } else { + runtime.import_recurrent_state_for_token_count( + session_id, + recurrent.as_ref(), + token_count, + )?; + } + recurrent_import_ms = started.elapsed().as_secs_f64() * 1000.0; + } + StagePrefixCachePayload::Disabled | StagePrefixCachePayload::ResidentKv => { + return Ok(false); + } + } + Ok(true) + })(); + let restored = match restore { + Ok(restored) => restored, + Err(error) => { + if deterministic_failure { + l2.remove(&key); + } + return Err(error); + } + }; + if !restored { + return Ok(None); + } + let logical_bytes = payload.byte_len(); + let payload_kind = payload.kind(); + let rewarm_enqueued = self.try_begin_record(&identity.page_id) + && matches!( + self.enqueue_exact_state_record(PendingExactStateRecord { + page_id: identity.page_id.clone(), + payload, + extra: ExactStateExtra { kv_desc }, + namespace: identity.namespace.clone(), + token_ids: identity.token_ids[..token_count as usize].to_vec(), + l3_fill_claim: None, + write_through_l3: false, + l2_promotion_digest: None, + }), + ExactStateRecordAdmission::Queued + ); + Ok(Some(ExactStateRestore { + page_id: identity.page_id.clone(), + token_count: token_count as usize, + payload_kind, + logical_bytes, + entries: l2.stats().entries as usize, + reconstruct_ms, + reconstruct_bytes, + reconstruct_blocks, + lookup_ms: fill_ms, + kv_import_ms, + recurrent_import_ms, + source: "l2", + fill_ms, + rewarm_enqueued, + })) + } +} + +fn mark_failure(deterministic: &mut bool, error: anyhow::Error) -> anyhow::Error { + *deterministic = true; + error +} + +#[cfg(test)] +mod tests { + use super::*; + + fn identity(tokens: Vec) -> PrefillKvIdentity { + PrefillKvIdentity { + identity: crate::kv_proto::PageIdentity::default(), + page_id: "page".to_string(), + namespace: "namespace".to_string(), + token_ids: tokens, + } + } + + fn location(token_count: u64, manifest_key: String) -> skippy_cache::L3Location { + skippy_cache::L3Location { + namespace_key: "namespace-key".to_string(), + prefix_key: "prefix-key".to_string(), + token_count, + manifest_key, + kv_desc_json: None, + kv_bytes: 0, + native_kv_passthrough: false, + } + } + + #[test] + fn second_l3_fill_promotes_and_serves_the_located_prefix() { + let l2 = StageL2::new(1 << 20, "model".to_string(), "state".to_string()); + let bytes = vec![7; 128]; + let digest = skippy_cache::segment_digest(&bytes); + let payload = skippy_cache::ExactStatePayload::full_state(bytes); + let identity = identity(vec![1, 2, 3, 4]); + let location = location(3, digest.clone()); + + assert!(!l2.consider_l3_fill(&digest, 3, payload.byte_len())); + assert!(l2.consider_l3_fill(&digest, 3, payload.byte_len())); + assert!(l2.promote( + &identity.namespace, + &identity.token_ids[..3], + &digest, + &payload, + &ExactStateExtra { kv_desc: None }, + )); + assert!(!l2.consider_l3_fill(&digest, 3, payload.byte_len())); + + let (_, hit) = l2 + .get(&identity, &location) + .expect("longer query should hit the located L3 prefix mirror"); + assert_eq!(hit.token_count, 3); + assert_eq!(hit.payload_digest, digest); + } + + #[test] + fn restart_class_prefix_promotes_on_first_fill_but_oversized_payload_does_not() { + let l2 = StageL2::new(1 << 20, "model".to_string(), "state".to_string()); + assert!(l2.consider_l3_fill("restart", RESTORE_WORTH_TOKENS, 1 << 20)); + assert!(!l2.consider_l3_fill("oversized", RESTORE_WORTH_TOKENS, PROMOTION_MAX_BYTES + 1,)); + assert!(!l2.consider_l3_fill("oversized", RESTORE_WORTH_TOKENS, PROMOTION_MAX_BYTES,)); + } + + #[test] + fn a_changed_durable_digest_invalidates_the_mirror() { + let l2 = StageL2::new(1 << 20, "model".to_string(), "state".to_string()); + let bytes = vec![7; 128]; + let digest = skippy_cache::segment_digest(&bytes); + let payload = skippy_cache::ExactStatePayload::full_state(bytes); + let identity = identity(vec![1, 2, 3]); + assert!(l2.promote( + &identity.namespace, + &identity.token_ids, + &digest, + &payload, + &ExactStateExtra { kv_desc: None }, + )); + + let changed = location(3, skippy_cache::segment_digest(b"changed")); + assert!(l2.get(&identity, &changed).is_none()); + assert_eq!(l2.stats().entries, 0); + } +} diff --git a/crates/skippy-server/src/kv_integration/mod.rs b/crates/skippy-server/src/kv_integration/mod.rs index 93db8aa401..26123ad372 100644 --- a/crates/skippy-server/src/kv_integration/mod.rs +++ b/crates/skippy-server/src/kv_integration/mod.rs @@ -27,6 +27,7 @@ mod cache_affinity; mod config; mod exact_state; mod identity; +mod l2_serving; pub mod lifecycle; mod model_capability; mod output_tokens; @@ -185,6 +186,9 @@ pub struct KvStageIntegration { /// lets one multi-GiB export sit next to another and doubles the RAM the /// cache can pin behind a request. pub(crate) exact_state_record_queue_bytes: Arc, + /// Optional bounded host-RAM tier. Qualified L3 fills enter L2; an L2 hit + /// promotes back into L1 through the existing worker. + pub(crate) l2: Option, /// Durable L3 floor under the radix cache: exact-state records write /// through to it on the worker, and radix misses fill back from it. pub(crate) l3: Option>, @@ -229,6 +233,12 @@ pub(crate) struct PendingExactStateRecord { /// so requests arriving during the asynchronous re-warm prefill normally /// instead of duplicating the disk read. pub(crate) l3_fill_claim: Option, + /// Only freshly exported request state writes through. Tier fills that + /// merely re-warm L1 must not rewrite their existing durable entry. + pub(crate) write_through_l3: bool, + /// Durable manifest digest to mirror into L2 on this worker job. `None` + /// leaves the payload out of L2. + pub(crate) l2_promotion_digest: Option, } #[derive(Debug)] @@ -793,6 +803,11 @@ impl KvStageIntegration { Err(std::sync::TryLockError::WouldBlock) => None, }; let radix = radix_stats.unwrap_or_default(); + let l2 = self + .l2 + .as_ref() + .map(|tier| tier.stats()) + .unwrap_or_default(); let activations = self .activations .lock() @@ -951,6 +966,21 @@ impl KvStageIntegration { "skippy.exact_cache.max_entries", json!(self.exact_max_entries), ), + ("skippy.kv.l2.enabled", json!(self.l2.is_some())), + ("skippy.kv.l2.budget_bytes", json!(l2.budget_bytes)), + ("skippy.kv.l2.bytes", json!(l2.bytes)), + ("skippy.kv.l2.logical_bytes", json!(l2.logical_bytes)), + ("skippy.kv.l2.entries", json!(l2.entries)), + ("skippy.kv.l2.segments", json!(l2.segments)), + ("skippy.kv.l2.hits", json!(l2.hits)), + ("skippy.kv.l2.misses", json!(l2.misses)), + ("skippy.kv.l2.inserts", json!(l2.inserts)), + ("skippy.kv.l2.evictions", json!(l2.evictions)), + ( + "skippy.kv.l2.admission_rejects", + json!(l2.admission_rejects), + ), + ("skippy.kv.l2.refused_bytes", json!(l2.refused_bytes)), ( "skippy.kv.output_token_entries", json!(output_token_entries), @@ -1140,6 +1170,8 @@ mod exact_state_record_queue_tests { namespace: "test".to_string(), token_ids: vec![1], l3_fill_claim: None, + write_through_l3: true, + l2_promotion_digest: None, } } diff --git a/crates/skippy-server/src/kv_integration/resident_prefix.rs b/crates/skippy-server/src/kv_integration/resident_prefix.rs index cf35b4c2d0..05a4dd75cc 100644 --- a/crates/skippy-server/src/kv_integration/resident_prefix.rs +++ b/crates/skippy-server/src/kv_integration/resident_prefix.rs @@ -824,6 +824,7 @@ mod proactive_eviction_tests { payload: StageKvCachePayload::ResidentKv, max_entries: 4, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 1, shared_prefix_stride_tokens: 1, shared_prefix_record_limit: 1, diff --git a/docs/USAGE.md b/docs/USAGE.md index 23da197d01..e13044376c 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -468,7 +468,7 @@ kv_cache_policy = "balanced" # macro preset: auto quality balanced saver # explicit cache_type_k/v always wins over preset kv_offload = "auto" # bool or "auto" — KV residency / offload policy kv_unified = "auto" # bool or "auto" — unified KV layout (schema-reserved) -cache_ram_mib = 0 # byte cap for KV cache in MiB; 0 = no cap (schema-reserved) +cache_ram_mib = 0 # host-RAM L2 budget in MiB; 0 = disabled; requires L3 cache_idle_slots = 0 # idle slot retention count (schema-reserved) prompt_cache = "auto" # bool or "auto" — reuse previous prompt KV swa_full = false # sliding-window attention (model-family specific) diff --git a/docs/skippy/CONFIGURATION.md b/docs/skippy/CONFIGURATION.md index a210406aa2..cd3f28c851 100644 --- a/docs/skippy/CONFIGURATION.md +++ b/docs/skippy/CONFIGURATION.md @@ -91,7 +91,7 @@ website configuration reference, with the same `Wiring status`. | 5.1 | KV cache policy preset | `model_fit.kv_cache_policy` | P0 | `plugin/config.rs` | policy expander into cache_type_k, cache_type_v, kv_offload, cache_ram_mib | single-stage, staged | restart/reload only | mesh policy default | enum auto, quality, balanced, saver; explicit cache_type_k/v wins over preset expansion | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | quality=f16/f16 with no forced RAM cap; balanced=preserve runtime defaults; saver=prefer lower-memory dtypes plus offload warning if unsupported; auto=family or topology policy decides | wired | | 5.1 | KV offload | `model_fit.kv_offload` | P1 | `plugin/config.rs` | SkippyModelLoadOptions, StageConfig, native `skippy_runtime_config.kv_offload` tri-state | single-stage, staged | restart/reload only | backend runtime default or kv_cache_policy expansion | boolean or auto; auto preserves llama.cpp's derived `offload_kqv` default | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; capture tests in `crates/skippy-runtime/src/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | controls KV residency and offload policy | wired | | 5.1 | Unified KV cache | `model_fit.kv_unified` | P1 | `plugin/config.rs` | SkippyModelLoadOptions, StageConfig, native `skippy_runtime_config.kv_unified` tri-state | single-stage, staged | restart/reload only | backend runtime default | boolean or auto; recurrent/hybrid architectures still force this true natively regardless of the requested value | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; capture tests in `crates/skippy-runtime/src/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | advanced cache layout flag | wired | -| 5.1 | Cache RAM budget | `model_fit.cache_ram_mib` | P1 | `plugin/config.rs` | Schema-reserved only; resolver fail-closed today | single-stage, staged | restart/reload only | unset by default | integer >= 0 MiB; zero or unset means no forced cap; not executable today | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | caps cache memory when runtime path supports it; schema-reserved until a real executable surface exists | unwired | +| 5.1 | Cache RAM budget | `model_fit.cache_ram_mib` | P1 | `plugin/config.rs`, `skippy-server`, `skippy-cache` | `StageKvCacheConfig.l2_max_bytes` and bounded host-RAM L2 exact-state tier | single-stage, staged | restart/reload only | disabled by default | integer >= 0 MiB; zero or unset disables L2; positive values require prefix caching and active L3 | `#model-fit-context-and-kv-cache` | resolver propagation tests in `crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs`; promotion and admission tests in `skippy-server` and `skippy-cache`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | repeated or high-value L3 fills promote to L2; L2 hits restore before disk and asynchronously rewarm L1 | wired | | 5.1 | Cache idle slots | `model_fit.cache_idle_slots` | P1 | `plugin/config.rs` | SkippyModelLoadOptions, StageConfig, `skippy-server::RuntimeState::max_idle_sessions` idle-pool bound | single-stage, staged | restart/reload only | runtime default (unbounded, capped only by `lane_count`) | integer >= 0; bounds the idle session pool size | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; `crates/skippy-server/src/runtime_state.rs` and `crates/skippy-server/src/runtime_state/lane_lifecycle.rs` unit tests; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | bounds retained idle sessions so `drop_session_timed` discards lanes past the configured cap instead of growing the pool unbounded | wired | | 5.1 | Prompt cache / cache prompt | `model_fit.prompt_cache` | P1 | `plugin/config.rs` | cache config or request adapter defaults when executable | single-stage, staged | restart/reload only | disabled unless operator enables reuse | boolean only; reject when selected runtime does not expose prompt-cache behavior | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | request prompt reuse toggle | wired | | 5.1 | Prefix cache enable | `model_fit.prefix_cache.enabled` | P1 | `plugin/config.rs` | cache config or SKIPPY_PREFIX_CACHE bridge | single-stage, staged | restart/reload only | cache env default or disabled | boolean only; when false, ignore remaining prefix_cache fields | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | top-level exact-prefix cache toggle | wired | diff --git a/docs/skippy/PROMPT_CACHE.md b/docs/skippy/PROMPT_CACHE.md index b17f9a204c..b5356722c9 100644 --- a/docs/skippy/PROMPT_CACHE.md +++ b/docs/skippy/PROMPT_CACHE.md @@ -68,6 +68,27 @@ limit. Both limits are reported on `stage.openai_generation_summary` as `skippy.exact_cache.max_bytes` and `skippy.exact_cache.hard_max_bytes`. +## Host-RAM L2 Cache + +Set `model_fit.cache_ram_mib` to a positive MiB value to enable the bounded +host-RAM exact-state tier for that model. The default value, `0`, leaves L2 +disabled. Prefix caching and the node-local L3 cache must also be enabled. + +Exact-state lookup proceeds from the in-process radix cache (L1), to host RAM +(L2), then to the node-local disk cache (L3). L3 remains authoritative: the +server locates the current durable manifest before serving an L2 mirror and +requires the mirror digest to match it. The cache worker promotes an L3 fill +on its second hit within ten minutes, or on the first hit for prefixes of at +least 4,096 tokens. Payloads larger than 64 MiB stay out of L2. A verified L2 +hit restores the request immediately and queues the same payload to rewarm L1. +Rewarm records do not rewrite an existing L3 entry. Unloading the stage drops +its L2 tier. + +`stage.openai_generation_summary` reports `skippy.kv.l2.enabled`, budget and +resident byte counts, logical bytes, entries, segments, hits, misses, inserts, +evictions, and admission refusals. Exact-hit telemetry identifies the restore +source as `l2` and includes fill time and whether the L1 rewarm was queued. + ## mesh-llm Defaults mesh-llm wires Skippy prefix cache through family policy. For supported model diff --git a/tools/xtask/data/console_print_allowlist.json b/tools/xtask/data/console_print_allowlist.json index 94c3d891da..6f9e188845 100644 --- a/tools/xtask/data/console_print_allowlist.json +++ b/tools/xtask/data/console_print_allowlist.json @@ -4251,19 +4251,19 @@ ], "crates/skippy-prompt/src/prompt_cli/stage_config.rs": [ { - "line": 358, + "line": 360, "macro_name": "eprintln!" }, { - "line": 398, + "line": 400, "macro_name": "eprintln!" }, { - "line": 464, + "line": 466, "macro_name": "eprintln!" }, { - "line": 478, + "line": 480, "macro_name": "eprint!" } ], diff --git a/website/src/docs/pages/config-reference.md b/website/src/docs/pages/config-reference.md index cd5cdd9c56..91630c91c5 100644 --- a/website/src/docs/pages/config-reference.md +++ b/website/src/docs/pages/config-reference.md @@ -143,7 +143,7 @@ for the activity policy and privacy boundary. | `model_fit.kv_cache_policy` | enum | `balanced` (default), `auto`, `quality`, `saver`; expands into cache dtypes | both | model reload | wired | none | | `model_fit.kv_offload` | bool-or-`auto` | `auto` | both | model reload | wired | none | | `model_fit.kv_unified` | bool-or-`auto` | `auto` | both | model reload | wired (recurrent/hybrid architectures still force this true natively) | none | -| `model_fit.cache_ram_mib` | integer | unset (no cap) | both | model reload | unwired (any positive value fails at model load) | none | +| `model_fit.cache_ram_mib` | integer | `0`/unset = host-RAM L2 disabled | both | model reload | wired; requires prefix caching and active L3 | none | | `model_fit.cache_idle_slots` | integer | unset uses the runtime lane count; `0` drops every reset lane, positive values cap retained idle sessions | both | model reload | wired | none | | `model_fit.prompt_cache` | bool-or-`auto` | `auto` | both | model reload | wired | none | | `model_fit.prefix_cache.enabled` | boolean | unset uses family defaults; `false` disables | both | model reload | wired | none | From 589df86ef70b4aa67cee8376be5a833e64ceea54 Mon Sep 17 00:00:00 2001 From: scama Date: Mon, 14 Sep 2026 07:55:28 +1000 Subject: [PATCH 06/16] feat(skippy): wire benefit admission into L3 serving --- crates/skippy-cache/src/l3.rs | 37 ++- crates/skippy-cache/src/manager.rs | 268 ++++++++++++++++++ crates/skippy-cache/src/policy/mod.rs | 1 + crates/skippy-cache/src/tier.rs | 206 +++++++++++++- .../src/frontend/generation/queue.rs | 16 ++ .../local_generation/token_generation.rs | 6 +- .../token_generation/exact_state_recording.rs | 6 +- .../token_generation/kv_restore.rs | 18 +- .../src/kv_integration/config.rs | 4 +- .../src/kv_integration/exact_state.rs | 111 +++++++- .../src/kv_integration/l2_serving.rs | 1 + .../skippy-server/src/kv_integration/mod.rs | 4 + docs/skippy/KV_CACHE_DISK.md | 2 + docs/skippy/PROMPT_CACHE.md | 14 + 14 files changed, 665 insertions(+), 29 deletions(-) diff --git a/crates/skippy-cache/src/l3.rs b/crates/skippy-cache/src/l3.rs index 353a5fa35f..e8a5498ac1 100644 --- a/crates/skippy-cache/src/l3.rs +++ b/crates/skippy-cache/src/l3.rs @@ -1571,7 +1571,7 @@ impl HandoffSegmentStore { } } - fn is_pinned(&self, payload_digest: &str) -> bool { + pub(crate) fn is_pinned(&self, payload_digest: &str) -> bool { self.pins .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) @@ -1579,6 +1579,41 @@ impl HandoffSegmentStore { .is_some_and(|count| *count > 0) } + /// Remove the requested inactive manifests and collect physical objects + /// that become unreferenced. Missing and pinned keys are skipped. This is + /// the commit side of an external eviction policy; the store's own LRU + /// reservation remains the hard-budget fallback. + pub(crate) fn evict_manifest_keys(&self, keys: &[String]) -> Result> { + if keys.is_empty() { + return Ok(Vec::new()); + } + self.invalidate_usage(); + let mut removed = Vec::new(); + for key in keys { + if self.is_pinned(key) { + continue; + } + let path = self.manifest_path(key); + match fs::remove_file(&path) { + Ok(()) => { + self.packed.remove_manifest_index(key)?; + self.evicted_manifests.fetch_add(1, Ordering::Relaxed); + removed.push(key.clone()); + } + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => { + return Err(error) + .with_context(|| format!("failed to evict selected manifest {key}")); + } + } + } + if !removed.is_empty() { + self.remove_dangling_prefix_links()?; + self.collect_unreferenced_segments()?; + } + Ok(removed) + } + /// What the store holds right now, for the status contract. pub fn usage(&self) -> Result { let limits = self.limits(); diff --git a/crates/skippy-cache/src/manager.rs b/crates/skippy-cache/src/manager.rs index 5a6a844695..7fddb777c0 100644 --- a/crates/skippy-cache/src/manager.rs +++ b/crates/skippy-cache/src/manager.rs @@ -15,6 +15,10 @@ use serde::Serialize; use crate::{ l3::{HandoffSegmentStore, StoreLimits, StoreReconciliation, StoreUsage, WriteRefusal}, + policy::{ + AdmissionDecisionKind, BenefitPolicy, CostSample, EntryKey, PolicyConfig, PolicyEntryState, + SegmentId, + }, tier::L3Tier, }; @@ -32,6 +36,10 @@ pub(crate) struct L3Activity { pub(crate) geometry_rejected: AtomicU64, pub(crate) bytes_read: AtomicU64, pub(crate) bytes_written: AtomicU64, + pub(crate) benefit_probation: AtomicU64, + pub(crate) benefit_persist: AtomicU64, + pub(crate) benefit_lru_fallback: AtomicU64, + pub(crate) benefit_evictions: AtomicU64, last_error: Mutex>, } @@ -53,6 +61,10 @@ impl L3Activity { corrupt_entries: usage.map_or(0, |usage| usage.quarantined_objects), bytes_read: self.bytes_read.load(Ordering::Relaxed), bytes_written: self.bytes_written.load(Ordering::Relaxed), + benefit_probation: self.benefit_probation.load(Ordering::Relaxed), + benefit_persist: self.benefit_persist.load(Ordering::Relaxed), + benefit_lru_fallback: self.benefit_lru_fallback.load(Ordering::Relaxed), + benefit_evictions: self.benefit_evictions.load(Ordering::Relaxed), geometry_rejected: self.geometry_rejected.load(Ordering::Relaxed), last_error: self .last_error @@ -74,6 +86,10 @@ pub struct L3ActivitySnapshot { pub corrupt_entries: u64, pub bytes_read: u64, pub bytes_written: u64, + pub benefit_probation: u64, + pub benefit_persist: u64, + pub benefit_lru_fallback: u64, + pub benefit_evictions: u64, pub geometry_rejected: u64, pub last_error: Option, } @@ -134,6 +150,34 @@ struct L3ManagerInner { transitions: Mutex>, operations: RwLock<()>, reconciliation: StoreReconciliation, + benefit_admission: Mutex, +} + +#[derive(Debug)] +struct L3BenefitAdmission { + policy: BenefitPolicy, + manifests: BTreeMap, + restore_cost_ewma: BTreeMap, +} + +pub(crate) struct BenefitRestoreObservation { + pub(crate) state_identity: String, + pub(crate) key: EntryKey, + pub(crate) manifest: String, + pub(crate) exclusive_bytes: u64, + pub(crate) shared: Vec<(SegmentId, u64)>, + pub(crate) cold_prefill_cost: Option, + pub(crate) restore_cost: f64, +} + +impl Default for L3BenefitAdmission { + fn default() -> Self { + Self { + policy: BenefitPolicy::new(PolicyConfig::default()), + manifests: BTreeMap::new(), + restore_cost_ewma: BTreeMap::new(), + } + } } /// The single physical owner of a node-local L3 root. @@ -185,11 +229,235 @@ impl L3CacheManager { transitions: Mutex::new(VecDeque::new()), operations: RwLock::new(()), reconciliation, + benefit_admission: Mutex::new(L3BenefitAdmission::default()), }); managers.insert(root, Arc::downgrade(&inner)); Ok(Self { inner }) } + pub(crate) fn benefit_tracks(&self, key: EntryKey) -> bool { + self.inner + .benefit_admission + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .policy + .entry(key) + .is_some() + } + + /// Observe a candidate or recurrence. `None` deliberately selects the + /// existing LRU write-through path: timing coverage is allowed to be + /// sparse without making durable caching disappear. + pub(crate) fn benefit_should_persist( + &self, + key: EntryKey, + exclusive_bytes: u64, + shared: Vec<(SegmentId, u64)>, + cost: Option, + ) -> bool { + let Some(cost) = cost.filter(CostSample::is_valid) else { + self.inner + .activity + .benefit_lru_fallback + .fetch_add(1, Ordering::Relaxed); + return true; + }; + let mut state = self + .inner + .benefit_admission + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.policy.entry(key).is_some() { + let decision = state.policy.record_hit(key, cost); + let persist = decision.is_some_and(|decision| { + matches!( + decision.kind, + AdmissionDecisionKind::Promote | AdmissionDecisionKind::AdmitPersist + ) + }) || state + .policy + .entry(key) + .is_some_and(|entry| entry.state == PolicyEntryState::Admitted); + self.inner + .activity + .benefit_persist + .fetch_add(u64::from(persist), Ordering::Relaxed); + return persist; + } + let decision = state + .policy + .consider_admission(key, exclusive_bytes, shared, cost, &[]); + let persist = matches!(decision.kind, AdmissionDecisionKind::AdmitPersist); + if persist { + self.inner + .activity + .benefit_persist + .fetch_add(1, Ordering::Relaxed); + } else if matches!(decision.kind, AdmissionDecisionKind::AdmitProbation) { + self.inner + .activity + .benefit_probation + .fetch_add(1, Ordering::Relaxed); + } + persist + } + + pub(crate) fn benefit_candidate_cost( + &self, + state_identity: &str, + cold_prefill_cost: Option, + ) -> Option { + let cold_prefill_cost = + cold_prefill_cost.filter(|cost| cost.is_finite() && *cost >= 0.0)?; + let state = self + .inner + .benefit_admission + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let restore_cost = *state.restore_cost_ewma.get(state_identity)?; + Some(CostSample { + cold_prefill_cost, + restore_cost, + }) + } + + pub(crate) fn benefit_observe_restore(&self, observation: BenefitRestoreObservation) { + let Some(cold_prefill_cost) = observation + .cold_prefill_cost + .filter(|cost| cost.is_finite() && *cost >= 0.0) + else { + return; + }; + if !observation.restore_cost.is_finite() || observation.restore_cost < 0.0 { + return; + } + let cost = CostSample { + cold_prefill_cost, + restore_cost: observation.restore_cost, + }; + let mut state = self + .inner + .benefit_admission + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state + .restore_cost_ewma + .entry(observation.state_identity) + .and_modify(|estimate| *estimate = *estimate * 0.8 + observation.restore_cost * 0.2) + .or_insert(observation.restore_cost); + if state.policy.entry(observation.key).is_none() { + let decision = state.policy.consider_admission( + observation.key, + observation.exclusive_bytes, + observation.shared, + cost, + &[], + ); + if !matches!(decision.kind, AdmissionDecisionKind::Reject) { + state + .manifests + .insert(observation.key, observation.manifest); + } + } + let _ = state.policy.record_hit(observation.key, cost); + } + + pub(crate) fn benefit_observe_memory_hit( + &self, + state_identity: &str, + key: EntryKey, + cold_prefill_cost: Option, + ) -> bool { + let Some(cost) = self.benefit_candidate_cost(state_identity, cold_prefill_cost) else { + return false; + }; + let mut state = self + .inner + .benefit_admission + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let already_persisted = state.manifests.contains_key(&key); + state.policy.record_hit(key, cost).is_some_and(|decision| { + matches!(decision.kind, AdmissionDecisionKind::Promote) && !already_persisted + }) + } + + pub(crate) fn benefit_prepare_write( + &self, + candidate: EntryKey, + growth_bytes: u64, + ) -> Result { + let usage = self.inner.store.usage()?.used_bytes; + let budget = self.inner.store.limits().budget_bytes; + if budget == 0 { + return Ok(true); + } + let bytes_to_free = usage.saturating_add(growth_bytes).saturating_sub(budget); + if bytes_to_free == 0 { + return Ok(true); + } + let mut state = self + .inner + .benefit_admission + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state.policy.observe_pressure(usage as f64 / budget as f64); + let mut pinned = state + .manifests + .iter() + .filter_map(|(key, manifest)| self.inner.store.is_pinned(manifest).then_some(*key)) + .collect::>(); + pinned.extend( + state + .policy + .entries + .keys() + .filter(|key| **key != candidate && !state.manifests.contains_key(key)) + .copied(), + ); + let victims = state + .policy + .choose_victims(bytes_to_free, &pinned) + .into_iter() + .filter_map(|(key, verdict)| { + matches!(verdict, crate::policy::EvictionVerdict::Evict).then_some(key) + }) + .collect::>(); + if victims.contains(&candidate) { + let _ = state.policy.remove(candidate, &pinned); + return Ok(false); + } + let manifests = victims + .iter() + .filter_map(|key| state.manifests.get(key).cloned()) + .collect::>(); + let removed = self.inner.store.evict_manifest_keys(&manifests)?; + self.inner + .activity + .benefit_evictions + .fetch_add(removed.len() as u64, Ordering::Relaxed); + for key in victims { + if state + .manifests + .get(&key) + .is_some_and(|manifest| removed.contains(manifest)) + { + state.manifests.remove(&key); + let _ = state.policy.remove(key, &pinned); + } + } + Ok(true) + } + + pub(crate) fn benefit_record_manifest(&self, key: EntryKey, manifest: String) { + self.inner + .benefit_admission + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .manifests + .insert(key, manifest); + } + pub fn tier(&self, state_identity: String, segment_bytes: usize) -> L3Tier { self.tier_for_model(state_identity.clone(), state_identity, segment_bytes) } diff --git a/crates/skippy-cache/src/policy/mod.rs b/crates/skippy-cache/src/policy/mod.rs index 0ce74a9bee..e58482e8ef 100644 --- a/crates/skippy-cache/src/policy/mod.rs +++ b/crates/skippy-cache/src/policy/mod.rs @@ -239,6 +239,7 @@ pub struct RemovalOutcome { /// The policy engine. Owns per-entry statistics and the shared-segment ledger; /// the caller drives it from cache events. +#[derive(Debug)] pub struct BenefitPolicy { pub(crate) config: PolicyConfig, pub(crate) entries: BTreeMap, diff --git a/crates/skippy-cache/src/tier.rs b/crates/skippy-cache/src/tier.rs index 589585cb55..963b82f11d 100644 --- a/crates/skippy-cache/src/tier.rs +++ b/crates/skippy-cache/src/tier.rs @@ -17,8 +17,11 @@ use crate::l3::{ HandoffManifest, HandoffSegmentRef, HandoffSegmentStore, MANIFEST_VERSION, PayloadGeometry, SegmentCodecIdentity, StoreLimits, StoreUsage, segment_digest, }; -use crate::manager::{L3ActivitySnapshot, L3CacheManager, L3EffectiveStatus}; +use crate::manager::{ + BenefitRestoreObservation, L3ActivitySnapshot, L3CacheManager, L3EffectiveStatus, +}; use crate::payload::{ExactStatePayload, ExactStatePayloadKind}; +use crate::policy::{CostSample, EntryKey, SegmentId}; /// Everything the status contract needs from the tier, in one read. #[derive(Debug, Clone, PartialEq, Eq, Serialize)] @@ -56,6 +59,20 @@ pub fn l3_namespace_key(namespace: &str) -> String { format!("blake3:{}", hasher.finalize().to_hex()) } +fn benefit_entry_key(namespace: &str, token_ids: &[i32]) -> EntryKey { + let digest = blake3::hash(l3_prefix_key(namespace, token_ids).as_bytes()); + u64::from_le_bytes(digest.as_bytes()[..8].try_into().expect("eight-byte key")) +} + +fn benefit_segment_id(digest: &str) -> SegmentId { + let digest = blake3::hash(digest.as_bytes()); + u64::from_le_bytes( + digest.as_bytes()[..8] + .try_into() + .expect("eight-byte segment id"), + ) +} + /// A located entry: the cheap index-probe result, addressing exactly one /// recorded manifest. Splitting locate from load lets callers single-flight /// the expensive load on the entry itself (namespace + recorded length + @@ -210,6 +227,66 @@ impl L3Tier { self.manager.activity_snapshot() } + /// Whether this prefix has an in-memory probation/admission history. + /// Callers use this to re-export an L1 hit only when the hit can promote + /// a probationary entry into durable storage. + pub fn benefit_tracks_prefix(&self, namespace: &str, token_ids: &[i32]) -> bool { + self.manager + .benefit_tracks(benefit_entry_key(namespace, token_ids)) + } + + /// Build a comparable serving cost once this stage has observed at least + /// one real L3 restore. Until then callers fall back to LRU write-through. + pub fn benefit_candidate_cost(&self, cold_prefill_cost: Option) -> Option { + self.manager + .benefit_candidate_cost(&self.state_identity, cold_prefill_cost) + } + + /// Feed an executed durable restore into both the stage restore-cost EWMA + /// and the per-entry reuse history. + pub fn benefit_observe_l3_restore( + &self, + namespace: &str, + token_ids: &[i32], + location: &L3Location, + cold_prefill_cost: Option, + restore_cost: f64, + ) { + let Ok(manifest) = self.store().load_manifest(&location.manifest_key) else { + return; + }; + let shared = manifest + .segments + .iter() + .map(|segment| (benefit_segment_id(&segment.digest), segment.bytes)) + .collect(); + self.manager + .benefit_observe_restore(BenefitRestoreObservation { + state_identity: self.state_identity.clone(), + key: benefit_entry_key(namespace, token_ids), + manifest: location.manifest_key.clone(), + exclusive_bytes: 0, + shared, + cold_prefill_cost, + restore_cost, + }); + } + + /// Record an L1 recurrence for a probationary durable candidate. `true` + /// means the caller should enqueue its existing RAM payload for L3. + pub fn benefit_observe_memory_hit( + &self, + namespace: &str, + token_ids: &[i32], + cold_prefill_cost: Option, + ) -> bool { + self.manager.benefit_observe_memory_hit( + &self.state_identity, + benefit_entry_key(namespace, token_ids), + cold_prefill_cost, + ) + } + /// The full status snapshot. One pass over the manifests; safe to call /// while serving, since it takes no lock a request path holds. pub fn status(&self) -> Result { @@ -243,8 +320,25 @@ impl L3Tier { kv_desc_json: Option, geometry: Option<&PayloadGeometry>, ) -> Result { + self.spill_with_cost(namespace, token_ids, payload, kv_desc_json, geometry, None)? + .context("LRU fallback unexpectedly declined an L3 spill") + } + + /// Offer a serving-path spill to benefit admission. A measured cost keeps + /// a first-seen entry in RAM probation and returns `None`; recurrence can + /// promote it. Missing measurements preserve the established LRU + /// write-through behavior. + pub fn spill_with_cost( + &self, + namespace: &str, + token_ids: &[i32], + payload: &ExactStatePayload, + kv_desc_json: Option, + geometry: Option<&PayloadGeometry>, + cost: Option, + ) -> Result> { let _operation = self.manager.operation_guard(); - let result = self.spill_inner(namespace, token_ids, payload, kv_desc_json, geometry); + let result = self.spill_inner(namespace, token_ids, payload, kv_desc_json, geometry, cost); if let Err(error) = &result { self.manager.activity_counters().record_error(error); } @@ -258,7 +352,8 @@ impl L3Tier { payload: &ExactStatePayload, kv_desc_json: Option, geometry: Option<&PayloadGeometry>, - ) -> Result { + cost: Option, + ) -> Result> { if payload.byte_len() == 0 { bail!( "refusing to spill an empty exact-state payload: no state component was exported" @@ -404,6 +499,40 @@ impl L3Tier { Ok(&wire[start..end]) }) .collect::>>()?; + manifest.segments = cuts + .iter() + .zip(&segment_slices) + .enumerate() + .map(|(index, (cut, bytes))| HandoffSegmentRef { + index: index as u32, + offset: cut.offset, + bytes: cut.len, + digest: segment_digest(bytes), + codec_identity: Some(cut.representation.identity(cut.len)), + meta_json: (!cut.label.is_empty()).then(|| cut.label.clone()), + }) + .collect(); + let exclusive_bytes = serde_json::to_vec(&manifest) + .map(|bytes| bytes.len() as u64) + .unwrap_or_default(); + let benefit_key = benefit_entry_key(namespace, token_ids); + let shared = manifest + .segments + .iter() + .map(|segment| (benefit_segment_id(&segment.digest), segment.bytes)) + .collect::>(); + if !self + .manager + .benefit_should_persist(benefit_key, exclusive_bytes, shared, cost) + { + return Ok(None); + } + if !self.manager.benefit_prepare_write( + benefit_key, + (wire.len() as u64).saturating_add(exclusive_bytes), + )? { + return Ok(None); + } let stored_segments = match self.store().try_put_segments(&segment_slices) { Ok(Ok(stored)) => stored, Ok(Err(refusal)) => { @@ -420,18 +549,10 @@ impl L3Tier { // these segments are unreferenced, and an eviction triggered by // another writer would collect them mid-build. let mut held = Vec::with_capacity(stored_segments.len()); - for ((index, cut), stored) in cuts.into_iter().enumerate().zip(stored_segments) { + for stored in stored_segments { if stored.put.new { new_bytes = new_bytes.saturating_add(stored.put.bytes); } - manifest.segments.push(HandoffSegmentRef { - index: index as u32, - offset: cut.offset, - bytes: cut.len, - digest: stored.digest.clone(), - codec_identity: Some(cut.representation.identity(cut.len)), - meta_json: (!cut.label.is_empty()).then_some(cut.label), - }); held.push(stored); } // Pin before publishing the manifest so another stage cannot evict @@ -471,7 +592,9 @@ impl L3Tier { .bytes_written .fetch_add(new_bytes, Ordering::Relaxed); self.manager.record_successful_write(); - Ok(payload_digest) + self.manager + .benefit_record_manifest(benefit_key, payload_digest.clone()); + Ok(Some(payload_digest)) } /// Locate the longest recorded prefix of the query, mirroring the radix @@ -773,6 +896,63 @@ mod tests { assert_eq!(identity.calibration_digest, None); } + #[test] + fn measured_candidates_wait_in_probation_until_reused() { + let tier = tier("benefit-probation", "blake3:benefit"); + let seed_tokens = tokens(4); + let seed_payload = ExactStatePayload::full_state(vec![7; 128]); + let seed = tier + .spill("ns", &seed_tokens, &seed_payload, None, None) + .expect("seed LRU spill"); + let seed_location = tier + .locate_longest("ns", &seed_tokens, 8) + .expect("seed locate") + .expect("seed location"); + assert_eq!(seed_location.manifest_key, seed); + tier.benefit_observe_l3_restore("ns", &seed_tokens, &seed_location, Some(400.0), 100.0); + + let candidate_tokens = tokens(8); + let candidate_payload = ExactStatePayload::full_state(vec![9; 256]); + let cost = tier.benefit_candidate_cost(Some(400.0)); + assert!( + tier.spill_with_cost( + "ns", + &candidate_tokens, + &candidate_payload, + None, + None, + cost, + ) + .expect("first offer") + .is_none() + ); + assert!(tier.benefit_tracks_prefix("ns", &candidate_tokens)); + assert!( + tier.spill_with_cost( + "ns", + &candidate_tokens, + &candidate_payload, + None, + None, + cost, + ) + .expect("first reuse") + .is_none() + ); + let admitted = tier + .spill_with_cost( + "ns", + &candidate_tokens, + &candidate_payload, + None, + None, + cost, + ) + .expect("second reuse") + .expect("candidate promoted"); + assert!(tier.store().load_manifest(&admitted).is_ok()); + } + #[test] fn native_passthrough_preserves_supported_runtime_kv_representations() { // ggml type ids used by the runtime: F32, F16, Q8_0 and Q4_0. The diff --git a/crates/skippy-server/src/frontend/generation/queue.rs b/crates/skippy-server/src/frontend/generation/queue.rs index af96b5905c..71b3a9c6ef 100644 --- a/crates/skippy-server/src/frontend/generation/queue.rs +++ b/crates/skippy-server/src/frontend/generation/queue.rs @@ -453,6 +453,14 @@ impl GenerationServiceEstimator { predicted_wait_ms_for_state(&state, self.concurrency.load(Ordering::Acquire)) } + pub(in crate::frontend) fn estimated_prefill_ms(&self, tokens: usize) -> Option { + let state = self.state.lock().ok()?; + state + .prefill_ms_per_token_ewma + .map(|per_token| per_token * tokens as f64) + .filter(|cost| cost.is_finite() && *cost >= 0.0) + } + pub(in crate::frontend) fn set_concurrency(&self, concurrency: usize) { self.concurrency .store(concurrency.max(1), Ordering::Release); @@ -950,4 +958,12 @@ mod service_estimator_tests { let samples = VecDeque::from([1.0, 1.0, 1.0, 10.0]); assert_eq!(conservative_ms_per_token(Some(1.5), &samples), Some(10.0)); } + + #[test] + fn prefill_estimate_uses_observed_stage_rate() { + let estimator = GenerationServiceEstimator::new(1); + assert_eq!(estimator.estimated_prefill_ms(128), None); + estimator.observe_completed(GenerationAdmissionWork::new(100, 0), 25.0, 0.0); + assert_eq!(estimator.estimated_prefill_ms(40), Some(10.0)); + } } diff --git a/crates/skippy-server/src/frontend/local_generation/token_generation.rs b/crates/skippy-server/src/frontend/local_generation/token_generation.rs index ca29dadfb9..920aadccb5 100644 --- a/crates/skippy-server/src/frontend/local_generation/token_generation.rs +++ b/crates/skippy-server/src/frontend/local_generation/token_generation.rs @@ -1255,7 +1255,11 @@ impl StageOpenAiBackend { cache_operation, )?; ensure_cache_operation_active(runtime, session_id, cache_operation)?; - kv.record_exact_state(runtime, session_id, &identity) + let cold_prefill_cost = self + .generation_service_estimator + .estimated_prefill_ms(boundary); + let l3_cost = kv.l3_benefit_cost(cold_prefill_cost); + kv.record_exact_state_with_cost(runtime, session_id, &identity, l3_cost) .map_err(openai_backend_error)?; prefill_cache_chunks( runtime, diff --git a/crates/skippy-server/src/frontend/local_generation/token_generation/exact_state_recording.rs b/crates/skippy-server/src/frontend/local_generation/token_generation/exact_state_recording.rs index 8952210279..a76f42be76 100644 --- a/crates/skippy-server/src/frontend/local_generation/token_generation/exact_state_recording.rs +++ b/crates/skippy-server/src/frontend/local_generation/token_generation/exact_state_recording.rs @@ -120,7 +120,11 @@ impl StageOpenAiBackend { let base = self.local_kv_message_base(session_id, ids); let identity = kv.prefill_identity(&self.config, &base, 0, checkpoint_tokens); - match kv.record_exact_state(runtime, session_id, &identity) { + let cold_prefill_cost = self + .generation_service_estimator + .estimated_prefill_ms(checkpoint_tokens.len()); + let l3_cost = kv.l3_benefit_cost(cold_prefill_cost); + match kv.record_exact_state_with_cost(runtime, session_id, &identity, l3_cost) { Ok(Some(record)) => { let mut attrs = self.openai_attrs(ids); attrs.insert( diff --git a/crates/skippy-server/src/frontend/local_generation/token_generation/kv_restore.rs b/crates/skippy-server/src/frontend/local_generation/token_generation/kv_restore.rs index 7ddd6f074e..4c1bb233fa 100644 --- a/crates/skippy-server/src/frontend/local_generation/token_generation/kv_restore.rs +++ b/crates/skippy-server/src/frontend/local_generation/token_generation/kv_restore.rs @@ -18,7 +18,15 @@ impl StageOpenAiBackend { let mut protected_resident_seq_id = None; let kv_restore_started = Instant::now(); let kv_restore_timer = self.telemetry.is_debug_enabled().then(PhaseTimer::start); - match kv.restore_exact_state(runtime, session_id, identities) { + let estimated_cold_prefill_ms = self + .generation_service_estimator + .estimated_prefill_ms(prefill_tokens.len()); + match kv.restore_exact_state_with_cold_cost( + runtime, + session_id, + identities, + estimated_cold_prefill_ms, + ) { Ok(Some(restored)) => { restored_prefill = true; cache_stats.status = "hit"; @@ -195,7 +203,13 @@ impl StageOpenAiBackend { ) { let base = self.local_kv_message_base(session_id, ids); let exact_identity = kv.prefill_identity(&self.config, &base, 0, prefill_tokens); - if let Ok(Some(record)) = kv.record_exact_state(runtime, session_id, &exact_identity) { + let cold_prefill_cost = self + .generation_service_estimator + .estimated_prefill_ms(prefill_tokens.len()); + let l3_cost = kv.l3_benefit_cost(cold_prefill_cost); + if let Ok(Some(record)) = + kv.record_exact_state_with_cost(runtime, session_id, &exact_identity, l3_cost) + { resident_recorded_pages = resident_recorded_pages.saturating_add(1); let mut attrs = self.openai_attrs(ids); attrs.insert( diff --git a/crates/skippy-server/src/kv_integration/config.rs b/crates/skippy-server/src/kv_integration/config.rs index fb574935bb..54e7158b10 100644 --- a/crates/skippy-server/src/kv_integration/config.rs +++ b/crates/skippy-server/src/kv_integration/config.rs @@ -580,12 +580,13 @@ fn store_exact_radix_record( .kv_desc .as_ref() .and_then(|desc| kv_page_geometry(desc, pending.payload.byte_len())); - let spill = l3.spill( + let spill = l3.spill_with_cost( &pending.namespace, &pending.token_ids, &pending.payload, kv_desc_json, geometry.as_ref(), + pending.l3_cost, ); emit_l3_state_transitions(l3); if let Err(error) = spill { @@ -871,6 +872,7 @@ mod tests { l3_fill_claim: None, write_through_l3: true, l2_promotion_digest: None, + l3_cost: None, } } diff --git a/crates/skippy-server/src/kv_integration/exact_state.rs b/crates/skippy-server/src/kv_integration/exact_state.rs index 41d137755c..baeb519915 100644 --- a/crates/skippy-server/src/kv_integration/exact_state.rs +++ b/crates/skippy-server/src/kv_integration/exact_state.rs @@ -42,14 +42,33 @@ fn preflight_native_kv_location( } impl KvStageIntegration { + pub(crate) fn l3_benefit_cost( + &self, + cold_prefill_cost: Option, + ) -> Option { + self.l3 + .as_ref() + .and_then(|l3| l3.benefit_candidate_cost(cold_prefill_cost)) + } + pub fn restore_exact_state( &self, runtime: &mut RuntimeState, session_id: &str, identities: &[PrefillKvIdentity], + ) -> Result> { + self.restore_exact_state_with_cold_cost(runtime, session_id, identities, None) + } + + pub fn restore_exact_state_with_cold_cost( + &self, + runtime: &mut RuntimeState, + session_id: &str, + identities: &[PrefillKvIdentity], + cold_prefill_cost: Option, ) -> Result> { runtime.restore_transaction(session_id, |runtime| { - self.restore_exact_state_inner(runtime, session_id, identities) + self.restore_exact_state_inner(runtime, session_id, identities, cold_prefill_cost) }) } @@ -58,6 +77,7 @@ impl KvStageIntegration { runtime: &mut RuntimeState, session_id: &str, identities: &[PrefillKvIdentity], + cold_prefill_cost: Option, ) -> Result> { if !self.should_lookup() || self.exact_state_payload().is_none() { return Ok(None); @@ -91,9 +111,13 @@ impl KvStageIntegration { // The durable tiers may still hold this prefix. // Runs inside the restore transaction, so a failed import // rolls the lane back exactly as a radix restore would. - if let Some(restored) = - self.restore_from_l3(runtime, session_id, identity, lookup_started)? - { + if let Some(restored) = self.restore_from_l3( + runtime, + session_id, + identity, + lookup_started, + cold_prefill_cost, + )? { return Ok(Some(restored)); } continue; @@ -242,8 +266,41 @@ impl KvStageIntegration { drop(lease); continue; } + let promote_to_l3 = self.l3.as_ref().is_some_and(|l3| { + l3.benefit_observe_memory_hit( + &identity.namespace, + &lookup.stored_tokens, + cold_prefill_cost, + ) + }); + let page_id = lookup.value.page_id.clone(); + let promotion_payload = promote_to_l3.then(|| lookup.value.payload.clone()); + let promotion_extra = promote_to_l3.then(|| lookup.value.extra.clone()); + let stored_tokens = lookup.stored_tokens.clone(); + drop(lease); + let promotion_enqueued = if let (Some(payload), Some(extra)) = + (promotion_payload, promotion_extra) + && self.try_begin_record(&page_id) + { + matches!( + self.enqueue_exact_state_record(PendingExactStateRecord { + page_id: page_id.clone(), + payload, + extra, + namespace: identity.namespace.clone(), + token_ids: stored_tokens, + l3_fill_claim: None, + write_through_l3: true, + l2_promotion_digest: None, + l3_cost: None, + }), + ExactStateRecordAdmission::Queued + ) + } else { + false + }; let restored = ExactStateRestore { - page_id: lookup.value.page_id, + page_id, token_count: token_count as usize, payload_kind: lookup.value.payload.kind(), logical_bytes: lookup.logical_bytes, @@ -256,9 +313,8 @@ impl KvStageIntegration { recurrent_import_ms, source: "radix", fill_ms: 0.0, - rewarm_enqueued: false, + rewarm_enqueued: promotion_enqueued, }; - drop(lease); return Ok(Some(restored)); } Ok(None) @@ -292,6 +348,16 @@ impl KvStageIntegration { runtime: &mut RuntimeState, session_id: &str, identity: &PrefillKvIdentity, + ) -> Result> { + self.record_exact_state_with_cost(runtime, session_id, identity, None) + } + + pub fn record_exact_state_with_cost( + &self, + runtime: &mut RuntimeState, + session_id: &str, + identity: &PrefillKvIdentity, + l3_cost: Option, ) -> Result> { let Some(exact_state_payload) = self.exact_state_payload() else { return Ok(None); @@ -320,7 +386,12 @@ impl KvStageIntegration { return Err(error); } }; - if already_recorded { + let probation_recurrence = already_recorded + && l3_cost.is_some() + && self.l3.as_ref().is_some_and(|l3| { + l3.benefit_tracks_prefix(&identity.namespace, &identity.token_ids) + }); + if already_recorded && !probation_recurrence { self.finish_record(&identity.page_id); return Ok(None); } @@ -397,6 +468,7 @@ impl KvStageIntegration { l3_fill_claim: None, write_through_l3: true, l2_promotion_digest: None, + l3_cost, }) { ExactStateRecordAdmission::Queued => { // Recording owns the radix/blob locks while it hashes a potentially @@ -447,6 +519,7 @@ impl KvStageIntegration { session_id: &str, identity: &PrefillKvIdentity, lookup_started: Instant, + cold_prefill_cost: Option, ) -> Result> { const MAX_PREFIX_PROBES: usize = 64; let Some(l3) = &self.l3 else { @@ -485,8 +558,15 @@ impl KvStageIntegration { return Ok(None); } } - let outcome = - self.fill_and_import(runtime, session_id, identity, lookup_started, l3, &location); + let outcome = self.fill_and_import( + runtime, + session_id, + identity, + lookup_started, + l3, + &location, + cold_prefill_cost, + ); // On success the claim travels with the re-warm record and the worker // releases it once the entry is radix-resident. On any other outcome // release it here. @@ -500,6 +580,7 @@ impl KvStageIntegration { outcome } + #[allow(clippy::too_many_arguments)] fn fill_and_import( &self, runtime: &mut RuntimeState, @@ -508,6 +589,7 @@ impl KvStageIntegration { lookup_started: Instant, l3: &std::sync::Arc, location: &skippy_cache::L3Location, + cold_prefill_cost: Option, ) -> Result> { // Native page capability is checked from manifest metadata before the // tier reads any segment bytes. Runtime ABI, platform and numerical @@ -619,6 +701,14 @@ impl KvStageIntegration { } let logical_bytes = fill.payload.byte_len(); let payload_kind = fill.payload.kind(); + let restore_cost = lookup_started.elapsed().as_secs_f64() * 1_000.0; + l3.benefit_observe_l3_restore( + &identity.namespace, + &identity.token_ids[..token_count as usize], + location, + cold_prefill_cost, + restore_cost, + ); // Re-warm the RAM tier off the request path. A drop is fine: the // disk copy stays authoritative. The fill claim rides along so the // worker releases it only once the entry is radix-resident. @@ -635,6 +725,7 @@ impl KvStageIntegration { l3_fill_claim: Some(l3_fill_claim_key(l3, location)), write_through_l3: false, l2_promotion_digest, + l3_cost: None, }); let rewarm_enqueued = matches!(admission, ExactStateRecordAdmission::Queued); Ok(Some(ExactStateRestore { diff --git a/crates/skippy-server/src/kv_integration/l2_serving.rs b/crates/skippy-server/src/kv_integration/l2_serving.rs index 1ae1e84796..c511e83688 100644 --- a/crates/skippy-server/src/kv_integration/l2_serving.rs +++ b/crates/skippy-server/src/kv_integration/l2_serving.rs @@ -328,6 +328,7 @@ impl KvStageIntegration { l3_fill_claim: None, write_through_l3: false, l2_promotion_digest: None, + l3_cost: None, }), ExactStateRecordAdmission::Queued ); diff --git a/crates/skippy-server/src/kv_integration/mod.rs b/crates/skippy-server/src/kv_integration/mod.rs index 26123ad372..bdbab91f54 100644 --- a/crates/skippy-server/src/kv_integration/mod.rs +++ b/crates/skippy-server/src/kv_integration/mod.rs @@ -239,6 +239,9 @@ pub(crate) struct PendingExactStateRecord { /// Durable manifest digest to mirror into L2 on this worker job. `None` /// leaves the payload out of L2. pub(crate) l2_promotion_digest: Option, + /// Measured cold-versus-restore cost for L3 benefit admission. Missing + /// telemetry preserves the established LRU write-through behavior. + pub(crate) l3_cost: Option, } #[derive(Debug)] @@ -1172,6 +1175,7 @@ mod exact_state_record_queue_tests { l3_fill_claim: None, write_through_l3: true, l2_promotion_digest: None, + l3_cost: None, } } diff --git a/docs/skippy/KV_CACHE_DISK.md b/docs/skippy/KV_CACHE_DISK.md index e22d2c24db..d85103e9b4 100644 --- a/docs/skippy/KV_CACHE_DISK.md +++ b/docs/skippy/KV_CACHE_DISK.md @@ -155,6 +155,8 @@ The status payload reports: `filesystem_available_bytes`, `minimum_free_bytes`, `manifests`, `unique_segments`, `evicted_manifests`, and `quarantined_objects`. - `activity`, `reconciliation` (see below), and `inventory` (per-model entries). + Activity includes benefit-admission probation, persistence, LRU-fallback, + and policy-selected eviction counters. At **startup**, any resolution warnings (deprecated legacy vars, zero auto budget, unavailable store) are emitted as `Warning` events in the node log. diff --git a/docs/skippy/PROMPT_CACHE.md b/docs/skippy/PROMPT_CACHE.md index b5356722c9..42246026c6 100644 --- a/docs/skippy/PROMPT_CACHE.md +++ b/docs/skippy/PROMPT_CACHE.md @@ -89,6 +89,20 @@ resident byte counts, logical bytes, entries, segments, hits, misses, inserts, evictions, and admission refusals. Exact-hit telemetry identifies the restore source as `l2` and includes fill time and whether the L1 rewarm was queued. +## Durable admission + +After a stage has measured an L3 restore, local OpenAI serving uses that +restore-cost EWMA with the generation service estimator's cold-prefill cost. +New entries remain in L1 probation and reach L3 after two observed reuses. +Admission scores reuse probability and saved prefill time per exclusive and +fractionally shared byte. Under disk pressure, the shared node manager removes +the lowest-benefit inactive manifests first and preserves active pins. Before +the first usable timing sample, or on serving paths without one, the existing +reference-aware LRU write-through remains the safe fallback. + +The L3 activity status reports `benefit_probation`, `benefit_persist`, +`benefit_lru_fallback`, and `benefit_evictions` counters. + ## mesh-llm Defaults mesh-llm wires Skippy prefix cache through family policy. For supported model From a15743fbfc4caf0649ad0a5d8f5f34ffa5e9a39a Mon Sep 17 00:00:00 2001 From: scama Date: Mon, 14 Sep 2026 09:05:04 +1000 Subject: [PATCH 07/16] feat(skippy): wire CacheGen into durable serving --- crates/mesh-llm-config/src/lib.rs | 10 + crates/mesh-llm-config/src/model.rs | 13 + .../control_behavior/runtime_controls.rs | 2 +- .../src/model/built_in_schema/declarations.rs | 5 + .../src/model/built_in_schema/presentation.rs | 12 + crates/mesh-llm-config/src/wiring_status.rs | 1 + .../src/wiring_status/runtime.rs | 1 + .../src/inference/skippy/family_policy.rs | 1 + .../inference/skippy/resolver/resolution.rs | 6 + .../inference/skippy/resolver/translation.rs | 3 + .../src/inference/skippy/resolver/types.rs | 5 +- .../src/runtime/config_state.rs | 2 +- .../src/runtime/config_state_tests/kv_disk.rs | 7 +- .../src/runtime/local_model_only.rs | 1 + crates/skippy-cache/src/identity.rs | 44 ++- crates/skippy-cache/src/l2/mod.rs | 3 + crates/skippy-cache/src/l3.rs | 136 ++++++- crates/skippy-cache/src/lib.rs | 15 +- crates/skippy-cache/src/tier.rs | 267 ++++++++++++-- .../src/runner/cachegen_gate.rs | 118 +----- .../src/prompt_cli/stage_config.rs | 2 + crates/skippy-protocol/src/config.rs | 14 + crates/skippy-protocol/src/lib.rs | 4 +- crates/skippy-runtime/src/kv_pages.rs | 104 +++++- crates/skippy-runtime/src/lib.rs | 1 + .../src/binary_transport/stage_execution.rs | 1 + .../src/frontend/local_generation/tests.rs | 1 + .../src/frontend/prefix_cache.rs | 2 + .../src/frontend/tests/support.rs | 1 + .../src/kv_integration/activation.rs | 1 + .../src/kv_integration/cache_affinity.rs | 1 + .../src/kv_integration/config.rs | 336 +++++++++++++++++- .../src/kv_integration/exact_state.rs | 77 +++- .../src/kv_integration/identity.rs | 1 + .../src/kv_integration/l2_serving.rs | 2 + .../skippy-server/src/kv_integration/mod.rs | 7 + .../src/kv_integration/resident_prefix.rs | 1 + .../src/runtime_state/lane_lifecycle.rs | 19 + docs/skippy/CACHEGEN_BACKEND_PLAN.md | 20 +- docs/skippy/CONFIGURATION.md | 1 + docs/skippy/KV_CACHE_DISK.md | 10 +- tools/xtask/data/console_print_allowlist.json | 8 +- website/src/docs/pages/config-reference.md | 1 + website/src/docs/pages/kv-caching.md | 7 + 44 files changed, 1072 insertions(+), 202 deletions(-) diff --git a/crates/mesh-llm-config/src/lib.rs b/crates/mesh-llm-config/src/lib.rs index a8d9d7e870..487f564be8 100644 --- a/crates/mesh-llm-config/src/lib.rs +++ b/crates/mesh-llm-config/src/lib.rs @@ -308,6 +308,10 @@ skippy_abi = "0.1.25" 16_384 ); assert_eq!(defaults.runtime.kv_cache.disk.budget_mib, None); + assert_eq!( + defaults.runtime.kv_cache.disk.codec, + crate::KvDiskCodec::Native + ); let fixed = parse_config_toml( r#" @@ -316,6 +320,7 @@ mode = "fixed" directory = "/fast-disk/mesh-kv-cache" budget_mib = 32768 minimum_free_mib = 16384 +codec = "cachegen" "#, ) .expect("fixed disk-cache config should parse"); @@ -324,6 +329,10 @@ minimum_free_mib = 16384 Some(KvDiskTierMode::Fixed) ); assert_eq!(fixed.runtime.kv_cache.disk.budget_mib, Some(32_768)); + assert_eq!( + fixed.runtime.kv_cache.disk.codec, + crate::KvDiskCodec::CacheGen + ); } #[test] @@ -420,6 +429,7 @@ selection = "vulcan" for path in [ "runtime.kv_cache.disk.mode", "runtime.kv_cache.disk.directory", + "runtime.kv_cache.disk.codec", ] { assert_eq!(setting(path).apply_mode, ConfigApplyMode::StaticOnLoad); assert_eq!( diff --git a/crates/mesh-llm-config/src/model.rs b/crates/mesh-llm-config/src/model.rs index 9665c2f1ec..ce11841429 100644 --- a/crates/mesh-llm-config/src/model.rs +++ b/crates/mesh-llm-config/src/model.rs @@ -195,6 +195,15 @@ pub enum KvDiskTierMode { Fixed, } +#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum KvDiskCodec { + #[default] + Native, + #[serde(rename = "cachegen")] + CacheGen, +} + #[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq, Eq)] pub struct RuntimeKvCacheConfig { #[serde(default)] @@ -213,6 +222,10 @@ pub struct KvDiskTierConfig { pub budget_mib: Option, #[serde(default)] pub minimum_free_mib: Option, + /// KV representation persisted in the disk tier. CacheGen remains an + /// explicit opt-in while backend/dtype qualification is incomplete. + #[serde(default)] + pub codec: KvDiskCodec, } impl KvDiskTierConfig { diff --git a/crates/mesh-llm-config/src/model/built_in_schema/control_behavior/runtime_controls.rs b/crates/mesh-llm-config/src/model/built_in_schema/control_behavior/runtime_controls.rs index 4fa148a886..e148928c8f 100644 --- a/crates/mesh-llm-config/src/model/built_in_schema/control_behavior/runtime_controls.rs +++ b/crates/mesh-llm-config/src/model/built_in_schema/control_behavior/runtime_controls.rs @@ -110,7 +110,7 @@ pub(super) fn apply_runtime_controls_behavior(setting: &mut ConfigSettingSchema, "runtime.activity.response" | "runtime.activity.advertisement" => { set_static_options(setting) } - "runtime.kv_cache.disk.mode" => set_static_options(setting), + "runtime.kv_cache.disk.mode" | "runtime.kv_cache.disk.codec" => set_static_options(setting), "runtime.kv_cache.disk.directory" => { set_text_format(setting, ConfigTextFormat::Path); push_non_empty_constraint(setting); diff --git a/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs b/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs index 625c3edcce..39fe9e5ff5 100644 --- a/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs +++ b/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs @@ -133,6 +133,11 @@ fn build_built_in_config_schema() -> ConfigSchema { ConfigValueSchema::Integer, true, ), + kv_disk_setting( + "runtime.kv_cache.disk.codec", + string_enum(["native", "cachegen"]), + false, + ), runtime_setting( "runtime.model_target_demand_upgrade_min_requests", ConfigValueSchema::Integer, diff --git a/crates/mesh-llm-config/src/model/built_in_schema/presentation.rs b/crates/mesh-llm-config/src/model/built_in_schema/presentation.rs index 0023497e5c..ebc926dd2e 100644 --- a/crates/mesh-llm-config/src/model/built_in_schema/presentation.rs +++ b/crates/mesh-llm-config/src/model/built_in_schema/presentation.rs @@ -167,6 +167,18 @@ fn kv_disk_presentation(rendered: &str) -> Option { .unit("MiB") .hint("number"), ), + "runtime.kv_cache.disk.codec" => Some( + sp( + "Disk cache codec", + "Persist native KV pages or qualified CacheGen archives. CacheGen currently activates only for validated Metal cache types.", + PROMPT_CACHE_CATEGORY, + 50, + ) + .choices(&[ + ("native", "Native", "Persist exact native KV pages."), + ("cachegen", "CacheGen", "Use CacheGen for qualified Metal KV layouts."), + ]), + ), _ => None, } } diff --git a/crates/mesh-llm-config/src/wiring_status.rs b/crates/mesh-llm-config/src/wiring_status.rs index f757699335..1bc5fde721 100644 --- a/crates/mesh-llm-config/src/wiring_status.rs +++ b/crates/mesh-llm-config/src/wiring_status.rs @@ -515,6 +515,7 @@ pub const WIRING_MANIFEST: &[WiringEntry] = &[ runtime::KV_CACHE_DISK_DIRECTORY, runtime::KV_CACHE_DISK_BUDGET_MIB, runtime::KV_CACHE_DISK_MINIMUM_FREE_MIB, + runtime::KV_CACHE_DISK_CODEC, WiringEntry { path: "runtime.model_target_demand_upgrade_min_requests", status: WiringStatus::Wired, diff --git a/crates/mesh-llm-config/src/wiring_status/runtime.rs b/crates/mesh-llm-config/src/wiring_status/runtime.rs index 990eaee6ea..1731343e71 100644 --- a/crates/mesh-llm-config/src/wiring_status/runtime.rs +++ b/crates/mesh-llm-config/src/wiring_status/runtime.rs @@ -7,6 +7,7 @@ pub(super) const KV_CACHE_DISK_BUDGET_MIB: WiringEntry = wired_kv_cache("runtime.kv_cache.disk.budget_mib"); pub(super) const KV_CACHE_DISK_MINIMUM_FREE_MIB: WiringEntry = wired_kv_cache("runtime.kv_cache.disk.minimum_free_mib"); +pub(super) const KV_CACHE_DISK_CODEC: WiringEntry = wired_kv_cache("runtime.kv_cache.disk.codec"); const fn wired_kv_cache(path: &'static str) -> WiringEntry { WiringEntry { diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs index ca91b85e3e..e30b8b7807 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs @@ -55,6 +55,7 @@ impl FamilyPolicy { max_entries: bounded_entries, max_bytes, l2_max_bytes: 0, + codec: skippy_protocol::StageKvCacheCodec::Native, min_tokens, shared_prefix_stride_tokens: 128, shared_prefix_record_limit: derive_shared_prefix_record_limit(bounded_entries), diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs index e737503d8f..519a82f525 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs @@ -22,6 +22,7 @@ use super::types::{ use crate::plugin::{ BoolOrAuto, ModelConfigDefaults, ModelConfigEntry, ModelFitConfig, ThroughputConfig, }; +use mesh_llm_config::KvDiskCodec; #[cfg(test)] pub(crate) fn resolve_skippy_config( @@ -276,6 +277,10 @@ fn resolve_model_fit_config( if l2_max_bytes > 0 && matches!(prefix_cache, ResolvedStageKvCache::Disabled) { anyhow::bail!("model_fit.cache_ram_mib requires prefix caching to be enabled"); } + let kv_cache_codec = match context.request.mesh_config.runtime.kv_cache.disk.codec { + KvDiskCodec::Native => skippy_protocol::StageKvCacheCodec::Native, + KvDiskCodec::CacheGen => skippy_protocol::StageKvCacheCodec::CacheGen, + }; Ok(ResolvedModelFitConfig { ctx_size, @@ -286,6 +291,7 @@ fn resolve_model_fit_config( kv_cache_policy: kv.effective_policy, prefix_cache, l2_max_bytes, + kv_cache_codec, kv_offload, kv_offload_resolved, kv_unified, diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs index 8a79e5f3cd..5edc1388cc 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs @@ -480,6 +480,7 @@ impl ResolvedSkippyConfig { max_entries: 0, max_bytes: 0, l2_max_bytes: 0, + codec: skippy_protocol::StageKvCacheCodec::Native, min_tokens: 0, shared_prefix_stride_tokens: 0, shared_prefix_record_limit: 0, @@ -491,6 +492,7 @@ impl ResolvedSkippyConfig { max_entries: 128, max_bytes: 0, l2_max_bytes: 0, + codec: skippy_protocol::StageKvCacheCodec::Native, min_tokens: 256, shared_prefix_stride_tokens: 128, shared_prefix_record_limit: 2, @@ -522,6 +524,7 @@ impl ResolvedSkippyConfig { } if let Some(cache) = resolved.as_mut() { cache.l2_max_bytes = self.model_fit.l2_max_bytes; + cache.codec = self.model_fit.kv_cache_codec; } Ok(resolved) } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs index 4e2217184b..4be29cd9df 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs @@ -1,6 +1,8 @@ use std::path::{Path, PathBuf}; -use skippy_protocol::{FlashAttentionType, StageKvCacheMode, StageKvCachePayload}; +use skippy_protocol::{ + FlashAttentionType, StageKvCacheCodec, StageKvCacheMode, StageKvCachePayload, +}; use skippy_runtime::package::PackageGenerationInfo; use skippy_server::{EmbeddedOpenAiRequestDefaults, SpeculativeDecodeConfig}; @@ -75,6 +77,7 @@ pub(crate) struct ResolvedModelFitConfig { pub(crate) kv_cache_policy: String, pub(crate) prefix_cache: ResolvedStageKvCache, pub(crate) l2_max_bytes: u64, + pub(crate) kv_cache_codec: StageKvCacheCodec, pub(crate) kv_offload: String, /// Parsed `kv_offload` for the native tri-state control. `None` covers /// both "auto" and any value that did not parse to a bool. diff --git a/crates/mesh-llm-host-runtime/src/runtime/config_state.rs b/crates/mesh-llm-host-runtime/src/runtime/config_state.rs index 651c4e06c4..83f8f6ac32 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/config_state.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/config_state.rs @@ -325,7 +325,7 @@ fn logging_dynamic_limits_changed(old: &LoggingConfig, new: &LoggingConfig) -> b } fn kv_disk_changes_require_restart(old: &KvDiskTierConfig, new: &KvDiskTierConfig) -> bool { - old.mode != new.mode || old.directory != new.directory + old.mode != new.mode || old.directory != new.directory || old.codec != new.codec } fn kv_disk_dynamic_limits_changed(old: &KvDiskTierConfig, new: &KvDiskTierConfig) -> bool { diff --git a/crates/mesh-llm-host-runtime/src/runtime/config_state_tests/kv_disk.rs b/crates/mesh-llm-host-runtime/src/runtime/config_state_tests/kv_disk.rs index fbd43a50af..1aad724804 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/config_state_tests/kv_disk.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/config_state_tests/kv_disk.rs @@ -1,5 +1,5 @@ use super::*; -use mesh_llm_config::{KvDiskTierConfig, KvDiskTierMode}; +use mesh_llm_config::{KvDiskCodec, KvDiskTierConfig, KvDiskTierMode}; use std::path::PathBuf; fn disk_config() -> KvDiskTierConfig { @@ -8,6 +8,7 @@ fn disk_config() -> KvDiskTierConfig { directory: Some(PathBuf::from("/var/lib/mesh-llm/kv-cache")), budget_mib: Some(32 * 1024), minimum_free_mib: Some(16 * 1024), + codec: Default::default(), } } @@ -22,6 +23,10 @@ fn disk_mode_and_directory_changes_require_restart() { let mut directory = old.clone(); directory.directory = Some(PathBuf::from("/var/lib/mesh-llm/other-cache")); assert!(kv_disk_changes_require_restart(&old, &directory)); + + let mut codec = old.clone(); + codec.codec = KvDiskCodec::CacheGen; + assert!(kv_disk_changes_require_restart(&old, &codec)); } #[test] diff --git a/crates/mesh-llm-host-runtime/src/runtime/local_model_only.rs b/crates/mesh-llm-host-runtime/src/runtime/local_model_only.rs index c6408f7d11..e3c398fa25 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local_model_only.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local_model_only.rs @@ -506,6 +506,7 @@ mod tests { max_entries: 8, max_bytes: 0, l2_max_bytes: 0, + codec: skippy_protocol::StageKvCacheCodec::Native, min_tokens: 8, shared_prefix_stride_tokens: 8, shared_prefix_record_limit: 2, diff --git a/crates/skippy-cache/src/identity.rs b/crates/skippy-cache/src/identity.rs index ad2506e248..d84d9c20ef 100644 --- a/crates/skippy-cache/src/identity.rs +++ b/crates/skippy-cache/src/identity.rs @@ -1,4 +1,4 @@ -use skippy_protocol::{FlashAttentionType, LoadMode, StageConfig}; +use skippy_protocol::{FlashAttentionType, LoadMode, StageConfig, StageKvCacheCodec}; pub const NATIVE_KV_RUNTIME_ABI_VERSION: &str = "stage-abi-0.1.52/native-kv-page-v3"; pub const NATIVE_KV_DTYPE: &str = "ggml-native-kv"; @@ -263,6 +263,16 @@ pub fn exact_state_identity_for_stage(config: &StageConfig, payload_kind: &str) hasher.update(&config.lane_count.to_le_bytes()); hasher.update(b"payload:"); hasher.update(payload_kind.as_bytes()); + // Native remains byte-for-byte compatible with the pre-selector identity. + // Opt-in lossy storage gets a distinct namespace so disabling CacheGen can + // never restore an archive written by an earlier process. + if config + .kv_cache + .as_ref() + .is_some_and(|cache| cache.codec == StageKvCacheCodec::CacheGen) + { + hasher.update(b"lossy-disk-codec:cachegen-kv-envelope-v1/cachegen-v1"); + } format!("blake3:{}", hasher.finalize().to_hex()) } @@ -969,6 +979,38 @@ mod identity_stability_tests { ); } + #[test] + fn cachegen_uses_a_distinct_exact_state_identity() { + let native = config_with_topology("topology-a"); + let cachegen = StageConfig { + kv_cache: Some(skippy_protocol::StageKvCacheConfig { + mode: skippy_protocol::StageKvCacheMode::LookupRecord, + payload: skippy_protocol::StageKvCachePayload::KvRecurrent, + max_entries: 64, + max_bytes: 0, + l2_max_bytes: 0, + codec: StageKvCacheCodec::CacheGen, + min_tokens: 64, + shared_prefix_stride_tokens: 128, + shared_prefix_record_limit: 2, + }), + ..native.clone() + }; + let mut explicit_native = cachegen.clone(); + explicit_native.kv_cache.as_mut().unwrap().codec = StageKvCacheCodec::Native; + + assert_ne!( + exact_state_identity_for_stage(&native, "kv-recurrent"), + exact_state_identity_for_stage(&cachegen, "kv-recurrent"), + "opting out of CacheGen must make earlier lossy entries unreachable" + ); + assert_eq!( + exact_state_identity_for_stage(&native, "kv-recurrent"), + exact_state_identity_for_stage(&explicit_native, "kv-recurrent"), + "the default native codec must preserve existing durable identities" + ); + } + /// Weight identity is the protection that replaced `topology_id`. /// /// Two runs can share a `model_id` while serving different tensors — a diff --git a/crates/skippy-cache/src/l2/mod.rs b/crates/skippy-cache/src/l2/mod.rs index 5d72cc8b88..0a73565c9c 100644 --- a/crates/skippy-cache/src/l2/mod.rs +++ b/crates/skippy-cache/src/l2/mod.rs @@ -1959,6 +1959,7 @@ mod tests { }, ], kv_bytes: 24, + kv_decoded_bytes: 24, recurrent_bytes: 8, kv_desc_json: None, token_count: 4, @@ -2001,6 +2002,7 @@ mod tests { payload_digest: "blake3:aa".to_string(), segments: Vec::new(), kv_bytes: 10, + kv_decoded_bytes: 10, recurrent_bytes: 0, kv_desc_json: None, token_count: 1, @@ -2732,6 +2734,7 @@ mod tests { meta_json: None, }], kv_bytes: 32, + kv_decoded_bytes: 32, recurrent_bytes: 0, kv_desc_json: None, token_count: 1, diff --git a/crates/skippy-cache/src/l3.rs b/crates/skippy-cache/src/l3.rs index e8a5498ac1..5b2f2bad9e 100644 --- a/crates/skippy-cache/src/l3.rs +++ b/crates/skippy-cache/src/l3.rs @@ -85,6 +85,9 @@ pub const CODEC_NATIVE_KV_PAGE: &str = "native-kv-page"; /// tensor types, geometry, and platform are bound by the exact-state identity /// and the serialized runtime page descriptor carried by the manifest. pub const CODEC_NATIVE_KV_PAGE_VERSION: u32 = 1; +/// CacheGen archive followed by an exact recurrent-state tail. +pub const CODEC_CACHEGEN_KV_ENVELOPE: &str = "cachegen-kv-envelope"; +pub const CODEC_CACHEGEN_KV_ENVELOPE_VERSION: u32 = 1; /// The codec used to encode a payload's segment bytes, stamped into the /// manifest so representations are explicit and negotiable. @@ -114,11 +117,20 @@ impl PayloadCodec { } } + pub fn cachegen_kv_envelope() -> Self { + Self { + name: CODEC_CACHEGEN_KV_ENVELOPE.to_string(), + version: CODEC_CACHEGEN_KV_ENVELOPE_VERSION, + } + } + /// Whether this build can assemble a payload encoded with this codec. /// Only the exact raw name and version are supported; any other name or a /// future raw version is unknown and must be refused before assembly. pub fn is_supported(&self) -> bool { - self.name == CODEC_RAW && self.version == CODEC_RAW_VERSION + (self.name == CODEC_RAW && self.version == CODEC_RAW_VERSION) + || (self.name == CODEC_CACHEGEN_KV_ENVELOPE + && self.version == CODEC_CACHEGEN_KV_ENVELOPE_VERSION) } } @@ -208,6 +220,10 @@ impl SegmentCodecIdentity { } } + pub fn cachegen_archive(decoded_len: u64, calibration_digest: String) -> Self { + crate::cachegen::lmcache::segment_identity(decoded_len, calibration_digest) + } + /// Whether this identity names the supported native KV passthrough /// representation. This is stricter than a name check: an identity with /// the wrong version, class, or calibration is not native passthrough. @@ -218,16 +234,27 @@ impl SegmentCodecIdentity { && self.calibration_digest.is_none() } + pub fn is_cachegen_archive(&self) -> bool { + self.name == crate::cachegen::lmcache::CACHEGEN_CODEC_NAME + && self.version == crate::cachegen::lmcache::CACHEGEN_CODEC_VERSION + && self.class == CodecClass::Lossy + && self + .calibration_digest + .as_deref() + .is_some_and(|digest| !digest.is_empty()) + } + /// Whether this build can assemble a segment encoded with this identity. /// Exact raw and native KV passthrough are implemented. Both assemble /// verbatim; anything else is unknown and must be refused before assembly. pub fn is_supported(&self) -> bool { let supported_representation = (self.name == CODEC_RAW && self.version == CODEC_RAW_VERSION) - || self.is_native_kv_page(); + || self.is_native_kv_page() + || self.is_cachegen_archive(); supported_representation - && self.class == CodecClass::Exact - && self.calibration_digest.is_none() + && (self.is_cachegen_archive() + || (self.class == CodecClass::Exact && self.calibration_digest.is_none())) } /// Internal consistency the capability negotiation relies on: an exact @@ -288,6 +315,10 @@ pub struct HandoffManifest { pub payload_digest: String, pub segments: Vec, pub kv_bytes: u64, + /// Decoded native KV length. Equal to `kv_bytes` for native pages; for a + /// CacheGen envelope `kv_bytes` is the archive boundary. + #[serde(default)] + pub kv_decoded_bytes: u64, pub recurrent_bytes: u64, /// Serialized `RuntimeKvPageDesc` for kv-recurrent payloads; opaque to /// this crate so the store does not depend on the runtime. @@ -321,6 +352,7 @@ impl HandoffManifest { payload_digest: String::new(), segments: Vec::new(), kv_bytes: 0, + kv_decoded_bytes: 0, recurrent_bytes: 0, kv_desc_json: None, token_count: 0, @@ -338,6 +370,15 @@ impl HandoffManifest { .is_some_and(SegmentCodecIdentity::is_native_kv_page) }) } + + pub fn uses_cachegen_kv(&self) -> bool { + self.segments.iter().any(|segment| { + segment + .codec_identity + .as_ref() + .is_some_and(SegmentCodecIdentity::is_cachegen_archive) + }) + } } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -1914,12 +1955,10 @@ fn reject_unsupported_codec(manifest: &HandoffManifest) -> Result<()> { match manifest.codec.as_ref() { Some(codec) if codec.is_supported() => Ok(()), Some(codec) => bail!( - "manifest {} uses unsupported codec {}/{}; this build assembles only {}/{}", + "manifest {} uses unsupported codec {}/{}", manifest.payload_digest, codec.name, - codec.version, - CODEC_RAW, - CODEC_RAW_VERSION + codec.version ), None => bail!( "manifest {} has no codec identity; this build requires an explicit codec", @@ -1944,6 +1983,26 @@ fn manifest_version_is_supported(version: u32) -> bool { /// segment index so the offending ref is identifiable in the message. fn reject_unsupported_segment_codecs(manifest: &HandoffManifest) -> Result<()> { let has_native_kv = manifest.uses_native_kv_passthrough(); + let has_cachegen_kv = manifest.uses_cachegen_kv(); + let kv_representation = if has_cachegen_kv { + "CacheGen KV" + } else { + "native KV" + }; + if has_native_kv && has_cachegen_kv { + bail!( + "manifest {} mixes native and CacheGen KV representations", + manifest.payload_digest + ); + } + let envelope_is_cachegen = + manifest.codec.as_ref() == Some(&PayloadCodec::cachegen_kv_envelope()); + if envelope_is_cachegen != has_cachegen_kv { + bail!( + "manifest {} payload codec disagrees with its CacheGen segments", + manifest.payload_digest + ); + } let has_runtime_descriptor = manifest .kv_desc_json .as_deref() @@ -1961,6 +2020,46 @@ fn reject_unsupported_segment_codecs(manifest: &HandoffManifest) -> Result<()> { manifest.payload_digest ); } + if has_cachegen_kv + && (manifest.version != MANIFEST_VERSION + || manifest.payload_kind != "kv-recurrent" + || manifest.kv_bytes == 0 + || manifest.kv_decoded_bytes == 0 + || !has_runtime_descriptor + || manifest.kv_bytes.checked_add(manifest.recurrent_bytes) + != Some(manifest.total_bytes) + || manifest.codec.as_ref() != Some(&PayloadCodec::cachegen_kv_envelope())) + { + bail!( + "manifest {} names CacheGen KV without a current, complete kv-recurrent envelope and runtime page descriptor", + manifest.payload_digest + ); + } + if has_cachegen_kv { + let mut cachegen_segments = manifest.segments.iter().filter(|segment| { + segment + .codec_identity + .as_ref() + .is_some_and(SegmentCodecIdentity::is_cachegen_archive) + }); + let cachegen_segment = cachegen_segments + .next() + .context("CacheGen manifest has no archive segment")?; + let cachegen_identity = cachegen_segment + .codec_identity + .as_ref() + .context("CacheGen archive segment has no codec identity")?; + if cachegen_segments.next().is_some() + || cachegen_segment.offset != 0 + || cachegen_segment.bytes != manifest.kv_bytes + || cachegen_identity.decoded_len != manifest.kv_decoded_bytes + { + bail!( + "manifest {} CacheGen archive must be one complete KV segment with matching encoded and decoded lengths", + manifest.payload_digest + ); + } + } for segment in &manifest.segments { let Some(identity) = segment.codec_identity.as_ref() else { // Legacy formats carry no per-segment identity by construction; @@ -1969,15 +2068,11 @@ fn reject_unsupported_segment_codecs(manifest: &HandoffManifest) -> Result<()> { }; if !identity.is_supported() { bail!( - "manifest {} segment {} uses unsupported codec {}/{}; this build assembles only {}/{} and {}/{}", + "manifest {} segment {} uses unsupported codec {}/{}", manifest.payload_digest, segment.index, identity.name, - identity.version, - CODEC_RAW, - CODEC_RAW_VERSION, - CODEC_NATIVE_KV_PAGE, - CODEC_NATIVE_KV_PAGE_VERSION + identity.version ); } if !identity.is_self_consistent(segment.bytes) { @@ -1991,7 +2086,7 @@ fn reject_unsupported_segment_codecs(manifest: &HandoffManifest) -> Result<()> { segment.bytes ); } - if has_native_kv { + if has_native_kv || has_cachegen_kv { let end = segment .offset .checked_add(segment.bytes) @@ -1999,15 +2094,20 @@ fn reject_unsupported_segment_codecs(manifest: &HandoffManifest) -> Result<()> { let covers_kv = segment.offset < manifest.kv_bytes; if covers_kv && end > manifest.kv_bytes { bail!( - "manifest {} segment {} crosses the native KV boundary at byte {}", + "manifest {} segment {} crosses the {kv_representation} boundary at byte {}", manifest.payload_digest, segment.index, manifest.kv_bytes ); } - if covers_kv != identity.is_native_kv_page() { + let expected_kv = if has_cachegen_kv { + identity.is_cachegen_archive() + } else { + identity.is_native_kv_page() + }; + if covers_kv != expected_kv { bail!( - "manifest {} segment {} representation disagrees with the native KV boundary at byte {}", + "manifest {} segment {} representation disagrees with the {kv_representation} boundary at byte {}", manifest.payload_digest, segment.index, manifest.kv_bytes diff --git a/crates/skippy-cache/src/lib.rs b/crates/skippy-cache/src/lib.rs index 6da4340630..c9ba383896 100644 --- a/crates/skippy-cache/src/lib.rs +++ b/crates/skippy-cache/src/lib.rs @@ -25,11 +25,12 @@ pub use l2::{ l2_cache_key, }; pub use l3::{ - CODEC_NATIVE_KV_PAGE, CODEC_NATIVE_KV_PAGE_VERSION, CODEC_RAW, CODEC_RAW_VERSION, CodecClass, - GeometryBlock, GeometryKind, HandoffManifest, HandoffSegmentRef, HandoffSegmentStore, - LEGACY_MANIFEST_VERSION, LEGACY_PAYLOAD_CODEC_MANIFEST_VERSION, MANIFEST_VERSION, ManifestPin, - PayloadCodec, PayloadGeometry, Reservation, SegmentCodecIdentity, SegmentHold, SegmentPut, - StoreLimits, StoreReconciliation, StoreUsage, StoredSegment, WriteRefusal, segment_digest, + CODEC_CACHEGEN_KV_ENVELOPE, CODEC_CACHEGEN_KV_ENVELOPE_VERSION, CODEC_NATIVE_KV_PAGE, + CODEC_NATIVE_KV_PAGE_VERSION, CODEC_RAW, CODEC_RAW_VERSION, CodecClass, GeometryBlock, + GeometryKind, HandoffManifest, HandoffSegmentRef, HandoffSegmentStore, LEGACY_MANIFEST_VERSION, + LEGACY_PAYLOAD_CODEC_MANIFEST_VERSION, MANIFEST_VERSION, ManifestPin, PayloadCodec, + PayloadGeometry, Reservation, SegmentCodecIdentity, SegmentHold, SegmentPut, StoreLimits, + StoreReconciliation, StoreUsage, StoredSegment, WriteRefusal, segment_digest, }; pub use l3_remote::{ FetchStats, KvFetchClient, serve_connection, serve_store, serve_store_with_timeout, @@ -51,7 +52,9 @@ pub use resident::{ }; pub use source::{ManifestSource, SegmentSource}; -pub use tier::{L3Fill, L3Location, L3Status, L3Tier, l3_namespace_key, l3_prefix_key}; +pub use tier::{ + CacheGenKvPayload, L3Fill, L3Location, L3Status, L3Tier, l3_namespace_key, l3_prefix_key, +}; /// llama.cpp's hard sequence-id capacity for one context. pub const LLAMA_MAX_SEQ: i32 = 256; diff --git a/crates/skippy-cache/src/tier.rs b/crates/skippy-cache/src/tier.rs index 963b82f11d..dce3154955 100644 --- a/crates/skippy-cache/src/tier.rs +++ b/crates/skippy-cache/src/tier.rs @@ -14,8 +14,8 @@ use anyhow::{Context, Result, bail}; use serde::Serialize; use crate::l3::{ - HandoffManifest, HandoffSegmentRef, HandoffSegmentStore, MANIFEST_VERSION, PayloadGeometry, - SegmentCodecIdentity, StoreLimits, StoreUsage, segment_digest, + HandoffManifest, HandoffSegmentRef, HandoffSegmentStore, MANIFEST_VERSION, PayloadCodec, + PayloadGeometry, SegmentCodecIdentity, StoreLimits, StoreUsage, segment_digest, }; use crate::manager::{ BenefitRestoreObservation, L3ActivitySnapshot, L3CacheManager, L3EffectiveStatus, @@ -90,6 +90,8 @@ pub struct L3Location { pub kv_desc_json: Option, pub kv_bytes: u64, pub native_kv_passthrough: bool, + pub cachegen_kv: bool, + pub kv_decoded_bytes: u64, } /// A successful fill from the tier. @@ -101,6 +103,24 @@ pub struct L3Fill { pub token_count: u64, pub kv_desc_json: Option, pub payload_bytes: u64, + pub cachegen_kv: bool, + pub kv_decoded_bytes: u64, +} + +/// Encoded KV component supplied by the serving worker. The calibration +/// digest binds the archive's lossy numerical representation in the manifest. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CacheGenKvPayload { + pub archive: Vec, + pub decoded_len: u64, + pub calibration_digest: String, +} + +struct SpillOptions<'a> { + kv_desc_json: Option, + geometry: Option<&'a PayloadGeometry>, + cost: Option, + cachegen: Option, } pub struct L3Tier { @@ -110,17 +130,25 @@ pub struct L3Tier { segment_bytes: usize, } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq)] enum SegmentRepresentation { Raw, NativeKvPage, + CacheGenKv { + decoded_len: u64, + calibration_digest: String, + }, } impl SegmentRepresentation { - fn identity(self, encoded_len: u64) -> SegmentCodecIdentity { + fn identity(&self, encoded_len: u64) -> SegmentCodecIdentity { match self { Self::Raw => SegmentCodecIdentity::raw(encoded_len), Self::NativeKvPage => SegmentCodecIdentity::native_kv_page(encoded_len), + Self::CacheGenKv { + decoded_len, + calibration_digest, + } => SegmentCodecIdentity::cachegen_archive(*decoded_len, calibration_digest.clone()), } } } @@ -147,7 +175,7 @@ fn append_fixed_cuts( offset: *offset, len, label: String::new(), - representation, + representation: representation.clone(), }); *offset = offset.saturating_add(len); remaining -= len; @@ -338,29 +366,72 @@ impl L3Tier { cost: Option, ) -> Result> { let _operation = self.manager.operation_guard(); - let result = self.spill_inner(namespace, token_ids, payload, kv_desc_json, geometry, cost); + let result = self.spill_inner( + namespace, + token_ids, + payload, + SpillOptions { + kv_desc_json, + geometry, + cost, + cachegen: None, + }, + ); if let Err(error) = &result { self.manager.activity_counters().record_error(error); } result } - fn spill_inner( + /// Persist a CacheGen archive plus the payload's exact recurrent tail. + /// The archive is already encoded on the background serving worker. + pub fn spill_cachegen_with_cost( &self, namespace: &str, token_ids: &[i32], payload: &ExactStatePayload, - kv_desc_json: Option, - geometry: Option<&PayloadGeometry>, + kv_desc_json: String, + cachegen: CacheGenKvPayload, cost: Option, ) -> Result> { + let _operation = self.manager.operation_guard(); + let result = self.spill_inner( + namespace, + token_ids, + payload, + SpillOptions { + kv_desc_json: Some(kv_desc_json), + geometry: None, + cost, + cachegen: Some(cachegen), + }, + ); + if let Err(error) = &result { + self.manager.activity_counters().record_error(error); + } + result + } + + fn spill_inner( + &self, + namespace: &str, + token_ids: &[i32], + payload: &ExactStatePayload, + options: SpillOptions<'_>, + ) -> Result> { + let SpillOptions { + kv_desc_json, + geometry, + cost, + mut cachegen, + } = options; if payload.byte_len() == 0 { bail!( "refusing to spill an empty exact-state payload: no state component was exported" ); } let token_count = token_ids.len() as u64; - let (kv, recurrent): (Vec, Vec) = match payload.kind() { + let (mut kv, recurrent): (Vec, Vec) = match payload.kind() { ExactStatePayloadKind::FullState => ( payload .full_state_bytes_timed() @@ -395,6 +466,25 @@ impl L3Tier { ), }; + let kv_decoded_bytes = kv.len() as u64; + if let Some(cachegen) = cachegen.as_mut() { + if payload.kind() != ExactStatePayloadKind::KvRecurrent + || cachegen.archive.is_empty() + || cachegen.decoded_len != kv_decoded_bytes + || cachegen.calibration_digest.is_empty() + || kv_desc_json.as_deref().is_none_or(str::is_empty) + { + bail!( + "invalid CacheGen spill: archive, descriptor, and decoded KV length are required" + ); + } + let decoded_len = usize::try_from(cachegen.decoded_len) + .context("CacheGen decoded KV length exceeds usize")?; + crate::cachegen::archive::validate_archive(&cachegen.archive, decoded_len) + .context("invalid CacheGen archive supplied for L3 spill")?; + kv = std::mem::take(&mut cachegen.archive); + } + // For full-state and recurrent-only exactly one component is populated, // and KV states reach gigabytes: concatenating would peak at twice the // payload for no benefit. Only a genuine composite needs the copy. @@ -419,8 +509,14 @@ impl L3Tier { manifest.total_bytes = wire.len() as u64; manifest.payload_digest = payload_digest.clone(); manifest.kv_bytes = kv_bytes; + manifest.kv_decoded_bytes = kv_decoded_bytes; manifest.recurrent_bytes = recurrent_bytes; - let native_kv_passthrough = payload.kind() == ExactStatePayloadKind::KvRecurrent + let cachegen_kv = cachegen.is_some(); + if cachegen_kv { + manifest.codec = Some(PayloadCodec::cachegen_kv_envelope()); + } + let native_kv_passthrough = !cachegen_kv + && payload.kind() == ExactStatePayloadKind::KvRecurrent && kv_bytes > 0 && kv_desc_json .as_deref() @@ -432,18 +528,21 @@ impl L3Tier { // geometry that does not describe these exact bytes is ignored rather // than trusted: mis-cutting would still reassemble, but silently write // the whole payload again every turn. - let geometry = geometry.filter(|geometry| { - let geometry_kv_bytes = geometry.total_bytes().saturating_sub(geometry.tail_bytes); - let matches = geometry.matches(wire.len() as u64) - && (!native_kv_passthrough || geometry_kv_bytes == kv_bytes); - if !matches { - self.manager - .activity_counters() - .geometry_rejected - .fetch_add(1, Ordering::Relaxed); - } - matches - }); + let geometry = (!cachegen_kv) + .then_some(geometry) + .flatten() + .filter(|geometry| { + let geometry_kv_bytes = geometry.total_bytes().saturating_sub(geometry.tail_bytes); + let matches = geometry.matches(wire.len() as u64) + && (!native_kv_passthrough || geometry_kv_bytes == kv_bytes); + if !matches { + self.manager + .activity_counters() + .geometry_rejected + .fetch_add(1, Ordering::Relaxed); + } + matches + }); let cuts = match geometry { Some(geometry) => geometry .plan(self.segment_bytes as u64) @@ -462,7 +561,25 @@ impl L3Tier { None => { let mut cuts = Vec::new(); let mut offset = 0u64; - if native_kv_passthrough { + if let Some(cachegen) = cachegen.as_ref() { + cuts.push(SegmentCut { + offset: 0, + len: kv_bytes, + representation: SegmentRepresentation::CacheGenKv { + decoded_len: cachegen.decoded_len, + calibration_digest: cachegen.calibration_digest.clone(), + }, + label: "cachegen-kv-archive".to_string(), + }); + offset = kv_bytes; + append_fixed_cuts( + &mut cuts, + &mut offset, + recurrent_bytes, + self.segment_bytes as u64, + SegmentRepresentation::Raw, + ); + } else if native_kv_passthrough { append_fixed_cuts( &mut cuts, &mut offset, @@ -666,6 +783,7 @@ impl L3Tier { ); } let native_kv_passthrough = manifest.uses_native_kv_passthrough(); + let cachegen_kv = manifest.uses_cachegen_kv(); return Ok(Some(L3Location { namespace_key, prefix_key, @@ -674,6 +792,8 @@ impl L3Tier { kv_desc_json: manifest.kv_desc_json.clone(), kv_bytes: manifest.kv_bytes, native_kv_passthrough, + cachegen_kv, + kv_decoded_bytes: manifest.kv_decoded_bytes, })); } Ok(None) @@ -741,11 +861,14 @@ impl L3Tier { } other => bail!("L3 manifest holds unknown payload kind {other}"), }; + let cachegen_kv = manifest.uses_cachegen_kv(); Ok(L3Fill { payload, token_count: manifest.token_count, kv_desc_json: manifest.kv_desc_json, payload_bytes, + cachegen_kv, + kv_decoded_bytes: manifest.kv_decoded_bytes, }) } @@ -785,6 +908,7 @@ impl L3Tier { #[cfg(test)] mod tests { use super::*; + use crate::cachegen::archive::{ComponentLayout, PageLayout, ValueType, encode_page}; use crate::l3::{ CODEC_NATIVE_KV_PAGE, CODEC_NATIVE_KV_PAGE_VERSION, GeometryBlock, GeometryKind, }; @@ -995,6 +1119,101 @@ mod tests { } } + #[test] + fn cachegen_archive_and_exact_recurrent_tail_round_trip_as_typed_envelope() { + let tier = tier("cachegen-envelope", "blake3:cachegen"); + let rows = 4u64; + let raw = (0..16u32) + .flat_map(|value| (value as f32 / 11.0).to_le_bytes()) + .collect::>(); + let layout = PageLayout { + payload_bytes: raw.len() as u64, + components: vec![ComponentLayout { + token_count: rows, + layer_count: 1, + k_type: ValueType::F32, + v_type: ValueType::F32, + k_row_bytes: 8, + v_row_bytes: 8, + v_element_bytes: 4, + k_idx_row_bytes: 0, + payload_offset: 0, + payload_bytes: raw.len() as u64, + v_transposed: false, + }], + }; + let archive = encode_page(&layout, &raw).expect("encode CacheGen page"); + let recurrent = vec![9u8; 17]; + let desc = runtime_kv_desc(0, rows, raw.len() as u64); + let digest = tier + .spill_cachegen_with_cost( + "ns", + &tokens(rows as usize), + &ExactStatePayload::kv_recurrent(raw.clone(), recurrent.clone()), + desc.clone(), + CacheGenKvPayload { + calibration_digest: segment_digest(&archive.bytes), + decoded_len: raw.len() as u64, + archive: archive.bytes.clone(), + }, + None, + ) + .expect("spill CacheGen") + .expect("LRU write-through"); + let manifest = tier.store().load_manifest(&digest).expect("manifest"); + assert!(manifest.uses_cachegen_kv()); + assert!(!manifest.uses_native_kv_passthrough()); + assert_eq!(manifest.codec, Some(PayloadCodec::cachegen_kv_envelope())); + assert_eq!(manifest.kv_decoded_bytes, raw.len() as u64); + assert_eq!(manifest.kv_bytes, archive.bytes.len() as u64); + + let location = tier + .locate_longest("ns", &tokens(rows as usize), 8) + .expect("locate") + .expect("location"); + assert!(location.cachegen_kv); + let fill = tier.load(&location).expect("load"); + assert!(fill.cachegen_kv); + assert_eq!(fill.kv_decoded_bytes, raw.len() as u64); + assert_eq!( + fill.payload.kv_bytes().unwrap().unwrap().as_ref(), + archive.bytes.as_slice() + ); + assert_eq!( + fill.payload.recurrent_state_bytes().unwrap().as_ref(), + recurrent + ); + + let mut mismatched = manifest; + mismatched.kv_decoded_bytes += 4; + assert!( + tier.store().assemble(&mismatched).is_err(), + "CacheGen identity must match the envelope's decoded KV length" + ); + } + + #[test] + fn cachegen_spill_rejects_a_malformed_archive_before_persisting() { + let tier = tier("cachegen-invalid-archive", "blake3:cachegen"); + let raw = vec![0u8; 64]; + let error = tier + .spill_cachegen_with_cost( + "ns", + &tokens(4), + &ExactStatePayload::kv_recurrent(raw.clone(), Vec::new()), + runtime_kv_desc(0, 4, raw.len() as u64), + CacheGenKvPayload { + archive: vec![1, 2, 3, 4], + decoded_len: raw.len() as u64, + calibration_digest: "blake3:invalid".to_string(), + }, + None, + ) + .expect_err("malformed CacheGen bytes must not enter L3"); + assert!(error.to_string().contains("invalid CacheGen archive")); + assert!(tier.store().list_manifests().unwrap().is_empty()); + } + #[test] fn fixed_native_and_recurrent_segments_do_not_cross_representation_boundary() { let tier = tier("native-fixed-boundary", "blake3:native"); diff --git a/crates/skippy-correctness/src/runner/cachegen_gate.rs b/crates/skippy-correctness/src/runner/cachegen_gate.rs index 75e1b63fb6..822f231c68 100644 --- a/crates/skippy-correctness/src/runner/cachegen_gate.rs +++ b/crates/skippy-correctness/src/runner/cachegen_gate.rs @@ -7,12 +7,9 @@ use std::{ }; use anyhow::{Context, Result, bail}; -use skippy_cache::cachegen::archive::{ - CacheGenArchive, ComponentLayout, PageLayout, ValueType, decode_page, encode_page, -}; use skippy_runtime::{ - GGML_TYPE_F16, GGML_TYPE_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, KV_PAGE_FLAG_V_TRANSPOSED, - RuntimeKvPageDesc, StageModel, StageSession, TokenSignal, + GGML_TYPE_F16, GGML_TYPE_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, RuntimeKvPageDesc, StageModel, + StageSession, TokenSignal, decode_cachegen_kv_page, encode_cachegen_kv_page, }; use crate::report::CacheGenGateReport; @@ -51,7 +48,7 @@ pub(in crate::runner) fn run_cachegen_gate( kv_desc.validate_payload(kv.len())?; let encode_started = Instant::now(); - let archive = encode_kv_archive(kv_desc, kv)?; + let archive = encode_cachegen_kv_page(kv_desc, kv)?; let encode_ms = elapsed_ms(encode_started); let native_storage_bytes = kv.len().saturating_add(recurrent.len()); let cachegen_storage_bytes = archive.bytes.len().saturating_add(recurrent.len()); @@ -61,7 +58,7 @@ pub(in crate::runner) fn run_cachegen_gate( // but never feed its multi-gigabyte output into the measured restore. The // accelerated path must consume the persisted archive directly or fail. let scalar_oracle_decode_started = Instant::now(); - let scalar_oracle = decode_kv_archive(kv_desc, &persisted.cachegen_archive)?; + let scalar_oracle = decode_cachegen_kv_page(kv_desc, &persisted.cachegen_archive)?; let scalar_oracle_decode_ms = elapsed_ms(scalar_oracle_decode_started); black_box(&scalar_oracle); drop(scalar_oracle); @@ -311,98 +308,6 @@ fn compare_cachegen_continuation( }) } -fn encode_kv_archive(desc: &RuntimeKvPageDesc, raw: &[u8]) -> Result { - desc.validate_payload(raw.len())?; - encode_page(&page_layout(desc)?, raw) -} - -fn decode_kv_archive(desc: &RuntimeKvPageDesc, archive: &[u8]) -> Result> { - let raw_len = usize::try_from(desc.payload_bytes).context("descriptor length exceeds usize")?; - desc.validate_payload(raw_len)?; - decode_page(archive, raw_len) -} - -fn page_layout(desc: &RuntimeKvPageDesc) -> Result { - let components = if desc.component_count == 0 { - vec![component_layout( - desc.token_count, - desc.layer_count, - desc.k_type, - desc.v_type, - desc.k_row_bytes, - desc.v_row_bytes, - desc.v_element_bytes, - desc.k_idx_row_bytes, - 0, - desc.payload_bytes, - desc.flags, - )?] - } else { - desc.components - .iter() - .take(desc.component_count as usize) - .map(|component| { - component_layout( - component.token_count, - component.layer_count, - component.k_type, - component.v_type, - component.k_row_bytes, - component.v_row_bytes, - component.v_element_bytes, - component.k_idx_row_bytes, - component.payload_offset, - component.payload_bytes, - component.flags, - ) - }) - .collect::>>()? - }; - Ok(PageLayout { - payload_bytes: desc.payload_bytes, - components, - }) -} - -#[allow(clippy::too_many_arguments)] -fn component_layout( - token_count: u64, - layer_count: u32, - k_type: u32, - v_type: u32, - k_row_bytes: u32, - v_row_bytes: u32, - v_element_bytes: u32, - k_idx_row_bytes: u32, - payload_offset: u64, - payload_bytes: u64, - flags: u64, -) -> Result { - Ok(ComponentLayout { - token_count, - layer_count, - k_type: cachegen_value_type(k_type)?, - v_type: cachegen_value_type(v_type)?, - k_row_bytes, - v_row_bytes, - v_element_bytes, - k_idx_row_bytes, - payload_offset, - payload_bytes, - v_transposed: flags & KV_PAGE_FLAG_V_TRANSPOSED != 0, - }) -} - -fn cachegen_value_type(value: u32) -> Result { - match value { - GGML_TYPE_F32 => Ok(ValueType::F32), - GGML_TYPE_F16 => Ok(ValueType::F16), - GGML_TYPE_Q8_0 => Ok(ValueType::Q8_0), - GGML_TYPE_Q4_0 => Ok(ValueType::Q4_0), - _ => bail!("CacheGen gate does not support runtime K/V type {value}"), - } -} - fn cache_type_name(value: u32) -> Result<&'static str> { match value { GGML_TYPE_F32 => Ok("f32"), @@ -522,6 +427,7 @@ fn max_or_zero(values: &[f64]) -> f64 { mod tests { use super::*; use skippy_cache::cachegen::lmcache::MAX_TOKENS_PER_CHUNK; + use skippy_runtime::KV_PAGE_FLAG_V_TRANSPOSED; fn f16_bytes(values: usize) -> Vec { (0..values) @@ -562,8 +468,8 @@ mod tests { fn archive_roundtrip_preserves_geometry_and_length() { let desc = descriptor(0); let raw = f16_bytes(48); - let archive = encode_kv_archive(&desc, &raw).expect("encode"); - let decoded = decode_kv_archive(&desc, &archive.bytes).expect("decode"); + let archive = encode_cachegen_kv_page(&desc, &raw).expect("encode"); + let decoded = decode_cachegen_kv_page(&desc, &archive.bytes).expect("decode"); assert_eq!(decoded.len(), raw.len()); assert_eq!(archive.tile_count, 4); assert_ne!(decoded, raw, "fixture must exercise lossy quantization"); @@ -573,8 +479,8 @@ mod tests { fn transposed_v_layout_is_restored_before_native_import() { let desc = descriptor(KV_PAGE_FLAG_V_TRANSPOSED); let raw = f16_bytes(48); - let archive = encode_kv_archive(&desc, &raw).expect("encode"); - let decoded = decode_kv_archive(&desc, &archive.bytes).expect("decode"); + let archive = encode_cachegen_kv_page(&desc, &raw).expect("encode"); + let decoded = decode_cachegen_kv_page(&desc, &archive.bytes).expect("decode"); assert_eq!(decoded.len(), raw.len()); assert_eq!(archive.tile_count, 4); } @@ -584,7 +490,7 @@ mod tests { let mut desc = descriptor(0); desc.payload_bytes += 2; let raw = f16_bytes(49); - assert!(encode_kv_archive(&desc, &raw).is_err()); + assert!(encode_cachegen_kv_page(&desc, &raw).is_err()); } #[test] @@ -592,8 +498,8 @@ mod tests { let token_count = MAX_TOKENS_PER_CHUNK as u64 + 4; let desc = descriptor_with_tokens(KV_PAGE_FLAG_V_TRANSPOSED, token_count); let raw = f16_bytes(desc.payload_bytes as usize / 2); - let archive = encode_kv_archive(&desc, &raw).expect("encode"); - let decoded = decode_kv_archive(&desc, &archive.bytes).expect("decode"); + let archive = encode_cachegen_kv_page(&desc, &raw).expect("encode"); + let decoded = decode_cachegen_kv_page(&desc, &archive.bytes).expect("decode"); assert_eq!(decoded.len(), raw.len()); assert_eq!(archive.tile_count, 8); } diff --git a/crates/skippy-prompt/src/prompt_cli/stage_config.rs b/crates/skippy-prompt/src/prompt_cli/stage_config.rs index 954541104a..90c0fc5692 100644 --- a/crates/skippy-prompt/src/prompt_cli/stage_config.rs +++ b/crates/skippy-prompt/src/prompt_cli/stage_config.rs @@ -133,6 +133,7 @@ fn prompt_stage_kv_cache_config( max_entries: 1, max_bytes: 0, l2_max_bytes: 0, + codec: skippy_protocol::StageKvCacheCodec::Native, min_tokens: args.kv_page_size_tokens.max(1), shared_prefix_stride_tokens: 128, shared_prefix_record_limit: 2, @@ -147,6 +148,7 @@ fn prompt_stage_kv_cache_config( max_entries: 128, max_bytes, l2_max_bytes: 0, + codec: skippy_protocol::StageKvCacheCodec::Native, min_tokens: args.kv_page_size_tokens.max(1), shared_prefix_stride_tokens: 128, shared_prefix_record_limit: 2, diff --git a/crates/skippy-protocol/src/config.rs b/crates/skippy-protocol/src/config.rs index 61148b0827..c9fa711695 100644 --- a/crates/skippy-protocol/src/config.rs +++ b/crates/skippy-protocol/src/config.rs @@ -329,6 +329,18 @@ pub enum StageKvCachePayload { FullState, } +/// Durable KV representation used below the in-process exact-state cache. +/// CacheGen is opt-in until every backend/dtype combination clears its +/// hardware quality and latency gate. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize, Serialize)] +#[serde(rename_all = "kebab-case")] +pub enum StageKvCacheCodec { + #[default] + Native, + #[serde(rename = "cachegen")] + CacheGen, +} + #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)] pub struct StageKvCacheConfig { #[serde(default = "default_kv_cache_mode")] @@ -343,6 +355,8 @@ pub struct StageKvCacheConfig { /// Zero keeps L2 disabled. #[serde(default)] pub l2_max_bytes: u64, + #[serde(default)] + pub codec: StageKvCacheCodec, #[serde(default = "default_kv_cache_min_tokens")] pub min_tokens: u64, #[serde(default = "default_kv_cache_shared_stride_tokens")] diff --git a/crates/skippy-protocol/src/lib.rs b/crates/skippy-protocol/src/lib.rs index 6d48d08bb1..90f9a89605 100644 --- a/crates/skippy-protocol/src/lib.rs +++ b/crates/skippy-protocol/src/lib.rs @@ -26,8 +26,8 @@ pub use admission::{ pub use config::{ ActivationDType, ActivationDescriptor, ActivationLayout, FlashAttentionType, GlmDsaPolicy, LoadMode, PeerConfig, SplitMode, StageActivationCodec, StageActivationCodecPolicy, StageConfig, - StageDevice, StageIdentity, StageKvCacheConfig, StageKvCacheMode, StageKvCachePayload, - StageTopology, StageTopologyEntry, + StageDevice, StageIdentity, StageKvCacheCodec, StageKvCacheConfig, StageKvCacheMode, + StageKvCachePayload, StageTopology, StageTopologyEntry, }; pub use messages::{ AckMessage, DecodeTokenMessage, ErrorMessage, FinalPrefillChunkMessage, MessageBase, diff --git a/crates/skippy-runtime/src/kv_pages.rs b/crates/skippy-runtime/src/kv_pages.rs index e8a0f90e74..23b3b4cde3 100644 --- a/crates/skippy-runtime/src/kv_pages.rs +++ b/crates/skippy-runtime/src/kv_pages.rs @@ -1,11 +1,113 @@ use std::ptr; use anyhow::{Result, ensure}; -use skippy_cache::cachegen::archive::{RecordKind, ValidatedArchive, validate_archive}; +use skippy_cache::cachegen::archive::{ + CacheGenArchive, ComponentLayout, PageLayout, RecordKind, ValidatedArchive, ValueType, + decode_page, encode_page, validate_archive, +}; use skippy_ffi::{CacheGenRecordV1, KvPageDesc as RawKvPageDesc}; use crate::error::{ensure_ok, free_error}; use crate::session::StageSession; +use crate::{ + GGML_TYPE_F16, GGML_TYPE_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, KV_PAGE_FLAG_V_TRANSPOSED, +}; + +/// Encode one complete runtime KV page into the portable CacheGen archive +/// consumed directly by the native device import ABI. +pub fn encode_cachegen_kv_page(desc: &RuntimeKvPageDesc, raw: &[u8]) -> Result { + desc.validate_payload(raw.len())?; + encode_page(&cachegen_page_layout(desc)?, raw) +} + +/// Scalar oracle used by correctness tooling. Serving restores call the +/// native device importer and never materialize this decoded allocation. +pub fn decode_cachegen_kv_page(desc: &RuntimeKvPageDesc, archive: &[u8]) -> Result> { + let raw_len = usize::try_from(desc.payload_bytes)?; + desc.validate_payload(raw_len)?; + decode_page(archive, raw_len) +} + +fn cachegen_page_layout(desc: &RuntimeKvPageDesc) -> Result { + let components = if desc.component_count == 0 { + vec![cachegen_component_layout( + desc.token_count, + desc.layer_count, + desc.k_type, + desc.v_type, + desc.k_row_bytes, + desc.v_row_bytes, + desc.v_element_bytes, + desc.k_idx_row_bytes, + 0, + desc.payload_bytes, + desc.flags, + )?] + } else { + desc.components + .iter() + .take(desc.component_count as usize) + .map(|component| { + cachegen_component_layout( + component.token_count, + component.layer_count, + component.k_type, + component.v_type, + component.k_row_bytes, + component.v_row_bytes, + component.v_element_bytes, + component.k_idx_row_bytes, + component.payload_offset, + component.payload_bytes, + component.flags, + ) + }) + .collect::>>()? + }; + Ok(PageLayout { + payload_bytes: desc.payload_bytes, + components, + }) +} + +#[allow(clippy::too_many_arguments)] +fn cachegen_component_layout( + token_count: u64, + layer_count: u32, + k_type: u32, + v_type: u32, + k_row_bytes: u32, + v_row_bytes: u32, + v_element_bytes: u32, + k_idx_row_bytes: u32, + payload_offset: u64, + payload_bytes: u64, + flags: u64, +) -> Result { + Ok(ComponentLayout { + token_count, + layer_count, + k_type: cachegen_value_type(k_type)?, + v_type: cachegen_value_type(v_type)?, + k_row_bytes, + v_row_bytes, + v_element_bytes, + k_idx_row_bytes, + payload_offset, + payload_bytes, + v_transposed: flags & KV_PAGE_FLAG_V_TRANSPOSED != 0, + }) +} + +fn cachegen_value_type(value: u32) -> Result { + match value { + GGML_TYPE_F32 => Ok(ValueType::F32), + GGML_TYPE_F16 => Ok(ValueType::F16), + GGML_TYPE_Q8_0 => Ok(ValueType::Q8_0), + GGML_TYPE_Q4_0 => Ok(ValueType::Q4_0), + _ => anyhow::bail!("CacheGen does not support runtime K/V type {value}"), + } +} fn cachegen_records(validated: &ValidatedArchive<'_>) -> Result> { validated diff --git a/crates/skippy-runtime/src/lib.rs b/crates/skippy-runtime/src/lib.rs index cf1151ff2f..291756c58f 100644 --- a/crates/skippy-runtime/src/lib.rs +++ b/crates/skippy-runtime/src/lib.rs @@ -37,6 +37,7 @@ pub(crate) use error::ensure_ok; pub use gguf_writer::{ ModelInfo, SlicePlan, write_gguf_from_parts, write_gguf_metadata_from_parts, }; +pub use kv_pages::{decode_cachegen_kv_page, encode_cachegen_kv_page}; pub use logging::{ LLAMA_LOG_LEVEL_DEBUG, MeasuredNativeBuffers, NativeLogEvent, NativeLogParserMode, NativeLogParserPolicy, configure_native_log_parser, disable_verbose_native_logs, diff --git a/crates/skippy-server/src/binary_transport/stage_execution.rs b/crates/skippy-server/src/binary_transport/stage_execution.rs index 66faf9da68..b791169d08 100644 --- a/crates/skippy-server/src/binary_transport/stage_execution.rs +++ b/crates/skippy-server/src/binary_transport/stage_execution.rs @@ -977,6 +977,7 @@ pub(in crate::binary_transport) fn prefix_cache_test_config() -> StageConfig { max_entries: 8, max_bytes: 0, l2_max_bytes: 0, + codec: skippy_protocol::StageKvCacheCodec::Native, min_tokens: 256, shared_prefix_stride_tokens: 128, shared_prefix_record_limit: 2, diff --git a/crates/skippy-server/src/frontend/local_generation/tests.rs b/crates/skippy-server/src/frontend/local_generation/tests.rs index 3cba546ce4..cd195c979c 100644 --- a/crates/skippy-server/src/frontend/local_generation/tests.rs +++ b/crates/skippy-server/src/frontend/local_generation/tests.rs @@ -148,6 +148,7 @@ fn recurrent_test_backend( max_entries: 8, max_bytes: 0, l2_max_bytes: 0, + codec: skippy_protocol::StageKvCacheCodec::Native, min_tokens: 1, shared_prefix_stride_tokens: 1, shared_prefix_record_limit: 0, diff --git a/crates/skippy-server/src/frontend/prefix_cache.rs b/crates/skippy-server/src/frontend/prefix_cache.rs index d39695595e..c3a6cc1c2c 100644 --- a/crates/skippy-server/src/frontend/prefix_cache.rs +++ b/crates/skippy-server/src/frontend/prefix_cache.rs @@ -1373,6 +1373,7 @@ mod tests { max_entries: 8, max_bytes: 0, l2_max_bytes: 0, + codec: skippy_protocol::StageKvCacheCodec::Native, min_tokens: 1, shared_prefix_stride_tokens: 1, shared_prefix_record_limit: 1, @@ -1451,6 +1452,7 @@ mod tests { max_entries: 8, max_bytes: 0, l2_max_bytes: 0, + codec: skippy_protocol::StageKvCacheCodec::Native, min_tokens: 1, shared_prefix_stride_tokens: 1, shared_prefix_record_limit: 1, diff --git a/crates/skippy-server/src/frontend/tests/support.rs b/crates/skippy-server/src/frontend/tests/support.rs index 15d99b9f48..d121d651bd 100644 --- a/crates/skippy-server/src/frontend/tests/support.rs +++ b/crates/skippy-server/src/frontend/tests/support.rs @@ -48,6 +48,7 @@ pub(super) fn prefix_cache_test_config() -> StageConfig { max_entries: 8, max_bytes: 0, l2_max_bytes: 0, + codec: skippy_protocol::StageKvCacheCodec::Native, min_tokens: 256, shared_prefix_stride_tokens: 128, shared_prefix_record_limit: 2, diff --git a/crates/skippy-server/src/kv_integration/activation.rs b/crates/skippy-server/src/kv_integration/activation.rs index b7f874738e..b4deb8881f 100644 --- a/crates/skippy-server/src/kv_integration/activation.rs +++ b/crates/skippy-server/src/kv_integration/activation.rs @@ -190,6 +190,7 @@ mod tests { max_entries: 8, max_bytes: 0, l2_max_bytes: 0, + codec: skippy_protocol::StageKvCacheCodec::Native, min_tokens: 256, shared_prefix_stride_tokens: 128, shared_prefix_record_limit: 2, diff --git a/crates/skippy-server/src/kv_integration/cache_affinity.rs b/crates/skippy-server/src/kv_integration/cache_affinity.rs index 612164f229..502d611d70 100644 --- a/crates/skippy-server/src/kv_integration/cache_affinity.rs +++ b/crates/skippy-server/src/kv_integration/cache_affinity.rs @@ -110,6 +110,7 @@ mod tests { max_entries: 8, max_bytes: 0, l2_max_bytes: 0, + codec: skippy_protocol::StageKvCacheCodec::Native, min_tokens: 64, shared_prefix_stride_tokens: 32, shared_prefix_record_limit: 2, diff --git a/crates/skippy-server/src/kv_integration/config.rs b/crates/skippy-server/src/kv_integration/config.rs index 54e7158b10..af8e669c25 100644 --- a/crates/skippy-server/src/kv_integration/config.rs +++ b/crates/skippy-server/src/kv_integration/config.rs @@ -10,7 +10,9 @@ use skippy_cache::{ ResidentActivationCache, ResidentCacheConfig, SparseCheckpointPolicy, StoreLimits, UnifiedRadixCache, exact_state_identity_for_stage, numerical_model_identity_for_stage, }; -use skippy_protocol::{StageConfig, StageKvCacheConfig, StageKvCacheMode, StageKvCachePayload}; +use skippy_protocol::{ + StageConfig, StageKvCacheCodec, StageKvCacheConfig, StageKvCacheMode, StageKvCachePayload, +}; use skippy_runtime::{ModelStateKind, RuntimeKvPageDesc}; use super::{ @@ -199,6 +201,7 @@ impl KvStageIntegration { let worker_exact_blobs = exact_blobs.clone(); let worker_l3 = l3.clone(); let worker_l2 = l2.clone(); + let worker_cachegen = cachegen_serving_enabled(config, cache_config.codec); let inflight_records: Arc>> = l3.as_ref().map_or_else( || Arc::new(Mutex::new(BTreeSet::new())), |tier| tier.manager().record_claims(tier.state_identity()), @@ -238,13 +241,16 @@ impl KvStageIntegration { worker_observer.as_ref(), pending, |pending| { - store_exact_radix_record( + store_exact_radix_record_with_codec( &worker_radix, &worker_exact_blobs, exact_max_entries, exact_byte_limits, worker_l2.as_ref(), - worker_l3.as_deref(), + DurableRecordTarget { + l3: worker_l3.as_deref(), + cachegen_enabled: worker_cachegen, + }, pending, ) }, @@ -304,6 +310,7 @@ impl KvStageIntegration { exact_state_record_queue_bytes, l2, l3, + cachegen_serving_enabled: worker_cachegen, inflight_fills, dense_without_recurrent, })) @@ -553,6 +560,7 @@ fn emit_l3_state_transitions(l3: &L3Tier) { } } +#[cfg(test)] fn store_exact_radix_record( radix: &Mutex>, blobs: &Mutex, @@ -562,6 +570,39 @@ fn store_exact_radix_record( l3: Option<&L3Tier>, pending: PendingExactStateRecord, ) -> Result<()> { + store_exact_radix_record_with_codec( + radix, + blobs, + max_entries, + limits, + l2, + DurableRecordTarget { + l3, + cachegen_enabled: false, + }, + pending, + ) +} + +#[derive(Clone, Copy)] +struct DurableRecordTarget<'a> { + l3: Option<&'a L3Tier>, + cachegen_enabled: bool, +} + +fn store_exact_radix_record_with_codec( + radix: &Mutex>, + blobs: &Mutex, + max_entries: usize, + limits: ExactStateByteLimits, + l2: Option<&super::l2_serving::StageL2>, + durable: DurableRecordTarget<'_>, + pending: PendingExactStateRecord, +) -> Result<()> { + let DurableRecordTarget { + l3, + cachegen_enabled, + } = durable; // Write through to the durable tier before the payload is deduplicated // into blocks, while its bytes are still contiguous. Best-effort: a full // or failing disk must not fail the in-memory record. The refusal reason @@ -580,14 +621,63 @@ fn store_exact_radix_record( .kv_desc .as_ref() .and_then(|desc| kv_page_geometry(desc, pending.payload.byte_len())); - let spill = l3.spill_with_cost( - &pending.namespace, - &pending.token_ids, - &pending.payload, - kv_desc_json, - geometry.as_ref(), - pending.l3_cost, - ); + let cachegen_spill = if cachegen_enabled { + match ( + pending.extra.kv_desc.as_ref(), + pending.payload.kv_bytes().ok().flatten(), + ) { + (Some(desc), Some(kv)) + if !kv.is_empty() && cachegen_descriptor_is_qualified(desc) => + { + match skippy_runtime::encode_cachegen_kv_page(desc, kv.as_ref()) { + Ok(archive) if archive.bytes.len() < kv.len() => { + let calibration_digest = skippy_cache::segment_digest(&archive.bytes); + Some(l3.spill_cachegen_with_cost( + &pending.namespace, + &pending.token_ids, + &pending.payload, + kv_desc_json.clone().unwrap_or_default(), + skippy_cache::CacheGenKvPayload { + archive: archive.bytes, + decoded_len: desc.payload_bytes, + calibration_digest, + }, + pending.l3_cost, + )) + } + Ok(_) => None, + Err(error) => { + static WARNED_CACHEGEN: std::sync::atomic::AtomicBool = + std::sync::atomic::AtomicBool::new(false); + if !WARNED_CACHEGEN.swap(true, std::sync::atomic::Ordering::AcqRel) { + let _ = mesh_llm_events::emit_event(OutputEvent::Warning { + message: "CacheGen encode declined; storing native KV page" + .to_string(), + context: Some(format!( + "page_id={} reason={error:#}", + pending.page_id + )), + }); + } + None + } + } + } + _ => None, + } + } else { + None + }; + let spill = cachegen_spill.unwrap_or_else(|| { + l3.spill_with_cost( + &pending.namespace, + &pending.token_ids, + &pending.payload, + kv_desc_json, + geometry.as_ref(), + pending.l3_cost, + ) + }); emit_l3_state_transitions(l3); if let Err(error) = spill { static WARNED: std::sync::atomic::AtomicBool = @@ -684,6 +774,48 @@ fn store_exact_radix_record( Ok(()) } +fn cachegen_serving_enabled(config: &StageConfig, codec: StageKvCacheCodec) -> bool { + if codec != StageKvCacheCodec::CacheGen { + return false; + } + let device = config + .selected_device + .as_ref() + .map(|device| device.backend_device.to_ascii_lowercase()) + .unwrap_or_default(); + let metal = device.contains("metal") || device.starts_with("mtl"); + let key = config.cache_type_k.trim().to_ascii_lowercase(); + let value = config.cache_type_v.trim().to_ascii_lowercase(); + let qualified = matches!( + (key.as_str(), value.as_str()), + ("f32", "f32") | ("f32", "f16") | ("f16", "f32") + ); + if metal && qualified { + return true; + } + let _ = mesh_llm_events::emit_event(OutputEvent::Warning { + message: "CacheGen disk codec unavailable for this model stage; using native KV pages" + .to_string(), + context: Some(format!( + "stage_id={} backend={} cache_type_k={} cache_type_v={}", + config.stage_id, device, config.cache_type_k, config.cache_type_v + )), + }); + false +} + +fn cachegen_descriptor_is_qualified(desc: &RuntimeKvPageDesc) -> bool { + if desc.component_count != 0 { + return false; + } + matches!( + (desc.k_type, desc.v_type), + (skippy_runtime::GGML_TYPE_F32, skippy_runtime::GGML_TYPE_F32) + | (skippy_runtime::GGML_TYPE_F32, skippy_runtime::GGML_TYPE_F16) + | (skippy_runtime::GGML_TYPE_F16, skippy_runtime::GGML_TYPE_F32) + ) +} + /// Evicts least-recently-used exact entries while the catalog holds more than /// `limit_bytes`, never dropping below `min_retained_entries`. A zero limit /// means the catalog has no byte ceiling and nothing is evicted. @@ -822,6 +954,7 @@ fn effective_cache_config(config: &StageConfig) -> Option { max_entries, max_bytes, l2_max_bytes: 0, + codec: skippy_protocol::StageKvCacheCodec::Native, min_tokens, shared_prefix_stride_tokens, shared_prefix_record_limit, @@ -853,7 +986,7 @@ fn parse_cache_mode(value: &str) -> Option { #[cfg(test)] mod tests { use super::*; - use skippy_protocol::{FlashAttentionType, LoadMode}; + use skippy_protocol::{FlashAttentionType, LoadMode, StageDevice}; fn limits(soft_bytes: u64, hard_bytes: u64) -> ExactStateByteLimits { ExactStateByteLimits { @@ -876,6 +1009,184 @@ mod tests { } } + fn kv_descriptor(k_type: u32, v_type: u32, token_count: u64) -> RuntimeKvPageDesc { + let row_bytes = 16_u32; + RuntimeKvPageDesc { + version: 1, + layer_start: 0, + layer_end: 1, + token_start: 0, + token_count, + layer_count: 1, + k_type, + v_type, + k_row_bytes: row_bytes, + v_row_bytes: row_bytes, + v_element_bytes: if v_type == skippy_runtime::GGML_TYPE_F32 { + 4 + } else { + 2 + }, + k_idx_row_bytes: 0, + payload_bytes: token_count * u64::from(row_bytes) * 2, + flags: 0, + codec: 0, + component_count: 0, + components: Box::default(), + } + } + + fn pending_kv( + page_id: &str, + tokens: &[i32], + desc: RuntimeKvPageDesc, + ) -> PendingExactStateRecord { + PendingExactStateRecord { + page_id: page_id.to_string(), + payload: skippy_cache::ExactStatePayload::kv_recurrent( + vec![0; desc.payload_bytes as usize], + Vec::new(), + ), + extra: super::super::ExactStateExtra { + kv_desc: Some(desc), + }, + namespace: "model".to_string(), + token_ids: tokens.to_vec(), + l3_fill_claim: None, + write_through_l3: true, + l2_promotion_digest: None, + l3_cost: None, + } + } + + fn test_l3(name: &str) -> (std::path::PathBuf, L3Tier) { + let root = std::env::temp_dir() + .join("skippy-server-cachegen-tests") + .join(format!("{name}-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&root); + let tier = L3Tier::open(&root, 0, "blake3:test-tier".to_string(), 4096).unwrap(); + (root, tier) + } + + #[test] + fn qualified_cachegen_write_persists_a_compressed_typed_envelope() { + let radix = Mutex::new(UnifiedRadixCache::new()); + let blobs = Mutex::new(CacheBlobStore::new(64)); + let (root, tier) = test_l3("qualified"); + let tokens = (0..256).collect::>(); + let desc = kv_descriptor( + skippy_runtime::GGML_TYPE_F32, + skippy_runtime::GGML_TYPE_F32, + tokens.len() as u64, + ); + + store_exact_radix_record_with_codec( + &radix, + &blobs, + 8, + limits(0, 0), + None, + DurableRecordTarget { + l3: Some(&tier), + cachegen_enabled: true, + }, + pending_kv("cachegen", &tokens, desc.clone()), + ) + .unwrap(); + + let location = tier + .locate_longest("model", &tokens, tokens.len()) + .unwrap() + .expect("CacheGen L3 location"); + assert!(location.cachegen_kv); + assert!(!location.native_kv_passthrough); + assert_eq!(location.kv_decoded_bytes, desc.payload_bytes); + assert!(location.kv_bytes < location.kv_decoded_bytes); + let fill = tier.load(&location).expect("CacheGen L3 fill"); + assert!(fill.cachegen_kv); + let archive = fill.payload.kv_bytes().unwrap().unwrap(); + skippy_cache::cachegen::archive::validate_archive( + archive.as_ref(), + desc.payload_bytes as usize, + ) + .expect("serving write must persist a valid CacheGen archive"); + let _ = std::fs::remove_dir_all(root); + } + + #[test] + fn unqualified_cache_layout_falls_back_to_native_l3() { + let radix = Mutex::new(UnifiedRadixCache::new()); + let blobs = Mutex::new(CacheBlobStore::new(64)); + let (root, tier) = test_l3("fallback"); + let tokens = (0..8).collect::>(); + let desc = kv_descriptor( + skippy_runtime::GGML_TYPE_F16, + skippy_runtime::GGML_TYPE_F16, + tokens.len() as u64, + ); + + store_exact_radix_record_with_codec( + &radix, + &blobs, + 8, + limits(0, 0), + None, + DurableRecordTarget { + l3: Some(&tier), + cachegen_enabled: true, + }, + pending_kv("native", &tokens, desc), + ) + .unwrap(); + + let location = tier + .locate_longest("model", &tokens, tokens.len()) + .unwrap() + .expect("native L3 location"); + assert!(!location.cachegen_kv); + assert!(location.native_kv_passthrough); + assert_eq!(location.kv_bytes, location.kv_decoded_bytes); + let _ = std::fs::remove_dir_all(root); + } + + #[test] + fn cachegen_selector_enables_only_qualified_metal_layouts() { + let mut config = test_config("future/model"); + config.selected_device = Some(StageDevice { + backend_device: "MTL0".to_string(), + stable_id: None, + index: Some(0), + vram_bytes: None, + }); + config.cache_type_k = "f32".to_string(); + config.cache_type_v = "f16".to_string(); + assert!(cachegen_serving_enabled( + &config, + StageKvCacheCodec::CacheGen + )); + + config.cache_type_k = "f16".to_string(); + assert!(!cachegen_serving_enabled( + &config, + StageKvCacheCodec::CacheGen + )); + config.cache_type_k = "f32".to_string(); + config.selected_device.as_mut().unwrap().backend_device = "CUDA0".to_string(); + assert!(!cachegen_serving_enabled( + &config, + StageKvCacheCodec::CacheGen + )); + config.selected_device = None; + assert!(!cachegen_serving_enabled( + &config, + StageKvCacheCodec::CacheGen + )); + assert!(!cachegen_serving_enabled( + &config, + StageKvCacheCodec::Native + )); + } + #[test] fn exact_payloads_live_on_radix_nodes_and_release_deduped_blocks_on_eviction() { let radix = Mutex::new(UnifiedRadixCache::new()); @@ -1607,6 +1918,7 @@ mod tests { max_entries: 512, max_bytes: 0, l2_max_bytes: 0, + codec: skippy_protocol::StageKvCacheCodec::Native, min_tokens: 1, shared_prefix_stride_tokens: 1, shared_prefix_record_limit: 1, diff --git a/crates/skippy-server/src/kv_integration/exact_state.rs b/crates/skippy-server/src/kv_integration/exact_state.rs index baeb519915..82450f3a57 100644 --- a/crates/skippy-server/src/kv_integration/exact_state.rs +++ b/crates/skippy-server/src/kv_integration/exact_state.rs @@ -19,10 +19,10 @@ fn resident_prefix_is_complete(matched_tokens: usize, requested_tokens: usize) - matched_tokens >= requested_tokens } -fn preflight_native_kv_location( +fn preflight_l3_kv_location( location: &skippy_cache::L3Location, ) -> Result> { - if !location.native_kv_passthrough { + if !location.native_kv_passthrough && !location.cachegen_kv { return Ok(None); } let json = location @@ -31,8 +31,12 @@ fn preflight_native_kv_location( .context("native KV manifest has no runtime page descriptor")?; let desc: skippy_runtime::RuntimeKvPageDesc = serde_json::from_str(json).context("native KV manifest has an invalid page descriptor")?; - let kv_bytes = - usize::try_from(location.kv_bytes).context("native KV payload length exceeds usize")?; + let declared_bytes = if location.cachegen_kv { + location.kv_decoded_bytes + } else { + location.kv_bytes + }; + let kv_bytes = usize::try_from(declared_bytes).context("KV payload length exceeds usize")?; desc.validate_payload(kv_bytes) .context("native KV manifest page descriptor is incompatible")?; if desc.token_start != 0 || desc.token_count != location.token_count { @@ -595,7 +599,7 @@ impl KvStageIntegration { // tier reads any segment bytes. Runtime ABI, platform and numerical // mode are already bound by the tier's exact-state identity; the page // descriptor completes the representation check. - let Ok(native_kv_desc) = preflight_native_kv_location(location) else { + let Ok(native_kv_desc) = preflight_l3_kv_location(location) else { return Ok(None); }; let fill_started = Instant::now(); @@ -614,9 +618,11 @@ impl KvStageIntegration { if fill.payload.byte_len() == 0 { return Ok(None); } - if location.native_kv_passthrough + if (location.native_kv_passthrough || location.cachegen_kv) && (fill.token_count != location.token_count - || fill.kv_desc_json != location.kv_desc_json) + || fill.kv_desc_json != location.kv_desc_json + || fill.cachegen_kv != location.cachegen_kv + || fill.kv_decoded_bytes != location.kv_decoded_bytes) { return Ok(None); } @@ -668,7 +674,17 @@ impl KvStageIntegration { // Same fail-closed checks as a radix restore: a // descriptor that does not describe these bytes, or a // page that is not the whole prefix, is a miss. - if desc.validate_payload(kv.len()).is_err() + let payload_valid = if fill.cachegen_kv { + usize::try_from(desc.payload_bytes) + .ok() + .is_some_and(|raw_len| { + skippy_cache::cachegen::archive::validate_archive(kv, raw_len) + .is_ok() + }) + } else { + desc.validate_payload(kv.len()).is_ok() + }; + if !payload_valid || desc.token_start != 0 || desc.token_count != token_count { @@ -682,7 +698,11 @@ impl KvStageIntegration { if let Some((kv, desc)) = kv_page { let import_started = Instant::now(); - runtime.import_kv_page(session_id, desc, kv.as_ref())?; + if fill.cachegen_kv { + runtime.import_cachegen_kv_page(session_id, desc, kv.as_ref())?; + } else { + runtime.import_kv_page(session_id, desc, kv.as_ref())?; + } kv_import_ms = import_started.elapsed().as_secs_f64() * 1000.0; } let import_started = Instant::now(); @@ -699,7 +719,16 @@ impl KvStageIntegration { } _ => return Ok(None), } - let logical_bytes = fill.payload.byte_len(); + let logical_bytes = if fill.cachegen_kv { + fill.kv_decoded_bytes.saturating_add( + fill.payload + .recurrent_state_bytes() + .map(|bytes| bytes.len() as u64) + .unwrap_or_default(), + ) + } else { + fill.payload.byte_len() + }; let payload_kind = fill.payload.kind(); let restore_cost = lookup_started.elapsed().as_secs_f64() * 1_000.0; l3.benefit_observe_l3_restore( @@ -712,6 +741,24 @@ impl KvStageIntegration { // Re-warm the RAM tier off the request path. A drop is fine: the // disk copy stays authoritative. The fill claim rides along so the // worker releases it only once the entry is radix-resident. + if fill.cachegen_kv { + return Ok(Some(ExactStateRestore { + page_id: identity.page_id.clone(), + token_count: token_count as usize, + payload_kind, + logical_bytes, + entries: 0, + reconstruct_ms: 0.0, + reconstruct_bytes: 0, + reconstruct_blocks: 0, + lookup_ms, + kv_import_ms, + recurrent_import_ms, + source: "l3-cachegen", + fill_ms, + rewarm_enqueued: false, + })); + } let l2_promotion_digest = self.l2.as_ref().and_then(|l2| { l2.consider_l3_fill(&location.manifest_key, token_count, fill.payload.byte_len()) .then(|| location.manifest_key.clone()) @@ -846,7 +893,7 @@ mod tests { use skippy_cache::{L3Location, UnifiedRadixCache}; - use super::{preflight_native_kv_location, resident_prefix_is_complete, try_touch_exact_state}; + use super::{preflight_l3_kv_location, resident_prefix_is_complete, try_touch_exact_state}; type TestRadix = UnifiedRadixCache< crate::kv_integration::RadixResidentEntry, @@ -912,6 +959,8 @@ mod tests { kv_desc_json: Some(serde_json::to_string(desc).unwrap()), kv_bytes: desc.payload_bytes, native_kv_passthrough: true, + cachegen_kv: false, + kv_decoded_bytes: desc.payload_bytes, } } @@ -941,7 +990,7 @@ mod tests { fn native_kv_descriptor_is_validated_before_segment_load() { let desc = native_desc(); assert_eq!( - preflight_native_kv_location(&native_location(&desc)).unwrap(), + preflight_l3_kv_location(&native_location(&desc)).unwrap(), Some(desc) ); @@ -949,10 +998,10 @@ mod tests { wrong_length.payload_bytes += 1; let mut location = native_location(&wrong_length); location.kv_bytes -= 1; - assert!(preflight_native_kv_location(&location).is_err()); + assert!(preflight_l3_kv_location(&location).is_err()); let mut wrong_prefix = native_desc(); wrong_prefix.token_start = 1; - assert!(preflight_native_kv_location(&native_location(&wrong_prefix)).is_err()); + assert!(preflight_l3_kv_location(&native_location(&wrong_prefix)).is_err()); } } diff --git a/crates/skippy-server/src/kv_integration/identity.rs b/crates/skippy-server/src/kv_integration/identity.rs index 16f5e045f2..840d0e32d9 100644 --- a/crates/skippy-server/src/kv_integration/identity.rs +++ b/crates/skippy-server/src/kv_integration/identity.rs @@ -226,6 +226,7 @@ mod tests { max_entries: 8, max_bytes: 0, l2_max_bytes: 0, + codec: skippy_protocol::StageKvCacheCodec::Native, min_tokens: 64, shared_prefix_stride_tokens: 32, shared_prefix_record_limit: 2, diff --git a/crates/skippy-server/src/kv_integration/l2_serving.rs b/crates/skippy-server/src/kv_integration/l2_serving.rs index c511e83688..1a90a41dc5 100644 --- a/crates/skippy-server/src/kv_integration/l2_serving.rs +++ b/crates/skippy-server/src/kv_integration/l2_serving.rs @@ -378,6 +378,8 @@ mod tests { kv_desc_json: None, kv_bytes: 0, native_kv_passthrough: false, + cachegen_kv: false, + kv_decoded_bytes: 0, } } diff --git a/crates/skippy-server/src/kv_integration/mod.rs b/crates/skippy-server/src/kv_integration/mod.rs index bdbab91f54..5d5a8b3ac5 100644 --- a/crates/skippy-server/src/kv_integration/mod.rs +++ b/crates/skippy-server/src/kv_integration/mod.rs @@ -192,6 +192,9 @@ pub struct KvStageIntegration { /// Durable L3 floor under the radix cache: exact-state records write /// through to it on the worker, and radix misses fill back from it. pub(crate) l3: Option>, + /// Whether this stage passed the backend and dtype gate for serving-path + /// CacheGen writes. Exposed in status so an opt-in fallback is visible. + pub(crate) cachegen_serving_enabled: bool, /// Manifest keys with an L3 fill in flight. Concurrent misses on one /// stored prefix must not each read it from disk: the loser prefills /// normally while the winner re-warms the radix for everyone. @@ -984,6 +987,10 @@ impl KvStageIntegration { json!(l2.admission_rejects), ), ("skippy.kv.l2.refused_bytes", json!(l2.refused_bytes)), + ( + "skippy.kv.l3.cachegen_enabled", + json!(self.cachegen_serving_enabled), + ), ( "skippy.kv.output_token_entries", json!(output_token_entries), diff --git a/crates/skippy-server/src/kv_integration/resident_prefix.rs b/crates/skippy-server/src/kv_integration/resident_prefix.rs index 05a4dd75cc..c133d6633d 100644 --- a/crates/skippy-server/src/kv_integration/resident_prefix.rs +++ b/crates/skippy-server/src/kv_integration/resident_prefix.rs @@ -825,6 +825,7 @@ mod proactive_eviction_tests { max_entries: 4, max_bytes: 0, l2_max_bytes: 0, + codec: skippy_protocol::StageKvCacheCodec::Native, min_tokens: 1, shared_prefix_stride_tokens: 1, shared_prefix_record_limit: 1, diff --git a/crates/skippy-server/src/runtime_state/lane_lifecycle.rs b/crates/skippy-server/src/runtime_state/lane_lifecycle.rs index 3f77b221ce..9711ab073b 100644 --- a/crates/skippy-server/src/runtime_state/lane_lifecycle.rs +++ b/crates/skippy-server/src/runtime_state/lane_lifecycle.rs @@ -323,6 +323,25 @@ impl RuntimeState { Ok(()) } + pub fn import_cachegen_kv_page( + &mut self, + session_id: &str, + desc: &RuntimeKvPageDesc, + archive: &[u8], + ) -> Result<()> { + let session = self.session(session_id)?; + session.import_cachegen_kv_page(desc, archive)?; + let token_end = desc + .token_start + .checked_add(desc.token_count) + .ok_or_else(|| anyhow::anyhow!("CacheGen KV page token range overflows"))?; + self.session_token_counts + .entry(session_id.to_string()) + .and_modify(|current| *current = (*current).max(token_end)) + .or_insert(token_end); + Ok(()) + } + pub fn save_resident_prefix( &mut self, session_id: &str, diff --git a/docs/skippy/CACHEGEN_BACKEND_PLAN.md b/docs/skippy/CACHEGEN_BACKEND_PLAN.md index 54880cfdd8..11db8c9536 100644 --- a/docs/skippy/CACHEGEN_BACKEND_PLAN.md +++ b/docs/skippy/CACHEGEN_BACKEND_PLAN.md @@ -1,8 +1,9 @@ # CacheGen Backend Qualification (#1652) Status: **LMCache-compatible direct Metal restore passes the local F32/F32 and -F32+F16 gates; F16 and the lower-width sampled types remain stopped on local -latency, and CacheGen remains opt-in pending selection-path wiring**. Owner: +F32+F16 gates; those layouts are wired into the opt-in disk-cache selection +path, while F16 and the lower-width sampled types remain stopped on local +latency**. Owner: jian yang. Reviewed against: #1652 scope, scama's directives of 2026-09-10 (v4 contract, CPU+Metal parity, six measurements, stop rule). @@ -118,10 +119,13 @@ container. Fixtures generated by the independent Python scalar transcription pin both 16-bin and 32-bin streams, and normal Rust tests require byte-for-byte encoder agreement and decoder agreement with those fixtures. -This remains an opt-in correctness path and is not selected by storage or the -request path. The matched 19K result below proves the reference recovers -continuation quality, while also proving that scalar arithmetic decode cannot -meet the restore-to-first-token gate. +The scalar decoder remains a correctness oracle and is never selected on the +request restore path. The opt-in disk-cache selector now runs the matching +encoder on the bounded record worker for qualified Metal layouts, persists the +portable archive, and dispatches restore to the native device importer. The +matched 19K result below proves the reference recovers continuation quality, +while also proving that scalar arithmetic decode cannot meet the +restore-to-first-token gate. ## Sequencing after the direct Metal gate @@ -136,8 +140,8 @@ meet the restore-to-first-token gate. 4. Add independently decodable substreams or an equivalent parallel entropy layout before retrying local-tier Metal promotion. Preserve the current scalar stream as the compatibility oracle for the new revision. -5. Add encode from resident K/V storage and copy only the compact archive back - to the persistence layer. +5. Move the now-wired worker-side encode into resident device storage so only + the compact archive crosses back to the persistence layer. 6. Run separate quality and latency gates for Q8_0/Q8_0, Q4_0/Q4_0, and the supported mixed K/V combinations now that every typed record is wired into the native backend contract. diff --git a/docs/skippy/CONFIGURATION.md b/docs/skippy/CONFIGURATION.md index cd3f28c851..d967f75e09 100644 --- a/docs/skippy/CONFIGURATION.md +++ b/docs/skippy/CONFIGURATION.md @@ -77,6 +77,7 @@ website configuration reference, with the same `Wiring status`. | #1576 | Disk cache directory | `runtime.kv_cache.disk.directory` | P0 | host runtime | node-scoped cache root | single-stage, staged | process restart | `$MESH_LLM_HOME/kv-cache` | explicit values must be absolute | `#node-local-disk-prompt-cache` | `mesh-llm-config` validation tests | Directory changes do not migrate or delete old data | wired | | #1576 | Fixed disk budget | `runtime.kv_cache.disk.budget_mib` | P0 | host runtime / `skippy-cache` | manager hard byte budget | single-stage, staged | applies dynamically | unset | required and > 0 only in fixed mode | `#node-local-disk-prompt-cache` | config precedence and manager limit tests | One shared physical cap per node root; never unbounded | wired | | #1576 | Minimum free storage | `runtime.kv_cache.disk.minimum_free_mib` | P0 | host runtime / `skippy-cache` | manager free-space reserve | single-stage, staged | applies dynamically | `16384` MiB | at least `1024` MiB | `#node-local-disk-prompt-cache` | config precedence and low-space transition tests | Writes decline while fills remain available | wired | +| #1652 | Disk cache codec | `runtime.kv_cache.disk.codec` | P1 | host runtime / `skippy-cache` / native runtime | per-stage durable write and restore | single-stage, staged | process restart | `native` | enum `native` or `cachegen`; CacheGen activates only for qualified Metal KV layouts | `#node-local-disk-prompt-cache` | config, archive, serving-worker, and native import tests | Unqualified or failed encodes write exact native pages | wired | ## Model fit, context, and KV cache diff --git a/docs/skippy/KV_CACHE_DISK.md b/docs/skippy/KV_CACHE_DISK.md index d85103e9b4..e312dd7ec1 100644 --- a/docs/skippy/KV_CACHE_DISK.md +++ b/docs/skippy/KV_CACHE_DISK.md @@ -36,8 +36,10 @@ the prompt cache. Turn it on deliberately. ## Configuration surface -The same four settings are expressible through a config file, environment -variables, and CLI flags. Sizes everywhere use **explicit IEC suffixes** +The capacity and location settings are expressible through a config file, +environment variables, and CLI flags. The typed payload codec is configured in +the file so its process-restart boundary is explicit. Sizes everywhere use +**explicit IEC suffixes** (`KiB`, `MiB`, `GiB`, `TiB`) on a positive whole number — for example `32GiB`. Bare numbers and decimal/`GB`-style units are rejected. @@ -46,6 +48,7 @@ Bare numbers and decimal/`GB`-style units are rejected. | Mode / fixed budget | `mode` (`off`/`auto`/`fixed`) + `budget_mib` | `MESH_LLM_KV_CACHE_DISK` (`off`/`auto`/`SIZE`) | `--kv-cache-disk off\|auto\|SIZE` | | Directory | `directory` (absolute) | `MESH_LLM_KV_CACHE_DISK_DIR` (absolute) | `--kv-cache-disk-dir ABSOLUTE_PATH` | | Minimum free reserve | `minimum_free_mib` | `MESH_LLM_KV_CACHE_MIN_FREE` (SIZE) | `--kv-cache-min-free SIZE` | +| Payload codec | `codec` (`native`/`cachegen`) | — | — | Notes: @@ -58,6 +61,9 @@ Notes: - Directories must be absolute. Relative paths (including bare-drive forms like `C:\cache` on non-Windows hosts) fail closed rather than resolving under the working directory. +- `codec = "native"` is the default. `codec = "cachegen"` is opt-in and only + writes CacheGen archives for explicitly qualified Metal F32/F32, F32/F16, + and F16/F32 KV layouts. Every other backend or layout writes native pages. ### Precedence diff --git a/tools/xtask/data/console_print_allowlist.json b/tools/xtask/data/console_print_allowlist.json index 6f9e188845..edb50682ac 100644 --- a/tools/xtask/data/console_print_allowlist.json +++ b/tools/xtask/data/console_print_allowlist.json @@ -4251,19 +4251,19 @@ ], "crates/skippy-prompt/src/prompt_cli/stage_config.rs": [ { - "line": 360, + "line": 362, "macro_name": "eprintln!" }, { - "line": 400, + "line": 402, "macro_name": "eprintln!" }, { - "line": 466, + "line": 468, "macro_name": "eprintln!" }, { - "line": 480, + "line": 482, "macro_name": "eprint!" } ], diff --git a/website/src/docs/pages/config-reference.md b/website/src/docs/pages/config-reference.md index 91630c91c5..58435cd521 100644 --- a/website/src/docs/pages/config-reference.md +++ b/website/src/docs/pages/config-reference.md @@ -121,6 +121,7 @@ produces a clear startup error rather than a partial start. | `runtime.kv_cache.disk.directory` | absolute path | `$MESH_LLM_HOME/kv-cache` | node-level | process restart | wired | `--kv-cache-disk-dir` | | `runtime.kv_cache.disk.budget_mib` | integer | required and > 0 only for `fixed` | node-level | applies dynamically | wired | fixed size passed to `--kv-cache-disk` | | `runtime.kv_cache.disk.minimum_free_mib` | integer | `16384`; minimum `1024` | node-level | applies dynamically | wired | `--kv-cache-min-free` | +| `runtime.kv_cache.disk.codec` | enum | `native` (default), `cachegen` | node-level | process restart | wired | config only | | `runtime.model_target_demand_upgrade_min_requests` | integer | `2` | node-level | process restart | wired | none | | `runtime.model_target_demand_upgrade_max_age_secs` | integer | `3600` | node-level | process restart | wired | none | | `advanced.server.alias` | string | unset; per-model alias overrides the default | both | model reload | wired; becomes the served identity used by `/v1/models` and routing | none | diff --git a/website/src/docs/pages/kv-caching.md b/website/src/docs/pages/kv-caching.md index df79ac2005..f7259a04b0 100644 --- a/website/src/docs/pages/kv-caching.md +++ b/website/src/docs/pages/kv-caching.md @@ -171,6 +171,7 @@ For an automatically sized cache: mode = "auto" directory = "/var/lib/mesh-llm/kv-cache" minimum_free_mib = 16384 +codec = "native" ``` For a fixed 32 GiB cap: @@ -181,6 +182,7 @@ mode = "fixed" directory = "/var/lib/mesh-llm/kv-cache" budget_mib = 32768 minimum_free_mib = 16384 +codec = "cachegen" ``` The directory must be absolute. If omitted, it is @@ -208,6 +210,11 @@ Invalid disk settings stop startup with a configuration error. Once a valid configuration is running, storage trouble fails open: Mesh logs the problem and serves the request with cold prefill. +`codec = "native"` is the default. The opt-in `cachegen` codec currently +activates only for the measured Metal F32/F32, F32/F16, and F16/F32 KV layouts; +other backends and layouts persist exact native pages. CacheGen archives are +encoded on the cache worker and restored directly into the native runtime. + ## Inspect and maintain the disk cache ```bash From 7a53fd907f26def0e2c03e0b63455da5df6cc385 Mon Sep 17 00:00:00 2001 From: scama Date: Tue, 15 Sep 2026 08:32:11 +1000 Subject: [PATCH 08/16] feat(skippy): derive KV defaults from publisher metadata --- .../src/inference/skippy/deployment.rs | 2 +- .../src/inference/skippy/family_policy.rs | 21 +- .../src/inference/skippy/kv_cache.rs | 128 +++++--- .../src/inference/skippy/mod.rs | 6 +- .../src/inference/skippy/package.rs | 11 + .../src/inference/skippy/resolver.rs | 1 + .../inference/skippy/resolver/resolution.rs | 87 ++--- .../inference/skippy/resolver/test_support.rs | 1 + .../src/inference/skippy/resolver/tests.rs | 70 ++++- .../src/inference/skippy/resolver/types.rs | 5 +- .../inference/skippy/split_certification.rs | 2 + .../src/runtime/local.rs | 14 +- .../src/runtime/local_package.rs | 42 +-- .../src/runtime/local_split/loading.rs | 7 +- .../src/runtime/local_split/test_support.rs | 41 ++- .../src/runtime/local_split/tests.rs | 67 ++-- .../src/runtime/split_planning.rs | 1 + .../src/runtime/stage_admission/tests.rs | 2 + crates/model-package/src/script.rs | 9 + .../src/scripts/split-model-job.sh | 51 +++ crates/skippy-model-package/README.md | 15 +- crates/skippy-model-package/src/cli.rs | 5 + crates/skippy-model-package/src/main.rs | 6 +- crates/skippy-model-package/src/package_v2.rs | 296 +++++++++++++++++- .../src/package_v2/tests.rs | 106 ++++++- crates/skippy-model-package/src/preflight.rs | 4 +- crates/skippy-model-package/src/verify_v2.rs | 16 +- .../src/verify_v2/tests.rs | 7 +- crates/skippy-model/src/package_carrier.rs | 6 + crates/skippy-package-format/src/lib.rs | 202 ++++++++++++ .../src/materialization.rs | 2 + .../src/stage_admission.rs | 2 + crates/skippy-package-format/src/tests.rs | 2 + docs/LAYER_PACKAGE_REPOS.md | 8 + docs/design/SKIPPY_PACKAGE_V2_SCHEMA.md | 31 +- tools/xtask/data/console_print_allowlist.json | 18 +- 36 files changed, 1027 insertions(+), 267 deletions(-) diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/deployment.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/deployment.rs index 38bc2884cb..0972a32dee 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/deployment.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/deployment.rs @@ -296,7 +296,7 @@ mod tests { continuous_batching: true, n_batch: None, n_ubatch: None, - kv_cache: KvCachePolicy::for_model_size(0), + kv_cache: KvCachePolicy::safe_default(), flash_attn_type: FlashAttentionType::Auto, kv_offload: None, kv_unified: None, diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs index e30b8b7807..b179c5d2e7 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs @@ -9,7 +9,6 @@ const DEFAULT_PREFIX_CACHE_MAX_ENTRIES: usize = 512; #[derive(Clone, Debug, Eq, PartialEq)] pub(crate) struct FamilyPolicy { - pub(crate) default_kv_cache_type: Option<&'static str>, pub(crate) prefix_cache: FamilyPrefixCachePolicy, } @@ -109,14 +108,8 @@ pub(crate) fn family_policy_for_model_path(path: impl AsRef) -> FamilyPoli generic_model_policy(metadata.as_ref()) } -fn generic_model_policy(meta: Option<&GgufCompactMeta>) -> FamilyPolicy { - // Inkling requires q4_0 native KV storage. This is read from the GGUF - // architecture field, never inferred from a repository or filename. - let default_kv_cache_type = meta - .is_some_and(|meta| meta.architecture == "inkling") - .then_some("q4_0"); +fn generic_model_policy(_meta: Option<&GgufCompactMeta>) -> FamilyPolicy { FamilyPolicy { - default_kv_cache_type, prefix_cache: FamilyPrefixCachePolicy::Auto { min_tokens: DEFAULT_PREFIX_CACHE_MIN_TOKENS, max_entries: DEFAULT_PREFIX_CACHE_MAX_ENTRIES, @@ -248,18 +241,6 @@ mod tests { } } - #[test] - fn gguf_architecture_only_controls_required_native_kv_storage_type() { - assert_eq!( - family_policy_for_compact_meta(&kv_meta("inkling")).default_kv_cache_type, - Some("q4_0") - ); - assert_eq!( - family_policy_for_compact_meta(&kv_meta("nemotron_h_moe")).default_kv_cache_type, - None - ); - } - #[test] fn stage_cache_cap_tracks_ctx_layers_and_kv_types() { let config = stage_config(); diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/kv_cache.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/kv_cache.rs index e4359e3187..fbafaa5459 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/kv_cache.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/kv_cache.rs @@ -40,31 +40,38 @@ pub(crate) struct KvCachePolicy { } impl KvCachePolicy { - const LARGE_MODEL_MIN_BYTES: u64 = 50 * 1024 * 1024 * 1024; + /// Conservative live-KV default when the package has no validated + /// publisher declaration. Model weight bytes and weight quantisation are + /// intentionally not inputs to this decision. + pub(crate) fn safe_default() -> Self { + Self { + k_type: KvCacheType::F16, + v_type: KvCacheType::F16, + } + } - /// Default KV cache policy, tiered by model size. - /// - /// Models >= 50 GB use Q4_0 K + Q4_0 V to keep KV cache small enough - /// that unified-memory machines don't thrash. On a 480B MoE split - /// across two Apple Silicon nodes the difference between Q8_0 and Q4_0 - /// is the difference between swap-thrashing at 1 tok/s and running at - /// 20+ tok/s. - /// - /// Smaller models use Q8_0 K + Q8_0 V which gives ~2× compression over - /// f16 with negligible quality loss. - /// - /// Users can override via `--cache-type-k` / `--cache-type-v`. - pub(crate) fn for_model_size(model_bytes: u64) -> Self { - if model_bytes >= Self::LARGE_MODEL_MIN_BYTES { - Self { - k_type: KvCacheType::Q4_0, - v_type: KvCacheType::Q4_0, - } - } else { - Self { + pub(crate) fn from_publisher_defaults( + defaults: Option<&skippy_package_format::PublisherModelDefaults>, + ) -> Self { + use skippy_package_format::PublisherDtype; + + let dtype = defaults + .and_then(|defaults| defaults.kv_cache_dtype.as_ref()) + .or_else(|| defaults.and_then(|defaults| defaults.compute_dtype.as_ref())) + .map(|declaration| declaration.dtype); + match dtype { + Some(PublisherDtype::Q8_0) => Self { k_type: KvCacheType::Q8_0, v_type: KvCacheType::Q8_0, - } + }, + Some(PublisherDtype::Q4_0) => Self { + k_type: KvCacheType::Q4_0, + v_type: KvCacheType::Q4_0, + }, + // BF16 maps conservatively to F16 until BF16 live KV is qualified. + // FP8 and F32 declarations also fall back because the embedded + // runtime does not currently expose a qualified matching type. + _ => Self::safe_default(), } } @@ -74,8 +81,8 @@ impl KvCachePolicy { /// Downgrade this *default* policy to one the model can actually load. /// - /// The size tiers above choose a quant purely from byte size; they do not - /// know whether the model satisfies llama.cpp's quantised-KV constraints + /// Publisher metadata does not know whether the model satisfies + /// llama.cpp's quantised-KV constraints /// (Flash Attention availability, per-head block alignment). Without this /// guard, an incompatible model (e.g. Grok, or a head_dim not divisible by /// the q8_0/q4_0 block size of 32) fails the context build outright rather @@ -121,19 +128,56 @@ impl KvCachePolicy { #[cfg(test)] mod tests { use super::*; + use skippy_package_format::{ + PublisherDtype, PublisherDtypeDeclaration, PublisherModelDefaults, + }; + + fn publisher_defaults( + kv_cache_dtype: Option, + compute_dtype: Option, + ) -> PublisherModelDefaults { + let declaration = |dtype| PublisherDtypeDeclaration { + dtype, + artifact_id: "publisher-config-json".to_string(), + json_path: "/dtype".to_string(), + }; + PublisherModelDefaults { + compute_dtype: compute_dtype.map(declaration), + kv_cache_dtype: kv_cache_dtype.map(declaration), + } + } #[test] - fn small_model_uses_q8_0() { - let policy = KvCachePolicy::for_model_size(10 * 1024 * 1024 * 1024); - assert_eq!(policy.k_type, KvCacheType::Q8_0); - assert_eq!(policy.v_type, KvCacheType::Q8_0); + fn missing_publisher_metadata_uses_f16() { + let policy = KvCachePolicy::from_publisher_defaults(None); + assert_eq!(policy.k_type, KvCacheType::F16); + assert_eq!(policy.v_type, KvCacheType::F16); + } + + #[test] + fn publisher_compute_bf16_maps_to_supported_f16() { + let defaults = publisher_defaults(None, Some(PublisherDtype::Bf16)); + assert_eq!( + KvCachePolicy::from_publisher_defaults(Some(&defaults)), + KvCachePolicy::safe_default() + ); + } + + #[test] + fn unqualified_publisher_fp8_falls_back_to_f16() { + let defaults = publisher_defaults(Some(PublisherDtype::Fp8), Some(PublisherDtype::Bf16)); + assert_eq!( + KvCachePolicy::from_publisher_defaults(Some(&defaults)), + KvCachePolicy::safe_default() + ); } #[test] - fn large_model_uses_q4_0() { - let policy = KvCachePolicy::for_model_size(50 * 1024 * 1024 * 1024); - assert_eq!(policy.k_type, KvCacheType::Q4_0); - assert_eq!(policy.v_type, KvCacheType::Q4_0); + fn explicit_publisher_kv_declaration_precedes_compute_dtype() { + let defaults = publisher_defaults(Some(PublisherDtype::Q8_0), Some(PublisherDtype::Bf16)); + let policy = KvCachePolicy::from_publisher_defaults(Some(&defaults)); + assert_eq!(policy.k_type, KvCacheType::Q8_0); + assert_eq!(policy.v_type, KvCacheType::Q8_0); } fn meta(architecture: &str, head_dim: u32) -> crate::models::gguf::GgufCompactMeta { @@ -150,14 +194,20 @@ mod tests { #[test] fn guard_keeps_quant_for_block_aligned_model() { - let policy = KvCachePolicy::for_model_size(10 * 1024 * 1024 * 1024); + let policy = KvCachePolicy { + k_type: KvCacheType::Q8_0, + v_type: KvCacheType::Q8_0, + }; let guarded = policy.guarded_for_model(Some(&meta("qwen3", 128))); assert_eq!(guarded, policy); } #[test] fn guard_falls_back_to_f16_for_unaligned_head_dim() { - let policy = KvCachePolicy::for_model_size(10 * 1024 * 1024 * 1024); + let policy = KvCachePolicy { + k_type: KvCacheType::Q8_0, + v_type: KvCacheType::Q8_0, + }; let guarded = policy.guarded_for_model(Some(&meta("phi2", 80))); assert_eq!(guarded.k_type, KvCacheType::F16); assert_eq!(guarded.v_type, KvCacheType::F16); @@ -165,7 +215,10 @@ mod tests { #[test] fn guard_falls_back_to_f16_for_grok() { - let policy = KvCachePolicy::for_model_size(60 * 1024 * 1024 * 1024); + let policy = KvCachePolicy { + k_type: KvCacheType::Q4_0, + v_type: KvCacheType::Q4_0, + }; let guarded = policy.guarded_for_model(Some(&meta("grok", 128))); assert_eq!(guarded.k_type, KvCacheType::F16); assert_eq!(guarded.v_type, KvCacheType::F16); @@ -173,7 +226,10 @@ mod tests { #[test] fn guard_is_noop_without_metadata() { - let policy = KvCachePolicy::for_model_size(10 * 1024 * 1024 * 1024); + let policy = KvCachePolicy { + k_type: KvCacheType::Q8_0, + v_type: KvCacheType::Q8_0, + }; assert_eq!(policy.guarded_for_model(None), policy); } } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs index 1065c6ae77..c6e78ee54b 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs @@ -51,9 +51,7 @@ use skippy_server::{ pub use certification::{ CertificationGateStatus, SkippyCertificationRequest, certify_layer_package, }; -pub(crate) use family_policy::{ - family_policy_for_compact_meta, family_policy_for_model_path, family_policy_for_stage_config, -}; +pub(crate) use family_policy::{family_policy_for_model_path, family_policy_for_stage_config}; pub(crate) use hooks::MeshAutoHookPolicy; pub(crate) use kv_cache::KvCachePolicy; #[cfg(test)] @@ -85,6 +83,7 @@ pub(crate) use package::{ pub(crate) use resolver::{ ResolvedEmbeddedOpenAiArgs, ResolvedSkippyConfig, SkippyConfigResolveRequest, effective_safety_margin_bytes, resolve_skippy_config_for_selector, + resolve_skippy_config_for_selector_with_publisher_defaults, }; pub(crate) use skippy_server::OpenAiGuardrailsStatus as SkippyOpenAiGuardrailsStatus; pub(crate) use split_certification::{require_split_certification, split_certification_label}; @@ -1546,6 +1545,7 @@ mod tests { activation_width: 4096, tensor_count: 100, generation: None, + publisher_defaults: None, } } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/package.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/package.rs index 438de395cf..f4073d057e 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/package.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/package.rs @@ -95,6 +95,8 @@ pub(crate) fn write_test_package_v2_fixture( entries: Vec::new(), }, sidecars: Vec::new(), + publisher_metadata: Vec::new(), + publisher_defaults: None, generation: None, native_abi_version: format!( "{}.{}.{}", @@ -259,6 +261,7 @@ pub struct SkippyPackageIdentity { pub activation_width: u32, pub tensor_count: u64, pub generation: Option, + pub publisher_defaults: Option, } #[derive(Clone, Debug, Eq, PartialEq, Serialize)] @@ -401,6 +404,7 @@ pub fn identity_from_package_v2(package_dir: &Path) -> Result Result, config_model_id: Option<&str>, ) -> Result { - resolve_skippy_config_with_context(ResolverContext::new_for_selector(request, config_model_id)) + resolve_skippy_config_for_selector_with_publisher_defaults(request, config_model_id, None) +} + +pub(crate) fn resolve_skippy_config_for_selector_with_publisher_defaults( + request: SkippyConfigResolveRequest<'_>, + config_model_id: Option<&str>, + publisher_defaults: Option<&skippy_package_format::PublisherModelDefaults>, +) -> Result { + resolve_skippy_config_with_context(ResolverContext::new_for_selector( + request, + config_model_id, + publisher_defaults, + )) } fn resolve_skippy_config_with_context( @@ -45,16 +57,10 @@ fn resolve_skippy_config_with_context( validate_supported_model_fit_controls(&context)?; validate_supported_hardware_controls(&context)?; - // Guard the size-tiered default so a model that cannot load quantised KV - // (Flash Attention off, or a head_dim not divisible by the block size) - // resolves to f16 instead of failing the context build. Explicit config / - // family defaults still take precedence in resolve_cache_type_* below and - // are intentionally not guarded here. - let kv_policy = KvCachePolicy::for_model_size(context.request.model_bytes) + let kv_policy = KvCachePolicy::from_publisher_defaults(context.publisher_defaults) .guarded_for_model(context.request.compact_meta); let hardware = resolve_hardware_config(&context)?; - let family_policy = family_policy_for_model_path(&hardware.resolved_model_path); - let model_fit = resolve_model_fit_config(&context, kv_policy, &family_policy)?; + let model_fit = resolve_model_fit_config(&context, kv_policy)?; let throughput = resolve_throughput_config(&context); let skippy = resolve_execution_config(&context); let speculative = resolve_speculative_config( @@ -147,12 +153,14 @@ struct ResolverContext<'a> { global_model_fit: Option<&'a ModelFitConfig>, model_throughput: Option<&'a ThroughputConfig>, global_throughput: Option<&'a ThroughputConfig>, + publisher_defaults: Option<&'a skippy_package_format::PublisherModelDefaults>, } impl<'a> ResolverContext<'a> { fn new_for_selector( request: SkippyConfigResolveRequest<'a>, config_model_id: Option<&str>, + publisher_defaults: Option<&'a skippy_package_format::PublisherModelDefaults>, ) -> Self { let model_entry = config_model_id.and_then(|selector| { request @@ -161,12 +169,13 @@ impl<'a> ResolverContext<'a> { .iter() .find(|entry| entry.model == selector) }); - Self::with_model_entry(request, model_entry) + Self::with_model_entry(request, model_entry, publisher_defaults) } fn with_model_entry( request: SkippyConfigResolveRequest<'a>, model_entry: Option<&'a ModelConfigEntry>, + publisher_defaults: Option<&'a skippy_package_format::PublisherModelDefaults>, ) -> Self { let mesh_config = request.mesh_config; let defaults = mesh_config.defaults.as_ref(); @@ -183,6 +192,7 @@ impl<'a> ResolverContext<'a> { global_model_fit, model_throughput, global_throughput, + publisher_defaults, } } } @@ -210,7 +220,6 @@ fn validate_supported_hardware_controls(context: &ResolverContext<'_>) -> Result fn resolve_model_fit_config( context: &ResolverContext<'_>, kv_policy: KvCachePolicy, - family_policy: &super::super::family_policy::FamilyPolicy, ) -> Result { let kv = resolve_kv_defaults(context, kv_policy); let throughput = resolve_throughput_defaults(context); @@ -246,8 +255,8 @@ fn resolve_model_fit_config( .and_then(|defaults| defaults.ubatch), BUILTIN_UBATCH, ); - let cache_type_k = resolve_cache_type_k(context, &kv, kv_policy, family_policy); - let cache_type_v = resolve_cache_type_v(context, &kv, kv_policy, family_policy); + let cache_type_k = resolve_cache_type_k(context, &kv, kv_policy); + let cache_type_v = resolve_cache_type_v(context, &kv, kv_policy); let kv_offload = resolve_kv_offload(context, &kv); let kv_offload_resolved = parse_kv_offload_string(&kv_offload); let kv_unified = resolve_kv_unified(context)?; @@ -330,7 +339,13 @@ fn resolve_kv_defaults(context: &ResolverContext<'_>, kv_policy: KvCachePolicy) let global_policy = context .global_model_fit .and_then(|fit| fit.kv_cache_policy.as_deref()); - let effective_policy = pick_string(model_policy, global_policy, Some("balanced")); + let effective_policy = model_policy.or(global_policy).unwrap_or_else(|| { + if context.publisher_defaults.is_some() { + "publisher" + } else { + "safe_f16" + } + }); KvDefaults { effective_policy: effective_policy.to_string(), @@ -339,30 +354,10 @@ fn resolve_kv_defaults(context: &ResolverContext<'_>, kv_policy: KvCachePolicy) } } -fn guarded_family_default_kv_cache_type( - context: &ResolverContext<'_>, - family_policy: &super::super::family_policy::FamilyPolicy, -) -> Option<&'static str> { - family_policy - .default_kv_cache_type - .and_then(|default| { - crate::models::gguf::GgufKvCacheQuant::from_llama_args(default, default) - }) - .map(|quant| { - context - .request - .compact_meta - .map(|meta| meta.compatible_default_kv_cache_quant(quant)) - .unwrap_or(quant) - }) - .map(|quant| quant.k.as_llama_arg()) -} - fn resolve_cache_type_k( context: &ResolverContext<'_>, kv: &KvDefaults, kv_policy: KvCachePolicy, - family_policy: &super::super::family_policy::FamilyPolicy, ) -> String { if let Some(explicit) = context .model_fit @@ -370,18 +365,6 @@ fn resolve_cache_type_k( { return explicit.to_string(); } - // Guard the family default against the model's quantised-KV compatibility - // so an unloadable family default degrades to f16 instead of failing the - // context build. Explicit config above and below stays unguarded. - if let Some(family_default) = guarded_family_default_kv_cache_type(context, family_policy) { - if let Some(explicit) = context - .global_model_fit - .and_then(|fit| non_auto_string(fit.cache_type_k.as_deref())) - { - return explicit.to_string(); - } - return family_default.to_string(); - } resolve_field_string( None, kv.model_macro @@ -401,7 +384,6 @@ fn resolve_cache_type_v( context: &ResolverContext<'_>, kv: &KvDefaults, kv_policy: KvCachePolicy, - family_policy: &super::super::family_policy::FamilyPolicy, ) -> String { if let Some(explicit) = context .model_fit @@ -409,15 +391,6 @@ fn resolve_cache_type_v( { return explicit.to_string(); } - if let Some(family_default) = guarded_family_default_kv_cache_type(context, family_policy) { - if let Some(explicit) = context - .global_model_fit - .and_then(|fit| non_auto_string(fit.cache_type_v.as_deref())) - { - return explicit.to_string(); - } - return family_default.to_string(); - } resolve_field_string( None, kv.model_macro diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/test_support.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/test_support.rs index cc7b92f268..c6198212f9 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/test_support.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/test_support.rs @@ -94,6 +94,7 @@ pub(super) fn fake_package_identity(layer_count: u32) -> SkippyPackageIdentity { activation_width: 4096, tensor_count: 100, generation: None, + publisher_defaults: None, } } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs index 2459917ade..f2c7735943 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs @@ -29,6 +29,64 @@ fn resolve_qwen_config_with_request_defaults( .expect("qwen config should resolve") } +fn publisher_q8_defaults() -> skippy_package_format::PublisherModelDefaults { + skippy_package_format::PublisherModelDefaults { + compute_dtype: Some(skippy_package_format::PublisherDtypeDeclaration { + dtype: skippy_package_format::PublisherDtype::Bf16, + artifact_id: "publisher-config-json".to_string(), + json_path: "/torch_dtype".to_string(), + }), + kv_cache_dtype: Some(skippy_package_format::PublisherDtypeDeclaration { + dtype: skippy_package_format::PublisherDtype::Q8_0, + artifact_id: "publisher-hf-quant-config-json".to_string(), + json_path: "/kv_cache_quant_algo".to_string(), + }), + } +} + +fn publisher_default_request(mesh_config: &MeshConfig) -> SkippyConfigResolveRequest<'_> { + SkippyConfigResolveRequest { + mesh_config, + model_id: "Qwen/Qwen3-0.6B:Q4_K_M", + model_path: Path::new("/models/qwen.gguf"), + model_bytes: 100 * 1024 * 1024 * 1024, + allocatable_memory_bytes: None, + request_defaults: None, + package_generation: None, + compact_meta: None, + } +} + +#[test] +fn publisher_kv_default_is_used_below_explicit_user_override() { + let defaults = publisher_q8_defaults(); + let automatic = resolve_skippy_config_for_selector_with_publisher_defaults( + publisher_default_request(&MeshConfig::default()), + None, + Some(&defaults), + ) + .unwrap(); + assert_eq!(automatic.model_fit.cache_type_k, "q8_0"); + assert_eq!(automatic.model_fit.cache_type_v, "q8_0"); + assert_eq!(automatic.model_fit.kv_cache_policy, "publisher"); + + let explicit_config = parse_config( + r#" +[defaults.model_fit] +cache_type_k = "f16" +cache_type_v = "f16" +"#, + ); + let explicit = resolve_skippy_config_for_selector_with_publisher_defaults( + publisher_default_request(&explicit_config), + None, + Some(&defaults), + ) + .unwrap(); + assert_eq!(explicit.model_fit.cache_type_k, "f16"); + assert_eq!(explicit.model_fit.cache_type_v, "f16"); +} + fn assert_request_override_keeps_load_time_config( without_request: &ResolvedSkippyConfig, with_request: &ResolvedSkippyConfig, @@ -1837,7 +1895,7 @@ fn oversized_chat_template_file_is_rejected_before_runtime_startup() { } #[test] -fn inkling_family_defaults_to_q4_kv() { +fn model_family_and_weight_size_do_not_quantize_live_kv() { let resolved = resolve_skippy_config(SkippyConfigResolveRequest { mesh_config: &MeshConfig::default(), model_id: "meshllm/inkling-UD-Q2_K_XL-layers", @@ -1850,16 +1908,12 @@ fn inkling_family_defaults_to_q4_kv() { }) .unwrap(); - assert_eq!(resolved.model_fit.cache_type_k, "q4_0"); - assert_eq!(resolved.model_fit.cache_type_v, "q4_0"); + assert_eq!(resolved.model_fit.cache_type_k, "f16"); + assert_eq!(resolved.model_fit.cache_type_v, "f16"); } -/// The family q4_0 default must be guarded against the model's own metadata: -/// an Inkling variant with per-head widths not divisible by the q4_0 block -/// size (32) cannot load quantised KV, so the resolver must degrade the -/// default to f16 rather than fail the context build. #[test] -fn inkling_family_kv_default_degrades_to_f16_for_incompatible_meta() { +fn safe_f16_default_remains_f16_for_incompatible_quantized_kv_meta() { let compact_meta = crate::models::gguf::GgufCompactMeta { architecture: "inkling".to_string(), context_length: 65_536, diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs index 4be29cd9df..f54ad22483 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs @@ -36,9 +36,8 @@ pub(crate) struct SkippyConfigResolveRequest<'a> { pub(crate) request_defaults: Option<&'a RequestDefaultsConfig>, pub(crate) package_generation: Option<&'a PackageGenerationInfo>, /// GGUF metadata for the model being resolved, when available. Used to - /// guard the size-tiered KV cache default against quantised-KV load - /// incompatibilities (Flash Attention / block alignment). `None` leaves the - /// default unguarded — the pre-existing behaviour. + /// guard publisher-declared quantised K/V against native load constraints + /// such as Flash Attention and block alignment. pub(crate) compact_meta: Option<&'a crate::models::gguf::GgufCompactMeta>, } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/split_certification.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/split_certification.rs index 1b7eef5123..079c95b69c 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/split_certification.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/split_certification.rs @@ -195,6 +195,7 @@ pub(crate) fn split_certification_label( activation_width: 0, tensor_count: 0, generation: None, + publisher_defaults: None, }; Some( if certified_family(&package).ok().flatten().is_some() { @@ -223,6 +224,7 @@ mod tests { activation_width: 1, tensor_count: 1, generation: None, + publisher_defaults: None, } } diff --git a/crates/mesh-llm-host-runtime/src/runtime/local.rs b/crates/mesh-llm-host-runtime/src/runtime/local.rs index 1288d57667..38487a507a 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local.rs @@ -692,12 +692,12 @@ pub(super) async fn start_local_openai_model( .flatten() }; - // Guard the size-tiered default against quantised-KV load incompatibilities - // (Flash Attention off, or a head_dim not divisible by the block size) so - // planning and the load agree and the context build does not fail. Explicit - // user overrides below are never guarded — they must fail loudly. - let kv_cache = skippy::KvCachePolicy::for_model_size(total_model_bytes) - .guarded_for_model(compact_meta.as_ref()); + let kv_cache = skippy::KvCachePolicy::from_publisher_defaults( + package + .as_ref() + .and_then(|package| package.publisher_defaults.as_ref()), + ) + .guarded_for_model(compact_meta.as_ref()); let effective_cache_type_k = spec .cache_type_k_override .unwrap_or(kv_cache.cache_type_k()); @@ -708,7 +708,7 @@ pub(super) async fn start_local_openai_model( effective_cache_type_k, effective_cache_type_v, ) - .unwrap_or(models::gguf::GgufKvCacheQuant::Q8_0); + .unwrap_or(models::gguf::GgufKvCacheQuant::F16); let measurement_key = MemoryPlanMeasurementKey::new(format!( "model={runtime_model_name:?};path={:?};bytes={local_model_bytes};capacity={my_vram};config={:?};config_model={:?};device={:?};pinned_gpu={:?};cache_k={effective_cache_type_k:?};cache_v={effective_cache_type_v:?};batch={:?};ubatch={:?};flash={:?}", spec.model_path, diff --git a/crates/mesh-llm-host-runtime/src/runtime/local_package.rs b/crates/mesh-llm-host-runtime/src/runtime/local_package.rs index 9886d41eb1..9449899758 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local_package.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local_package.rs @@ -88,42 +88,26 @@ pub(super) fn split_runtime_kv_bytes_per_token( /// Resolve the K/V cache types that split stages will actually load with. /// -/// Stage loading applies the family default (for example Inkling's Q4_0 K/V) -/// ahead of the size-tiered `KvCachePolicy`. Planning must resolve K/V the same -/// way, or it budgets for a cheaper cache than the stages allocate and -/// over-packs the topology into an out-of-memory load. +/// Planning uses the same package-backed default as stage loading so it +/// budgets the allocation that will actually be created. pub(super) fn split_effective_kv_cache_quant( package: &skippy::SkippyPackageIdentity, compact_meta: &models::gguf::GgufCompactMeta, cache_type_k_override: Option<&str>, cache_type_v_override: Option<&str>, ) -> models::gguf::GgufKvCacheQuant { - // Guard the size-tiered default against the model's quantised-KV - // compatibility (Flash Attention / block alignment) so planning budgets for - // the same cache the stage load can actually allocate. The family default - // gets the same metadata guard: a family that defaults to quantised K/V - // (Inkling -> q4_0) must fall back to f16 when the actual GGUF metadata - // cannot load it, or planning and load both select an unloadable cache. - // Explicit overrides below are never guarded — an override that cannot - // load must fail loudly. - let size_policy = skippy::KvCachePolicy::for_model_size(package.source_model_bytes) - .guarded_for_model(Some(compact_meta)); - let family_default = skippy::family_policy_for_compact_meta(compact_meta) - .default_kv_cache_type - .and_then(|default| models::gguf::GgufKvCacheQuant::from_llama_args(default, default)) - .map(|quant| compact_meta.compatible_default_kv_cache_quant(quant)) - .map(|quant| quant.k.as_llama_arg()); - - // Explicit user overrides win, then the family default, then model size. - let effective_k = cache_type_k_override - .or(family_default) - .unwrap_or(size_policy.cache_type_k()); - let effective_v = cache_type_v_override - .or(family_default) - .unwrap_or(size_policy.cache_type_v()); + let package_policy = + skippy::KvCachePolicy::from_publisher_defaults(package.publisher_defaults.as_ref()) + .guarded_for_model(Some(compact_meta)); + let effective_k = cache_type_k_override.unwrap_or(package_policy.cache_type_k()); + let effective_v = cache_type_v_override.unwrap_or(package_policy.cache_type_v()); models::gguf::GgufKvCacheQuant::from_llama_args(effective_k, effective_v).unwrap_or_else(|| { - split_kv_cache_quant(&size_policy, cache_type_k_override, cache_type_v_override) + split_kv_cache_quant( + &package_policy, + cache_type_k_override, + cache_type_v_override, + ) }) } pub(super) async fn resolve_split_runtime_package( @@ -175,7 +159,7 @@ pub(super) fn split_kv_cache_quant( split_kv_policy.cache_type_k(), split_kv_policy.cache_type_v(), ) - .unwrap_or(models::gguf::GgufKvCacheQuant::Q8_0); + .unwrap_or(models::gguf::GgufKvCacheQuant::F16); match (cache_type_k_override, cache_type_v_override) { (None, None) => policy_quant, diff --git a/crates/mesh-llm-host-runtime/src/runtime/local_split/loading.rs b/crates/mesh-llm-host-runtime/src/runtime/local_split/loading.rs index f27f090b32..7126a316ec 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local_split/loading.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local_split/loading.rs @@ -823,7 +823,7 @@ pub(super) async fn split_generation_load_settings<'a>( .first() .context("split topology did not produce stage 0")?; let load_mode = split_generation_load_mode(spec.package); - let mut resolved = skippy::resolve_skippy_config_for_selector( + let mut resolved = skippy::resolve_skippy_config_for_selector_with_publisher_defaults( skippy::SkippyConfigResolveRequest { mesh_config: spec.mesh_config, model_id: spec.model_ref, @@ -833,11 +833,12 @@ pub(super) async fn split_generation_load_settings<'a>( request_defaults: None, package_generation: spec.package.generation.as_ref(), // Split stage load uses the compact metadata scanned during planning - // so the resolver guards both the size-tiered default and the family - // K/V default exactly like the split planner does. + // so publisher-declared quantised K/V gets the same native + // compatibility guard as the split planner. compact_meta: Some(spec.compact_meta), }, spec.config_model_id, + spec.package.publisher_defaults.as_ref(), )?; resolved.materialize_projector_url().await?; resolved.model_fit.ctx_size = spec.ctx_size; diff --git a/crates/mesh-llm-host-runtime/src/runtime/local_split/test_support.rs b/crates/mesh-llm-host-runtime/src/runtime/local_split/test_support.rs index 2f39ce8371..8a469b52e4 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local_split/test_support.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local_split/test_support.rs @@ -32,6 +32,7 @@ pub(super) fn package(layer_count: u32) -> skippy::SkippyPackageIdentity { activation_width: 2048, tensor_count: 100, generation: None, + publisher_defaults: None, } } @@ -564,19 +565,10 @@ stop = ["END"] } /// Split stage loading must resolve with the compact metadata scanned during -/// planning: the architecture-driven K/V default gets the same compatibility -/// guard as the planner, so a default the actual GGUF cannot load (here: -/// Inkling → q4_0 with per-head widths not divisible by the q4_0 block size) -/// degrades to f16 at stage load instead of failing the context build. -/// -/// The package is deliberately small (10 GB) so the size-tiered policy alone -/// would pick q8_0: the observed q4_0-vs-f16 swing can only come from actual -/// Inkling metadata, pinning the plumbing rather than a model-name heuristic. -/// Split load specifications require this metadata, so both the -/// initial-load and coordinator-replan constructors must carry it; dropping -/// the final resolver handoff would regress this test to q4_0. +/// planning. A publisher-declared quantised K/V type that the GGUF cannot load +/// must degrade to f16 at stage load instead of failing the context build. #[tokio::test] -async fn split_stage_load_guards_metadata_kv_default_with_planned_metadata() { +async fn split_stage_load_guards_publisher_kv_default_with_planned_metadata() { let node = mesh::Node::new_for_tests(NodeRole::Host { http_port: 9338 }) .await .unwrap(); @@ -587,6 +579,14 @@ async fn split_stage_load_guards_metadata_kv_default_with_planned_metadata() { let mut identity = package(66); identity.package_ref = "hf://Mesh-LLM/test-inkling-package".to_string(); identity.source_model_bytes = 10 * 1024 * 1024 * 1024; + identity.publisher_defaults = Some(skippy_package_format::PublisherModelDefaults { + compute_dtype: None, + kv_cache_dtype: Some(skippy_package_format::PublisherDtypeDeclaration { + dtype: skippy_package_format::PublisherDtype::Q4_0, + artifact_id: "publisher-config".to_string(), + json_path: "/kv_cache_dtype".to_string(), + }), + }); let local_id = node.id(); let generation = SplitTopologyGeneration::new( "guard-topology".into(), @@ -600,7 +600,7 @@ async fn split_stage_load_guards_metadata_kv_default_with_planned_metadata() { ); // Per-head widths of 100 are not a multiple of the q4_0 block size (32), - // so the Inkling architecture's quantised default cannot load. + // so the publisher declaration cannot load. let incompatible_meta = crate::models::gguf::GgufCompactMeta { architecture: "inkling".to_string(), context_length: 65_536, @@ -613,8 +613,8 @@ async fn split_stage_load_guards_metadata_kv_default_with_planned_metadata() { ..Default::default() }; - // With the planned metadata, the unloadable architecture default degrades to - // f16 — the same cache the split planner budgets for. + // With the planned metadata, the unloadable publisher default degrades to + // f16, matching the split planner's budget. let guarded_spec = SplitGenerationLoadSpec { node: &node, mesh_config: &mesh_config, @@ -649,12 +649,11 @@ async fn split_stage_load_guards_metadata_kv_default_with_planned_metadata() { .expect("guarded split settings should resolve"); assert_eq!( guarded.runtime_options.config.cache_type_k, "f16", - "incompatible architecture default must degrade to f16 at stage load" + "incompatible publisher default must degrade to f16 at stage load" ); assert_eq!(guarded.runtime_options.config.cache_type_v, "f16"); - // Without metadata the model name carries no architecture authority, so - // the generic q8_0 size tier remains in effect. + // Without publisher metadata, model name and weight size do not quantize KV. let unguarded = skippy::resolve_skippy_config_for_selector( skippy::SkippyConfigResolveRequest { mesh_config: &mesh_config, @@ -670,10 +669,10 @@ async fn split_stage_load_guards_metadata_kv_default_with_planned_metadata() { ) .expect("unguarded resolver settings should resolve"); assert_eq!( - unguarded.model_fit.cache_type_k, "q8_0", - "no-metadata stage load keeps the generic size-tier default" + unguarded.model_fit.cache_type_k, "f16", + "no-metadata stage load uses the safe f16 default" ); - assert_eq!(unguarded.model_fit.cache_type_v, "q8_0"); + assert_eq!(unguarded.model_fit.cache_type_v, "f16"); } #[tokio::test] diff --git a/crates/mesh-llm-host-runtime/src/runtime/local_split/tests.rs b/crates/mesh-llm-host-runtime/src/runtime/local_split/tests.rs index f275e0c6d0..93a67e5134 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local_split/tests.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local_split/tests.rs @@ -1955,7 +1955,7 @@ fn split_topology_minimum_rejects_single_stage_split_candidate() { } #[test] -fn split_planning_uses_family_kv_defaults_for_inkling() { +fn split_planning_uses_safe_f16_without_publisher_metadata() { let mut meta = crate::models::gguf::GgufCompactMeta { architecture: "inkling".to_string(), context_length: 65_536, @@ -1969,22 +1969,20 @@ fn split_planning_uses_family_kv_defaults_for_inkling() { }; meta.kv_head_counts = vec![8; 66]; - // Inkling's reviewed family default keeps both planning and stage loading - // on quantized Q4_0 K/V rather than silently expanding to F16. let mut identity = package(66); identity.source_model_bytes = 318 * 1024 * 1024 * 1024; let planned = split_runtime_kv_bytes_per_token(&identity, &meta, None, None).unwrap(); - let expected_q4 = crate::models::gguf::GgufKvCacheQuant::from_llama_args("q4_0", "q4_0") + let expected_f16 = crate::models::gguf::GgufKvCacheQuant::from_llama_args("f16", "f16") .unwrap() .kv_cache_bytes_per_token(&meta) .unwrap(); - assert_eq!(planned, expected_q4); + assert_eq!(planned, expected_f16); - // Explicit user overrides still win over the family default. + // Explicit user overrides still win over the safe package fallback. let overridden = - split_runtime_kv_bytes_per_token(&identity, &meta, Some("f16"), Some("f16")).unwrap(); - assert!(overridden > planned); + split_runtime_kv_bytes_per_token(&identity, &meta, Some("q8_0"), Some("q8_0")).unwrap(); + assert!(overridden < planned); } #[test] @@ -2013,12 +2011,10 @@ fn split_planning_allows_zero_kv_only_for_proven_pure_recurrent_metadata() { assert!(split_runtime_kv_bytes_per_token(&identity, &dense_missing_heads, None, None).is_err()); } -/// The family default must get the same metadata guard as the size-tiered -/// policy: an Inkling variant whose per-head widths are not q4_0-block-aligned -/// cannot load quantised K/V, so planning must budget f16 bytes instead of -/// selecting an unloadable family default. +/// A package's quantised K/V declaration must be guarded by the actual GGUF +/// layout so planning never budgets an unloadable cache type. #[test] -fn split_planning_guards_family_kv_default_against_incompatible_meta() { +fn split_planning_guards_publisher_kv_default_against_incompatible_meta() { let mut meta = crate::models::gguf::GgufCompactMeta { architecture: "inkling".to_string(), context_length: 65_536, @@ -2036,6 +2032,14 @@ fn split_planning_guards_family_kv_default_against_incompatible_meta() { let mut identity = package(66); identity.source_model_bytes = 318 * 1024 * 1024 * 1024; + identity.publisher_defaults = Some(skippy_package_format::PublisherModelDefaults { + compute_dtype: None, + kv_cache_dtype: Some(skippy_package_format::PublisherDtypeDeclaration { + dtype: skippy_package_format::PublisherDtype::Q4_0, + artifact_id: "publisher-config".to_string(), + json_path: "/kv_cache_dtype".to_string(), + }), + }); let planned = split_runtime_kv_bytes_per_token(&identity, &meta, None, None).unwrap(); let expected_f16 = crate::models::gguf::GgufKvCacheQuant::from_llama_args("f16", "f16") @@ -2044,7 +2048,7 @@ fn split_planning_guards_family_kv_default_against_incompatible_meta() { .unwrap(); assert_eq!( planned, expected_f16, - "incompatible family default must degrade to f16 in split planning" + "incompatible publisher default must degrade to f16 in split planning" ); // An explicit override is never guarded — it still selects q4_0 even @@ -2063,7 +2067,7 @@ fn split_planning_guards_family_kv_default_against_incompatible_meta() { /// Set `INKLING_METADATA_GGUF` to a package's `shared/metadata.gguf` to run it; /// skipped otherwise so CI stays hermetic. #[test] -fn real_inkling_metadata_plans_family_kv_not_size_tiered() { +fn real_inkling_metadata_uses_safe_f16_without_publisher_defaults() { let Ok(path) = std::env::var("INKLING_METADATA_GGUF") else { eprintln!("skip: INKLING_METADATA_GGUF not set"); return; @@ -2081,43 +2085,16 @@ fn real_inkling_metadata_plans_family_kv_not_size_tiered() { meta.context_length ); - let policy = crate::inference::skippy::family_policy_for_compact_meta(&meta); - eprintln!( - "FAMILY default_kv_cache_type={:?}", - policy.default_kv_cache_type - ); - let mut identity = package(meta.layer_count); identity.source_model_bytes = 318 * 1024 * 1024 * 1024; let planned = split_runtime_kv_bytes_per_token(&identity, &meta, None, None).unwrap(); - let expected_q4 = crate::models::gguf::GgufKvCacheQuant::from_llama_args("q4_0", "q4_0") + let expected_f16 = crate::models::gguf::GgufKvCacheQuant::from_llama_args("f16", "f16") .unwrap() .kv_cache_bytes_per_token(&meta) .unwrap(); - let size_tiered = { - let p = - crate::inference::skippy::KvCachePolicy::for_model_size(identity.source_model_bytes); - split_kv_cache_quant(&p, None, None) - .kv_cache_bytes_per_token(&meta) - .unwrap() - }; - let ctx = u64::from(meta.context_length.max(1)); - eprintln!( - "KV/token planned={planned} size_tiered={size_tiered} ratio={:.2}x | @ctx{ctx}: planned={:.1}GiB size_tiered={:.1}GiB under_budget={:.1}GiB", - planned as f64 / size_tiered.max(1) as f64, - (planned * ctx) as f64 / (1024.0 * 1024.0 * 1024.0), - (size_tiered * ctx) as f64 / (1024.0 * 1024.0 * 1024.0), - ((planned - size_tiered.min(planned)) * ctx) as f64 / (1024.0 * 1024.0 * 1024.0), - ); - assert_eq!( - policy.default_kv_cache_type, - Some("q4_0"), - "inkling must resolve a q4_0 family K/V default" - ); - assert_eq!( - planned, expected_q4, - "family-aware planning must use the Inkling Q4_0 K/V default" + planned, expected_f16, + "model architecture and weight size must not quantize live K/V" ); } diff --git a/crates/mesh-llm-host-runtime/src/runtime/split_planning.rs b/crates/mesh-llm-host-runtime/src/runtime/split_planning.rs index d3fe54a7df..3ee98a5f76 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/split_planning.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/split_planning.rs @@ -704,6 +704,7 @@ mod tests { activation_width: 896, tensor_count: 100, generation: None, + publisher_defaults: None, } } diff --git a/crates/mesh-llm-host-runtime/src/runtime/stage_admission/tests.rs b/crates/mesh-llm-host-runtime/src/runtime/stage_admission/tests.rs index 450b690d6d..74b6f1adac 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/stage_admission/tests.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/stage_admission/tests.rs @@ -47,6 +47,8 @@ fn manifest() -> PackageManifest { ], }, sidecars: Vec::new(), + publisher_metadata: Vec::new(), + publisher_defaults: None, generation: None, native_abi_version: "0.1.52".to_string(), generator_version: "test".to_string(), diff --git a/crates/model-package/src/script.rs b/crates/model-package/src/script.rs index 91010c52a1..d5498f1805 100644 --- a/crates/model-package/src/script.rs +++ b/crates/model-package/src/script.rs @@ -214,6 +214,15 @@ mod tests { .contains(r#"WRITE_PACKAGE_PROJECTOR_ARGS+=(--projector "$PROJECTOR_PATH")"#) ); assert!(EMBEDDED_SCRIPT.contains(r#""${WRITE_PACKAGE_PROJECTOR_ARGS[@]}""#)); + assert!(EMBEDDED_SCRIPT.contains("Pinned source revision")); + assert!(EMBEDDED_SCRIPT.contains("config.json")); + assert!(EMBEDDED_SCRIPT.contains("hf_quant_config.json")); + assert!( + EMBEDDED_SCRIPT.contains( + r#"WRITE_PACKAGE_METADATA_ARGS+=(--publisher-metadata "$METADATA_PATH")"# + ) + ); + assert!(EMBEDDED_SCRIPT.contains(r#""${WRITE_PACKAGE_METADATA_ARGS[@]}""#)); assert!(EMBEDDED_SCRIPT.contains(r#"time "$SLICER" write-package "$WRITE_PACKAGE_INPUT""#)); assert!(!EMBEDDED_SCRIPT.contains(r#"time $SLICER write-package "$SOURCE_PATH""#)); } diff --git a/crates/model-package/src/scripts/split-model-job.sh b/crates/model-package/src/scripts/split-model-job.sh index 63ac241961..3f008eedd1 100755 --- a/crates/model-package/src/scripts/split-model-job.sh +++ b/crates/model-package/src/scripts/split-model-job.sh @@ -228,6 +228,19 @@ df -h / || true echo " Preparing Hugging Face uploader..." python3 -m venv "$VENV_DIR" > /dev/null "$VENV_DIR/bin/pip" install -q huggingface_hub +SOURCE_REVISION="$("$VENV_DIR/bin/python3" <<'PYTHON' +from huggingface_hub import HfApi +import os + +api = HfApi(token=os.environ.get("HF_TOKEN")) +info = api.model_info(os.environ["SOURCE_REPO"], revision=os.environ.get("SOURCE_REVISION", "main")) +if not info.sha: + raise RuntimeError("Hugging Face did not return an immutable source revision") +print(info.sha) +PYTHON +)" +export SOURCE_REVISION +echo " Pinned source revision: $SOURCE_REVISION" "$VENV_DIR/bin/python3" << 'PYTHON' from huggingface_hub import HfApi import os @@ -340,6 +353,43 @@ PYTHON echo " Projector: $PROJECTOR_PATH" WRITE_PACKAGE_PROJECTOR_ARGS+=(--projector "$PROJECTOR_PATH") done <<< "${SOURCE_PROJECTOR_FILES:-}" +PUBLISHER_METADATA_DIR="${LOCAL_WORK_DIR}/publisher-metadata" +mkdir -p "$PUBLISHER_METADATA_DIR" +export PUBLISHER_METADATA_DIR +mapfile -t PUBLISHER_METADATA_PATHS < <("$VENV_DIR/bin/python3" <<'PYTHON' +from huggingface_hub import hf_hub_download +from huggingface_hub.errors import EntryNotFoundError +from pathlib import Path +import os + +destination = Path(os.environ["PUBLISHER_METADATA_DIR"]) +for name in ( + "config.json", + "generation_config.json", + "tokenizer_config.json", + "chat_template.jinja", + "hf_quant_config.json", +): + try: + cached = Path(hf_hub_download( + repo_id=os.environ["SOURCE_REPO"], + filename=name, + revision=os.environ["SOURCE_REVISION"], + cache_dir=os.environ["HF_HUB_CACHE"], + token=os.environ.get("HF_TOKEN"), + )) + except EntryNotFoundError: + continue + output = destination / name + output.write_bytes(cached.read_bytes()) + print(output) +PYTHON +) +WRITE_PACKAGE_METADATA_ARGS=() +for METADATA_PATH in "${PUBLISHER_METADATA_PATHS[@]}"; do + echo " Publisher metadata: $METADATA_PATH" + WRITE_PACKAGE_METADATA_ARGS+=(--publisher-metadata "$METADATA_PATH") +done echo " Hugging Face cache: $HF_HUB_CACHE" echo " Package workspace: $PACKAGE_DIR" echo " Temporary workspace: $TMPDIR" @@ -363,6 +413,7 @@ time "$SLICER" write-package "$WRITE_PACKAGE_INPUT" \ --out-dir "$PACKAGE_DIR" \ --after-artifact-command "$ARTIFACT_UPLOAD_HOOK" \ "${WRITE_PACKAGE_PROJECTOR_ARGS[@]}" \ + "${WRITE_PACKAGE_METADATA_ARGS[@]}" \ "${WRITE_PACKAGE_IDENTITY_ARGS[@]}" WRITE_PACKAGE_STATUS=$? set -e diff --git a/crates/skippy-model-package/README.md b/crates/skippy-model-package/README.md index 053432537d..1db57d2a1c 100644 --- a/crates/skippy-model-package/README.md +++ b/crates/skippy-model-package/README.md @@ -89,7 +89,11 @@ separate native inspection change. Equal bytes in distinct source allocations are **not** aliases. Pass `--projector path/to/mmproj*.gguf` to copy and verify explicit projector -sidecars. This writer does not infer generation policy/defaults from tensor names +sidecars. Pass `--publisher-metadata path/to/config.json` (repeatable) to copy +supported Hugging Face configuration files into `metadata/`, bind each file to +the source repository and immutable revision, and derive typed compute/KV +defaults. `config.json` geometry that conflicts with the GGUF fails before +payload emission. This writer does not infer generation policy/defaults from tensor names or implement offline conversion. The existing `plan`, `write`, `write-stages`, `validate`, `validate-package`, `preflight`, and GLM-DSA commands still serve their existing slice/v1 contracts; they are not v2 certification or serving paths. @@ -124,6 +128,9 @@ substituted tensors, inconsistent source identities, duplicate artifacts/sidecar unproven alias claims, v1 manifests and corrupt/truncated files fail with nonzero exit status. Success prints JSON with the package ID, `source_completeness_verified` and checked source/artifact/tensor/projector counts; it does not modify the package. +Publisher metadata artifacts receive the same size and SHA-256 verification; +their repository/revision labels remain caller-supplied provenance rather than +an independently authenticated Hub claim. This unit verifies the writer's **byte-preserving whole-shard representation**. Repacked/transformed containers, tensor-only digests, non-projector sidecars and @@ -141,8 +148,10 @@ explicit provenance: skippy-model-package write-package ./model.gguf \ --out-dir model-package/ \ --model-id org/repo:Q4_K_M \ - --source-revision abc123 \ - --source-file Qwen3-8B-Q4_K_M.gguf + --source-repo org/repo \ + --source-revision aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa \ + --source-file Qwen3-8B-Q4_K_M.gguf \ + --publisher-metadata ./config.json ``` This keeps canonical package identity tied to real model coordinates rather diff --git a/crates/skippy-model-package/src/cli.rs b/crates/skippy-model-package/src/cli.rs index 6e5db5db5a..fc63efb2af 100644 --- a/crates/skippy-model-package/src/cli.rs +++ b/crates/skippy-model-package/src/cli.rs @@ -52,6 +52,11 @@ pub(crate) enum Command { out_dir: PathBuf, #[arg(long = "projector")] projectors: Vec, + /// Immutable publisher configuration files to copy into metadata/. + /// Supported basenames: config.json, generation_config.json, + /// tokenizer_config.json, chat_template.jinja, hf_quant_config.json. + #[arg(long = "publisher-metadata")] + publisher_metadata: Vec, #[arg(long)] after_artifact_command: Option, #[arg(long)] diff --git a/crates/skippy-model-package/src/main.rs b/crates/skippy-model-package/src/main.rs index 726543edfe..0991ec1f19 100644 --- a/crates/skippy-model-package/src/main.rs +++ b/crates/skippy-model-package/src/main.rs @@ -97,6 +97,7 @@ fn run(args: Args) -> Result<()> { model, out_dir, projectors, + publisher_metadata, after_artifact_command, transform_artifact_command, model_id, @@ -107,7 +108,10 @@ fn run(args: Args) -> Result<()> { } => package_v2::write_package( model, out_dir, - projectors, + package_v2::PackageSidecars { + projectors, + publisher_metadata, + }, ArtifactHook { command: after_artifact_command, }, diff --git a/crates/skippy-model-package/src/package_v2.rs b/crates/skippy-model-package/src/package_v2.rs index ee49cfdf0c..0d5df44485 100644 --- a/crates/skippy-model-package/src/package_v2.rs +++ b/crates/skippy-model-package/src/package_v2.rs @@ -8,8 +8,9 @@ use std::time::{SystemTime, UNIX_EPOCH}; use anyhow::{Context, Result, ensure}; use skippy_model::package_carrier::resolve_package_carrier; use skippy_package_format::{ - Artifact, ArtifactCatalog, PACKAGE_SCHEMA_VERSION, PackageManifest, Sidecar, SidecarKind, - SourceModel, Tensor, TensorCatalog, + Artifact, ArtifactCatalog, PACKAGE_SCHEMA_VERSION, PackageManifest, PublisherDtype, + PublisherDtypeDeclaration, PublisherMetadata, PublisherMetadataRole, PublisherModelDefaults, + Sidecar, SidecarKind, SourceModel, Tensor, TensorCatalog, }; use skippy_runtime::{ModelInfo, TensorInfo, write_gguf_metadata_from_parts}; @@ -27,15 +28,25 @@ mod layout; use layout::{PlannedArtifact, PlannedArtifactKind, plan_artifacts}; +#[derive(Default)] +pub(crate) struct PackageSidecars { + pub(crate) projectors: Vec, + pub(crate) publisher_metadata: Vec, +} + pub(crate) fn write_package( model: String, out_dir: PathBuf, - projectors: Vec, + sidecars: PackageSidecars, artifact_hook: ArtifactHook, artifact_transform: ArtifactHook, explicit: ExplicitSourceIdentity, resume_existing_artifacts: bool, ) -> Result<()> { + let PackageSidecars { + projectors, + publisher_metadata, + } = sidecars; ensure!( artifact_transform.command.is_none(), "v2 creation preserves source bytes; transform the independent source before packaging, not package artifacts" @@ -46,12 +57,25 @@ pub(crate) fn write_package( ensure_native_inventory_matches(&inventory, &source)?; let planned = plan_artifacts(&source.tensors)?; let mut manifest = manifest_from_source(&input, &inventory)?; + let mut publisher_defaults_manifest = manifest.clone(); + let mut publisher_artifact_ids = BTreeSet::new(); + for metadata_path in &publisher_metadata { + let metadata = publisher_metadata_descriptor(metadata_path, &manifest)?; + ensure!( + publisher_artifact_ids.insert(metadata.artifact_id.clone()), + "publisher metadata filename {:?} appears more than once", + metadata.source_path + ); + apply_publisher_defaults(&mut publisher_defaults_manifest, metadata_path, &metadata)?; + } + let publisher_defaults = publisher_defaults_manifest.publisher_defaults; fs::create_dir_all(&out_dir)?; ensure!( !out_dir.join("model-package.json").exists(), "output already contains model-package.json; use a new directory for v2 creation" ); - let mut progress = PackageProgress::new(planned.len() + projectors.len() + 2); + let mut progress = + PackageProgress::new(planned.len() + projectors.len() + publisher_metadata.len() + 2); let no_hook = ArtifactHook { command: None }; let source_tensors = source_tensors_by_name(&inventory)?; let common_names = planned @@ -153,6 +177,34 @@ pub(crate) fn write_package( }); manifest.artifact_catalog.entries.push(artifact); } + for metadata_path in &publisher_metadata { + let (artifact, metadata) = copy_publisher_metadata( + metadata_path, + &manifest, + &out_dir, + resume_existing_artifacts, + )?; + progress.start_step(&artifact.path)?; + run_artifact_hook( + &artifact_hook, + &out_dir.join(&artifact.path), + &artifact.path, + )?; + if artifact_hook.command.is_some() && out_dir.join(&artifact.path).exists() { + ensure!( + file_sha256(&out_dir.join(&artifact.path))? == artifact.sha256, + "publisher metadata changed after artifact hook" + ); + } + progress.finish_step(&format!( + "{} {}", + artifact.path, + format_bytes(artifact.byte_size) + ))?; + manifest.publisher_metadata.push(metadata); + manifest.artifact_catalog.entries.push(artifact); + } + manifest.publisher_defaults = publisher_defaults; manifest.package_id = manifest.computed_package_id()?; manifest.validate()?; progress.start_step("model-package.json")?; @@ -213,6 +265,8 @@ fn manifest_from_source( entries: Vec::new(), }, sidecars: Vec::new(), + publisher_metadata: Vec::new(), + publisher_defaults: None, generation: None, native_abi_version: format!( "{}.{}.{}", @@ -555,5 +609,239 @@ fn copy_projector(source: &Path, index: usize, out_dir: &Path, resume: bool) -> }) } +fn copy_publisher_metadata( + source: &Path, + manifest: &PackageManifest, + out_dir: &Path, + resume: bool, +) -> Result<(Artifact, PublisherMetadata)> { + let metadata = publisher_metadata_descriptor(source, manifest)?; + let source_name = metadata.source_path.as_str(); + let id = metadata.artifact_id.clone(); + let relative = format!("metadata/{source_name}"); + let destination = out_dir.join(&relative); + copy_artifact(source, &destination, resume)?; + let source_digest = file_sha256(source)?; + ensure!( + file_sha256(&destination)? == source_digest, + "written publisher metadata differs from source" + ); + let byte_size = fs::metadata(&destination)?.len(); + Ok(( + Artifact { + id, + path: relative, + byte_size, + sha256: source_digest, + }, + metadata, + )) +} + +fn publisher_metadata_descriptor( + source: &Path, + manifest: &PackageManifest, +) -> Result { + let source_name = source + .file_name() + .and_then(|name| name.to_str()) + .context("publisher metadata filename must be valid UTF-8")?; + let role = match source_name { + "config.json" => PublisherMetadataRole::ModelConfig, + "generation_config.json" => PublisherMetadataRole::GenerationConfig, + "tokenizer_config.json" => PublisherMetadataRole::TokenizerConfig, + "chat_template.jinja" => PublisherMetadataRole::ChatTemplate, + "hf_quant_config.json" => PublisherMetadataRole::HfQuantConfig, + other => anyhow::bail!("unsupported publisher metadata filename {other:?}"), + }; + let repo = manifest + .source_model + .repo + .clone() + .context("publisher metadata requires source repository provenance")?; + let revision = manifest + .source_model + .revision + .clone() + .context("publisher metadata requires an immutable source revision")?; + ensure!( + revision.len() == 40 && revision.bytes().all(|byte| byte.is_ascii_hexdigit()), + "publisher metadata requires a 40-character immutable Hugging Face commit revision" + ); + Ok(PublisherMetadata { + role, + artifact_id: format!("publisher-{}", source_name.replace(['.', '_'], "-")), + source_repo: repo, + source_revision: revision, + source_path: source_name.to_string(), + }) +} + +fn apply_publisher_defaults( + manifest: &mut PackageManifest, + source: &Path, + metadata: &PublisherMetadata, +) -> Result<()> { + if !matches!( + metadata.role, + PublisherMetadataRole::ModelConfig | PublisherMetadataRole::HfQuantConfig + ) { + return Ok(()); + } + let document: serde_json::Value = serde_json::from_slice(&fs::read(source)?) + .with_context(|| format!("parse publisher metadata {}", source.display()))?; + ensure!( + document.is_object(), + "publisher metadata {} must contain a JSON object", + source.display() + ); + if metadata.role == PublisherMetadataRole::ModelConfig { + validate_config_geometry(manifest, &document)?; + } + let defaults = manifest + .publisher_defaults + .get_or_insert(PublisherModelDefaults { + compute_dtype: None, + kv_cache_dtype: None, + }); + + if metadata.role == PublisherMetadataRole::ModelConfig + && let Some((json_path, value)) = first_string_at(&document, &["/torch_dtype", "/dtype"]) + && let Some(dtype) = parse_publisher_dtype(value) + { + merge_dtype_declaration( + &mut defaults.compute_dtype, + PublisherDtypeDeclaration { + dtype, + artifact_id: metadata.artifact_id.clone(), + json_path: json_path.to_string(), + }, + "compute dtype", + )?; + } + + let kv_paths = [ + "/kv_cache_quant_algo", + "/kv_cache_dtype", + "/kv_cache_scheme", + "/quantization_config/kv_cache_quant_algo", + "/quantization_config/kv_cache_dtype", + "/quantization_config/kv_cache_scheme", + "/compression_config/kv_cache_quant_algo", + "/compression_config/kv_cache_dtype", + "/compression_config/kv_cache_scheme", + ]; + if let Some((json_path, value)) = first_dtype_at(&document, &kv_paths) { + let dtype = parse_publisher_dtype(value).with_context(|| { + format!( + "unsupported publisher KV-cache dtype {value:?} at {json_path} in {}", + source.display() + ) + })?; + merge_dtype_declaration( + &mut defaults.kv_cache_dtype, + PublisherDtypeDeclaration { + dtype, + artifact_id: metadata.artifact_id.clone(), + json_path: json_path.to_string(), + }, + "KV-cache dtype", + )?; + } + Ok(()) +} + +fn first_string_at<'a>( + document: &'a serde_json::Value, + paths: &'a [&'a str], +) -> Option<(&'a str, &'a str)> { + paths.iter().find_map(|path| { + document + .pointer(path) + .and_then(serde_json::Value::as_str) + .map(|value| (*path, value)) + }) +} + +fn first_dtype_at<'a>( + document: &'a serde_json::Value, + paths: &'a [&'a str], +) -> Option<(&'a str, &'a str)> { + paths.iter().find_map(|path| { + let value = document.pointer(path)?; + value + .as_str() + .or_else(|| value.get("dtype").and_then(serde_json::Value::as_str)) + .or_else(|| value.get("type").and_then(serde_json::Value::as_str)) + .map(|dtype| (*path, dtype)) + }) +} + +fn parse_publisher_dtype(value: &str) -> Option { + match value.trim().to_ascii_lowercase().replace('-', "_").as_str() { + "f16" | "fp16" | "float16" => Some(PublisherDtype::F16), + "bf16" | "bfloat16" => Some(PublisherDtype::Bf16), + "f32" | "fp32" | "float32" => Some(PublisherDtype::F32), + "fp8" => Some(PublisherDtype::Fp8), + "fp8_e4m3" | "fp8_e4m3fn" => Some(PublisherDtype::Fp8E4m3), + "fp8_e5m2" => Some(PublisherDtype::Fp8E5m2), + "q8_0" | "int8" => Some(PublisherDtype::Q8_0), + "q4_0" | "int4" => Some(PublisherDtype::Q4_0), + _ => None, + } +} + +fn merge_dtype_declaration( + slot: &mut Option, + declaration: PublisherDtypeDeclaration, + label: &str, +) -> Result<()> { + if let Some(existing) = slot { + ensure!( + existing.dtype == declaration.dtype, + "conflicting publisher {label} declarations: {:?} at {} and {:?} at {}", + existing.dtype, + existing.json_path, + declaration.dtype, + declaration.json_path + ); + return Ok(()); + } + *slot = Some(declaration); + Ok(()) +} + +fn validate_config_geometry(manifest: &PackageManifest, config: &serde_json::Value) -> Result<()> { + let architecture = manifest + .model_metadata + .get("general.architecture") + .and_then(serde_json::Value::as_str) + .context("GGUF model metadata is missing general.architecture")?; + for (config_key, gguf_suffix) in [ + ("num_hidden_layers", "block_count"), + ("hidden_size", "embedding_length"), + ("num_attention_heads", "attention.head_count"), + ("num_key_value_heads", "attention.head_count_kv"), + ("max_position_embeddings", "context_length"), + ] { + let Some(config_value) = config.get(config_key).and_then(serde_json::Value::as_u64) else { + continue; + }; + let gguf_key = format!("{architecture}.{gguf_suffix}"); + let Some(gguf_value) = manifest + .model_metadata + .get(&gguf_key) + .and_then(serde_json::Value::as_u64) + else { + continue; + }; + ensure!( + config_value == gguf_value, + "publisher config {config_key}={config_value} conflicts with GGUF {gguf_key}={gguf_value}" + ); + } + Ok(()) +} + #[cfg(test)] mod tests; diff --git a/crates/skippy-model-package/src/package_v2/tests.rs b/crates/skippy-model-package/src/package_v2/tests.rs index aba19f4f0c..04428f16c4 100644 --- a/crates/skippy-model-package/src/package_v2/tests.rs +++ b/crates/skippy-model-package/src/package_v2/tests.rs @@ -32,7 +32,7 @@ fn write(source: &Path, out: &Path, resume: bool) -> Result<()> { write_package( source.display().to_string(), out.to_path_buf(), - Vec::new(), + PackageSidecars::default(), ArtifactHook { command: None }, ArtifactHook { command: None }, explicit(source), @@ -367,7 +367,7 @@ fn refuses_transform_hooks_and_existing_completion_marker() { let result = write_package( source.display().to_string(), out.clone(), - Vec::new(), + PackageSidecars::default(), ArtifactHook { command: None }, ArtifactHook { command: Some("must-not-run".into()), @@ -400,7 +400,10 @@ fn verified_resume_and_projector_sidecar_round_trip() { write_package( source.display().to_string(), out.clone(), - vec![projector], + PackageSidecars { + projectors: vec![projector], + publisher_metadata: Vec::new(), + }, ArtifactHook { command: None }, ArtifactHook { command: None }, explicit(&source), @@ -422,6 +425,101 @@ fn verified_resume_and_projector_sidecar_round_trip() { assert_eq!(manifest.artifact_catalog.entries.len(), 3); } +#[test] +fn publisher_metadata_is_copied_hashed_and_typed() { + let temp = tempfile::tempdir().unwrap(); + let source = temp.path().join("source.gguf"); + fixture(&source, &[tensor("first", 0)], None); + let config = temp.path().join("config.json"); + fs::write( + &config, + r#"{"torch_dtype":"bfloat16","num_hidden_layers":2}"#, + ) + .unwrap(); + let quant_config = temp.path().join("hf_quant_config.json"); + fs::write(&quant_config, r#"{"kv_cache_quant_algo":"FP8"}"#).unwrap(); + let out = temp.path().join("package"); + let mut source_identity = explicit(&source); + source_identity.source_repo = Some("fixture/model".to_string()); + source_identity.source_revision = Some("a".repeat(40)); + + write_package( + source.display().to_string(), + out.clone(), + PackageSidecars { + projectors: Vec::new(), + publisher_metadata: vec![config.clone(), quant_config], + }, + ArtifactHook { command: None }, + ArtifactHook { command: None }, + source_identity, + false, + ) + .unwrap(); + + let manifest = read_manifest(&out); + assert_eq!(manifest.publisher_metadata.len(), 2); + let metadata = &manifest.publisher_metadata[0]; + assert_eq!(metadata.source_repo, "fixture/model"); + assert_eq!(metadata.source_revision, "a".repeat(40)); + assert_eq!(metadata.source_path, "config.json"); + let artifact = manifest + .artifact_catalog + .entries + .iter() + .find(|artifact| artifact.id == metadata.artifact_id) + .unwrap(); + assert_eq!(artifact.path, "metadata/config.json"); + assert_eq!( + file_sha256(&out.join(&artifact.path)).unwrap(), + artifact.sha256 + ); + let defaults = manifest.publisher_defaults.unwrap(); + let declaration = defaults.compute_dtype.unwrap(); + assert_eq!( + declaration.dtype, + skippy_package_format::PublisherDtype::Bf16 + ); + assert_eq!(declaration.artifact_id, metadata.artifact_id); + assert_eq!( + defaults.kv_cache_dtype.unwrap().dtype, + skippy_package_format::PublisherDtype::Fp8 + ); + crate::verify_v2::verify_package(&out, &source, None, &[]).unwrap(); +} + +#[test] +fn publisher_config_conflicting_with_gguf_geometry_is_rejected() { + let temp = tempfile::tempdir().unwrap(); + let source = temp.path().join("source.gguf"); + fixture(&source, &[tensor("first", 0)], None); + let config = temp.path().join("config.json"); + fs::write(&config, r#"{"num_hidden_layers":99}"#).unwrap(); + let out = temp.path().join("package"); + let mut source_identity = explicit(&source); + source_identity.source_repo = Some("fixture/model".to_string()); + source_identity.source_revision = Some("a".repeat(40)); + + let error = write_package( + source.display().to_string(), + out, + PackageSidecars { + projectors: Vec::new(), + publisher_metadata: vec![config], + }, + ArtifactHook { command: None }, + ArtifactHook { command: None }, + source_identity, + false, + ) + .unwrap_err(); + assert!( + error + .to_string() + .contains("conflicts with GGUF llama.block_count=2") + ); +} + #[cfg(unix)] #[test] fn upload_hook_can_delete_verified_copies_without_losing_inventory() { @@ -440,7 +538,7 @@ fn upload_hook_can_delete_verified_copies_without_losing_inventory() { write_package( source.display().to_string(), out.clone(), - Vec::new(), + PackageSidecars::default(), ArtifactHook { command: Some(hook), }, diff --git a/crates/skippy-model-package/src/preflight.rs b/crates/skippy-model-package/src/preflight.rs index 0691a566f5..deae598fd5 100644 --- a/crates/skippy-model-package/src/preflight.rs +++ b/crates/skippy-model-package/src/preflight.rs @@ -656,7 +656,7 @@ fn partition_layers(layer_count: u32, stages: u32) -> Vec<(u32, u32)> { mod tests { use super::*; use crate::package::ArtifactHook; - use crate::package_v2::write_package; + use crate::package_v2::{PackageSidecars, write_package}; use crate::test_gguf::{explicit, fixture, tensor}; fn write_v2_fixture(root: &Path) -> std::path::PathBuf { @@ -673,7 +673,7 @@ mod tests { write_package( source.display().to_string(), package.clone(), - Vec::new(), + PackageSidecars::default(), ArtifactHook { command: None }, ArtifactHook { command: None }, explicit(&source), diff --git a/crates/skippy-model-package/src/verify_v2.rs b/crates/skippy-model-package/src/verify_v2.rs index 228bebac6a..66f441f073 100644 --- a/crates/skippy-model-package/src/verify_v2.rs +++ b/crates/skippy-model-package/src/verify_v2.rs @@ -82,12 +82,20 @@ pub(crate) fn verify_package( .iter() .map(|sidecar| sidecar.artifact_id.as_str()) .collect::>(); + let publisher_metadata_ids = manifest + .publisher_metadata + .iter() + .map(|metadata| metadata.artifact_id.as_str()) + .collect::>(); let mut written = BTreeMap::new(); let mut written_layer_ordinals = BTreeMap::new(); let mut used = BTreeSet::new(); used.insert(metadata_artifact_id.clone()); for artifact in &manifest.artifact_catalog.entries { - if artifact.id == metadata_artifact_id || sidecar_ids.contains(artifact.id.as_str()) { + if artifact.id == metadata_artifact_id + || sidecar_ids.contains(artifact.id.as_str()) + || publisher_metadata_ids.contains(artifact.id.as_str()) + { continue; } used.insert(artifact.id.clone()); @@ -155,6 +163,12 @@ pub(crate) fn verify_package( ); } verify_projectors(&manifest, &projectors, &artifacts, &mut used)?; + for metadata in &manifest.publisher_metadata { + ensure!( + used.insert(metadata.artifact_id.clone()), + "publisher metadata artifact cannot also have another package role" + ); + } ensure!( used.len() == artifacts.len(), "artifact catalog contains unaccounted artifacts" diff --git a/crates/skippy-model-package/src/verify_v2/tests.rs b/crates/skippy-model-package/src/verify_v2/tests.rs index 852cdb7a2a..57db671d32 100644 --- a/crates/skippy-model-package/src/verify_v2/tests.rs +++ b/crates/skippy-model-package/src/verify_v2/tests.rs @@ -1,7 +1,7 @@ use super::*; use crate::cli::{Args, Command}; use crate::package::ArtifactHook; -use crate::package_v2::write_package; +use crate::package_v2::{PackageSidecars, write_package}; use crate::test_gguf::{FixtureTensor, explicit, fixture, tensor}; use clap::Parser; @@ -31,7 +31,10 @@ impl Case { write_package( self.source.display().to_string(), self.package.clone(), - projectors, + PackageSidecars { + projectors, + publisher_metadata: Vec::new(), + }, ArtifactHook { command: None }, ArtifactHook { command: None }, explicit(&self.source), diff --git a/crates/skippy-model/src/package_carrier.rs b/crates/skippy-model/src/package_carrier.rs index b14f8626cd..edd10d814b 100644 --- a/crates/skippy-model/src/package_carrier.rs +++ b/crates/skippy-model/src/package_carrier.rs @@ -166,6 +166,11 @@ fn payload_artifacts(manifest: &PackageManifest) -> Result> { .iter() .map(|sidecar| sidecar.artifact_id.as_str()) .collect::>(); + let publisher_metadata = manifest + .publisher_metadata + .iter() + .map(|metadata| metadata.artifact_id.as_str()) + .collect::>(); let mut artifacts = manifest .artifact_catalog .entries @@ -173,6 +178,7 @@ fn payload_artifacts(manifest: &PackageManifest) -> Result> { .filter(|artifact| { artifact.id != manifest.source_model.metadata_artifact_id && !sidecars.contains(artifact.id.as_str()) + && !publisher_metadata.contains(artifact.id.as_str()) }) .collect::>(); artifacts.sort_by(|left, right| left.id.cmp(&right.id)); diff --git a/crates/skippy-package-format/src/lib.rs b/crates/skippy-package-format/src/lib.rs index c59ebbe2b7..afa305a815 100644 --- a/crates/skippy-package-format/src/lib.rs +++ b/crates/skippy-package-format/src/lib.rs @@ -23,6 +23,8 @@ pub struct PackageManifest { pub artifact_catalog: ArtifactCatalog, pub tensor_catalog: TensorCatalog, pub sidecars: Vec, + pub publisher_metadata: Vec, + pub publisher_defaults: Option, pub generation: Option, pub native_abi_version: String, pub generator_version: String, @@ -41,6 +43,10 @@ struct PackageRoot { artifact_catalog: ArtifactCatalog, #[serde(default, skip_serializing_if = "Vec::is_empty")] sidecars: Vec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + publisher_metadata: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + publisher_defaults: Option, #[serde(default, skip_serializing_if = "Option::is_none")] generation: Option, native_abi_version: String, @@ -59,6 +65,8 @@ impl From<&PackageManifest> for PackageRoot { layer_count: manifest.layer_count, artifact_catalog: manifest.artifact_catalog.clone(), sidecars: manifest.sidecars.clone(), + publisher_metadata: manifest.publisher_metadata.clone(), + publisher_defaults: manifest.publisher_defaults.clone(), generation: manifest.generation.clone(), native_abi_version: manifest.native_abi_version.clone(), generator_version: manifest.generator_version.clone(), @@ -82,6 +90,8 @@ impl PackageRoot { entries: Vec::new(), }, sidecars: self.sidecars, + publisher_metadata: self.publisher_metadata, + publisher_defaults: self.publisher_defaults, generation: self.generation, native_abi_version: self.native_abi_version, generator_version: self.generator_version, @@ -128,6 +138,8 @@ impl PackageManifest { let artifacts = collect_artifacts(&self.artifact_catalog.entries, &mut issues); validate_metadata_artifact_binding(self, &artifacts, &mut issues); validate_sidecars(&self.sidecars, &artifacts, &mut issues); + validate_publisher_metadata(self, &artifacts, &mut issues); + validate_publisher_defaults(self, &artifacts, &mut issues); if let Some(generation) = &self.generation { validate_generation(generation, self.layer_count, &mut issues); } @@ -157,6 +169,7 @@ impl PackageManifest { .entries .sort_by(|left, right| left.id.cmp(&right.id)); normalized.sidecars.sort(); + normalized.publisher_metadata.sort(); let digest = Sha256::digest(serde_json::to_vec(&normalized)?); let hex = digest .iter() @@ -293,6 +306,60 @@ pub enum SidecarKind { Mmproj, } +/// Immutable publisher files used to derive typed model defaults. +/// +/// These files are package-level metadata. They are not loader sidecars and +/// therefore never participate in stage selection. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PublisherMetadata { + pub role: PublisherMetadataRole, + pub artifact_id: String, + pub source_repo: String, + pub source_revision: String, + pub source_path: String, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PublisherMetadataRole { + ModelConfig, + GenerationConfig, + TokenizerConfig, + ChatTemplate, + HfQuantConfig, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PublisherModelDefaults { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub compute_dtype: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub kv_cache_dtype: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PublisherDtypeDeclaration { + pub dtype: PublisherDtype, + pub artifact_id: String, + pub json_path: String, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PublisherDtype { + F16, + Bf16, + F32, + Fp8, + Fp8E4m3, + Fp8E5m2, + Q8_0, + Q4_0, +} + /// Generation capability declarations carried as package data. /// /// Describes speculative-decoding capabilities of the source model only. @@ -925,6 +992,141 @@ fn validate_sidecars( } } +fn validate_publisher_metadata( + manifest: &PackageManifest, + artifacts: &BTreeMap<&str, &Artifact>, + issues: &mut Vec, +) { + let mut seen = BTreeSet::new(); + for (index, metadata) in manifest.publisher_metadata.iter().enumerate() { + let prefix = format!("publisher_metadata[{index}]"); + if !seen.insert((metadata.role, metadata.source_path.as_str())) { + push_issue( + issues, + ValidationCode::DuplicateSidecar, + prefix.clone(), + format!( + "publisher metadata semantic identity ({:?}, {:?}) appears more than once", + metadata.role, metadata.source_path + ), + ); + } + if !artifacts.contains_key(metadata.artifact_id.as_str()) { + push_issue( + issues, + ValidationCode::UnknownArtifact, + format!("{prefix}.artifact_id"), + format!("artifact {:?} does not exist", metadata.artifact_id), + ); + } + validate_nonempty( + &format!("{prefix}.source_repo"), + &metadata.source_repo, + issues, + ); + validate_nonempty( + &format!("{prefix}.source_revision"), + &metadata.source_revision, + issues, + ); + if metadata.source_revision.len() != 40 + || !metadata + .source_revision + .bytes() + .all(|byte| byte.is_ascii_hexdigit()) + { + push_issue( + issues, + ValidationCode::SourceIdentityMismatch, + format!("{prefix}.source_revision"), + "publisher metadata revision must be an immutable 40-character commit", + ); + } + validate_relative_path( + &format!("{prefix}.source_path"), + &metadata.source_path, + issues, + ); + if manifest.source_model.repo.as_deref() != Some(metadata.source_repo.as_str()) { + push_issue( + issues, + ValidationCode::SourceIdentityMismatch, + format!("{prefix}.source_repo"), + "publisher metadata repository differs from source_model.repo", + ); + } + if manifest.source_model.revision.as_deref() != Some(metadata.source_revision.as_str()) { + push_issue( + issues, + ValidationCode::SourceIdentityMismatch, + format!("{prefix}.source_revision"), + "publisher metadata revision differs from source_model.revision", + ); + } + } +} + +fn validate_publisher_defaults( + manifest: &PackageManifest, + artifacts: &BTreeMap<&str, &Artifact>, + issues: &mut Vec, +) { + let Some(defaults) = &manifest.publisher_defaults else { + return; + }; + for (name, declaration) in [ + ("compute_dtype", defaults.compute_dtype.as_ref()), + ("kv_cache_dtype", defaults.kv_cache_dtype.as_ref()), + ] { + let Some(declaration) = declaration else { + continue; + }; + let path = format!("publisher_defaults.{name}"); + if !artifacts.contains_key(declaration.artifact_id.as_str()) { + push_issue( + issues, + ValidationCode::UnknownArtifact, + format!("{path}.artifact_id"), + format!("artifact {:?} does not exist", declaration.artifact_id), + ); + } + let metadata = manifest + .publisher_metadata + .iter() + .find(|metadata| metadata.artifact_id == declaration.artifact_id); + if metadata.is_none() { + push_issue( + issues, + ValidationCode::SourceIdentityMismatch, + format!("{path}.artifact_id"), + "typed publisher default is not bound to publisher_metadata", + ); + } + if let Some(metadata) = metadata { + let role_is_valid = match name { + "compute_dtype" => metadata.role == PublisherMetadataRole::ModelConfig, + "kv_cache_dtype" => matches!( + metadata.role, + PublisherMetadataRole::ModelConfig | PublisherMetadataRole::HfQuantConfig + ), + _ => false, + }; + if !role_is_valid { + push_issue( + issues, + ValidationCode::SourceIdentityMismatch, + format!("{path}.artifact_id"), + format!( + "publisher default {name} cannot be sourced from {:?}", + metadata.role + ), + ); + } + } + validate_nonempty(&format!("{path}.json_path"), &declaration.json_path, issues); + } +} + fn validate_nonempty(path: &str, value: &str, issues: &mut Vec) { if value.trim().is_empty() { push_issue( diff --git a/crates/skippy-package-format/src/materialization.rs b/crates/skippy-package-format/src/materialization.rs index 3627dd16e3..467ee21f42 100644 --- a/crates/skippy-package-format/src/materialization.rs +++ b/crates/skippy-package-format/src/materialization.rs @@ -360,6 +360,8 @@ mod tests { ], }, sidecars: Vec::new(), + publisher_metadata: Vec::new(), + publisher_defaults: None, generation: None, native_abi_version: "0.1.49".to_string(), generator_version: "test".to_string(), diff --git a/crates/skippy-package-format/src/stage_admission.rs b/crates/skippy-package-format/src/stage_admission.rs index c287d66c31..af2335203b 100644 --- a/crates/skippy-package-format/src/stage_admission.rs +++ b/crates/skippy-package-format/src/stage_admission.rs @@ -606,6 +606,8 @@ mod tests { name: Some("vision".to_string()), }, ], + publisher_metadata: Vec::new(), + publisher_defaults: None, generation: None, native_abi_version: "0.1.49".to_string(), generator_version: "test".to_string(), diff --git a/crates/skippy-package-format/src/tests.rs b/crates/skippy-package-format/src/tests.rs index 68af6ee889..c9bcfb5169 100644 --- a/crates/skippy-package-format/src/tests.rs +++ b/crates/skippy-package-format/src/tests.rs @@ -1113,6 +1113,8 @@ fn fixture() -> PackageManifest { }], }, sidecars: Vec::new(), + publisher_metadata: Vec::new(), + publisher_defaults: None, generation: None, native_abi_version: "7".to_string(), generator_version: "0.76.0-rc9".to_string(), diff --git a/docs/LAYER_PACKAGE_REPOS.md b/docs/LAYER_PACKAGE_REPOS.md index e5d5412ae2..e6da4449e4 100644 --- a/docs/LAYER_PACKAGE_REPOS.md +++ b/docs/LAYER_PACKAGE_REPOS.md @@ -303,6 +303,14 @@ preserves the source pipeline tag in the package model card. This is how a combined vision/audio projector such as Inkling's `mmproj-BF16.gguf` travels with its Q2 layer package. +The job also resolves the requested source revision to an immutable Hugging +Face commit, downloads supported publisher configuration files when present, +and passes them to the package writer. They are uploaded under `metadata/` and +recorded in the manifest with semantic roles, hashes, and source provenance. +Existing packages without these entries continue to load with the safe F16 KV +default; they require a post-merge metadata backfill before publisher defaults +can become effective. + ## Publishing flow The HF Jobs script performs the publishing work: diff --git a/docs/design/SKIPPY_PACKAGE_V2_SCHEMA.md b/docs/design/SKIPPY_PACKAGE_V2_SCHEMA.md index 33ba48469d..9440fe38b1 100644 --- a/docs/design/SKIPPY_PACKAGE_V2_SCHEMA.md +++ b/docs/design/SKIPPY_PACKAGE_V2_SCHEMA.md @@ -35,7 +35,7 @@ The serialized `model-package.json` root contains: - `package_id`: canonical `sha256:` identity; - source and model identities; - an artifact catalog; -- optional projector and generation sidecars; +- optional projector sidecars, typed generation data, and publisher metadata; - native ABI and package-generator versions; - creation time for provenance. @@ -53,8 +53,8 @@ package. The package id is computed by the shared crate as follows: 1. Clone the manifest and replace `package_id` with the empty string. -2. Sort source files by path, artifacts by id, and sidecars by - `(kind, name, artifact_id)`. +2. Sort source files by path, artifacts by id, sidecars by + `(kind, name, artifact_id)`, and publisher metadata by its typed fields. 3. Serialize the normalized root with the shared Rust schema. 4. Hash the serialized bytes with SHA-256 and prefix the lowercase digest with `sha256:`. @@ -92,7 +92,7 @@ keys: Each locator array has exactly one entry per carrier tensor and follows GGUF tensor-directory order. Payload artifacts are indexed by artifact id after -excluding the metadata carrier and sidecars. The runtime rejects an unknown +excluding the metadata carrier, loader sidecars, and publisher metadata. The runtime rejects an unknown locator version, wrong array type or length, invalid part index, invalid alignment, or an extent outside the declared artifact size. @@ -125,6 +125,29 @@ Multiple projectors therefore require stable distinct names; the package writer uses each projector's deterministic artifact id as its name. Generation remains a typed manifest field rather than a generic sidecar. +## Publisher Metadata and Live KV Defaults + +Publisher files are package-level metadata rather than loader sidecars. The +writer accepts `config.json`, `generation_config.json`, +`tokenizer_config.json`, `chat_template.jinja`, and `hf_quant_config.json` via +repeatable `--publisher-metadata` arguments and copies them under `metadata/`. +Each entry records its semantic role, artifact id, source repository, immutable +source revision, and source path; the artifact catalog binds its size and +SHA-256 digest. + +The writer derives only typed defaults used by runtime policy. It records the +publisher compute dtype and any explicit KV-cache dtype together with the +artifact id and JSON path that supplied the value. Common architecture geometry +in `config.json` is compared with authoritative GGUF metadata, and a conflict +fails package creation before payload artifacts are emitted. Weight +quantization and GGUF size never imply a live KV dtype. + +Runtime precedence is explicit user K/V type, then the package's validated KV +declaration, then its compute dtype mapped to a supported live type, then +F16/F16. BF16 maps to F16 until BF16 live KV is qualified; FP8 declarations +also fall back to F16 until the embedded runtime exposes a qualified FP8 type. +Packages without publisher metadata remain readable and use F16/F16. + ## Loading Rule The runtime validates the JSON root, fetches and verifies its declared metadata diff --git a/tools/xtask/data/console_print_allowlist.json b/tools/xtask/data/console_print_allowlist.json index 247ebc9106..bfaff93581 100644 --- a/tools/xtask/data/console_print_allowlist.json +++ b/tools/xtask/data/console_print_allowlist.json @@ -1967,19 +1967,11 @@ ], "crates/mesh-llm-host-runtime/src/runtime/local_split/tests.rs": [ { - "line": 2068, + "line": 2072, "macro_name": "eprintln!" }, { - "line": 2073, - "macro_name": "eprintln!" - }, - { - "line": 2085, - "macro_name": "eprintln!" - }, - { - "line": 2106, + "line": 2077, "macro_name": "eprintln!" } ], @@ -3889,17 +3881,17 @@ "macro_name": "println!" }, { - "line": 137, + "line": 141, "macro_name": "println!" }, { - "line": 157, + "line": 161, "macro_name": "println!" } ], "crates/skippy-model-package/src/package_v2.rs": [ { - "line": 169, + "line": 221, "macro_name": "println!" } ], From 5063e222a9ea96b93c40ffa968bbc90b4c437b62 Mon Sep 17 00:00:00 2001 From: scama Date: Tue, 15 Sep 2026 08:50:49 +1000 Subject: [PATCH 09/16] fix(skippy): bind split roster to consolidated recipe --- .../src/inference/skippy/split-certified.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/split-certified.json b/crates/mesh-llm-host-runtime/src/inference/skippy/split-certified.json index 16f038c692..be5e899f64 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/split-certified.json +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/split-certified.json @@ -2,8 +2,8 @@ "schema_version": 1, "native_recipe": { "llama_upstream_sha": "3057bb66c86c46d5781e50e85462a760ba7d1feb", - "skippy_abi": "0.1.54", - "patch_queue_sha256": "38de596e11761ec5cf778332462003416bd6af2e5f01ed2415308bd4d52a2888" + "skippy_abi": "0.1.55", + "patch_queue_sha256": "3b6788383b952fa6bf4d3e3217309b58ba994258ce7cc0f09e3a4dedb332fab1" }, "models": [ { From c03e1f80bae71308d8289426fc0255a36f1d1c0f Mon Sep 17 00:00:00 2001 From: scama Date: Tue, 15 Sep 2026 09:25:54 +1000 Subject: [PATCH 10/16] fix(skippy): resequence CacheGen patches after main --- .../src/inference/skippy/split-certified.json | 2 +- ....patch => 0038-skippy-define-CacheGen-page-import-ABI.patch} | 0 ...tch => 0039-skippy-expose-CacheGen-backend-capability.patch} | 0 ... 0040-skippy-dispatch-CacheGen-pages-into-resident-KV.patch} | 0 ...041-ggml-metal-decode-CacheGen-pages-into-resident-KV.patch} | 0 ...0042-ggml-cuda-decode-CacheGen-pages-into-resident-KV.patch} | 0 ....patch => 0043-ggml-metal-align-staged-CacheGen-tiles.patch} | 0 ...atch => 0044-ggml-optimize-CacheGen-arithmetic-decode.patch} | 0 ...atch => 0045-ggml-decode-CacheGen-into-F32-KV-tensors.patch} | 0 ... 0046-ggml-restore-quantized-CacheGen-pages-on-device.patch} | 0 ...ctly.patch => 0047-ggml-metal-stage-CacheGen-directly.patch} | 0 ...=> 0048-ggml-decode-packed-CacheGen-symbols-on-device.patch} | 0 ...ch => 0049-ggml-cuda-stage-CacheGen-payloads-directly.patch} | 0 ...h => 0050-ggml-cuda-use-native-CacheGen-shuffle-masks.patch} | 0 14 files changed, 1 insertion(+), 1 deletion(-) rename third_party/llama.cpp/patches/{0034-skippy-define-CacheGen-page-import-ABI.patch => 0038-skippy-define-CacheGen-page-import-ABI.patch} (100%) rename third_party/llama.cpp/patches/{0035-skippy-expose-CacheGen-backend-capability.patch => 0039-skippy-expose-CacheGen-backend-capability.patch} (100%) rename third_party/llama.cpp/patches/{0036-skippy-dispatch-CacheGen-pages-into-resident-KV.patch => 0040-skippy-dispatch-CacheGen-pages-into-resident-KV.patch} (100%) rename third_party/llama.cpp/patches/{0037-ggml-metal-decode-CacheGen-pages-into-resident-KV.patch => 0041-ggml-metal-decode-CacheGen-pages-into-resident-KV.patch} (100%) rename third_party/llama.cpp/patches/{0038-ggml-cuda-decode-CacheGen-pages-into-resident-KV.patch => 0042-ggml-cuda-decode-CacheGen-pages-into-resident-KV.patch} (100%) rename third_party/llama.cpp/patches/{0039-ggml-metal-align-staged-CacheGen-tiles.patch => 0043-ggml-metal-align-staged-CacheGen-tiles.patch} (100%) rename third_party/llama.cpp/patches/{0040-ggml-optimize-CacheGen-arithmetic-decode.patch => 0044-ggml-optimize-CacheGen-arithmetic-decode.patch} (100%) rename third_party/llama.cpp/patches/{0041-ggml-decode-CacheGen-into-F32-KV-tensors.patch => 0045-ggml-decode-CacheGen-into-F32-KV-tensors.patch} (100%) rename third_party/llama.cpp/patches/{0042-ggml-restore-quantized-CacheGen-pages-on-device.patch => 0046-ggml-restore-quantized-CacheGen-pages-on-device.patch} (100%) rename third_party/llama.cpp/patches/{0043-ggml-metal-stage-CacheGen-directly.patch => 0047-ggml-metal-stage-CacheGen-directly.patch} (100%) rename third_party/llama.cpp/patches/{0044-ggml-decode-packed-CacheGen-symbols-on-device.patch => 0048-ggml-decode-packed-CacheGen-symbols-on-device.patch} (100%) rename third_party/llama.cpp/patches/{0045-ggml-cuda-stage-CacheGen-payloads-directly.patch => 0049-ggml-cuda-stage-CacheGen-payloads-directly.patch} (100%) rename third_party/llama.cpp/patches/{0046-ggml-cuda-use-native-CacheGen-shuffle-masks.patch => 0050-ggml-cuda-use-native-CacheGen-shuffle-masks.patch} (100%) diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/split-certified.json b/crates/mesh-llm-host-runtime/src/inference/skippy/split-certified.json index be5e899f64..c3f963143a 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/split-certified.json +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/split-certified.json @@ -3,7 +3,7 @@ "native_recipe": { "llama_upstream_sha": "3057bb66c86c46d5781e50e85462a760ba7d1feb", "skippy_abi": "0.1.55", - "patch_queue_sha256": "3b6788383b952fa6bf4d3e3217309b58ba994258ce7cc0f09e3a4dedb332fab1" + "patch_queue_sha256": "f6e4fe6971508333b246e6204f952405dc24ee21211339e349b51dd90f8304a7" }, "models": [ { diff --git a/third_party/llama.cpp/patches/0034-skippy-define-CacheGen-page-import-ABI.patch b/third_party/llama.cpp/patches/0038-skippy-define-CacheGen-page-import-ABI.patch similarity index 100% rename from third_party/llama.cpp/patches/0034-skippy-define-CacheGen-page-import-ABI.patch rename to third_party/llama.cpp/patches/0038-skippy-define-CacheGen-page-import-ABI.patch diff --git a/third_party/llama.cpp/patches/0035-skippy-expose-CacheGen-backend-capability.patch b/third_party/llama.cpp/patches/0039-skippy-expose-CacheGen-backend-capability.patch similarity index 100% rename from third_party/llama.cpp/patches/0035-skippy-expose-CacheGen-backend-capability.patch rename to third_party/llama.cpp/patches/0039-skippy-expose-CacheGen-backend-capability.patch diff --git a/third_party/llama.cpp/patches/0036-skippy-dispatch-CacheGen-pages-into-resident-KV.patch b/third_party/llama.cpp/patches/0040-skippy-dispatch-CacheGen-pages-into-resident-KV.patch similarity index 100% rename from third_party/llama.cpp/patches/0036-skippy-dispatch-CacheGen-pages-into-resident-KV.patch rename to third_party/llama.cpp/patches/0040-skippy-dispatch-CacheGen-pages-into-resident-KV.patch diff --git a/third_party/llama.cpp/patches/0037-ggml-metal-decode-CacheGen-pages-into-resident-KV.patch b/third_party/llama.cpp/patches/0041-ggml-metal-decode-CacheGen-pages-into-resident-KV.patch similarity index 100% rename from third_party/llama.cpp/patches/0037-ggml-metal-decode-CacheGen-pages-into-resident-KV.patch rename to third_party/llama.cpp/patches/0041-ggml-metal-decode-CacheGen-pages-into-resident-KV.patch diff --git a/third_party/llama.cpp/patches/0038-ggml-cuda-decode-CacheGen-pages-into-resident-KV.patch b/third_party/llama.cpp/patches/0042-ggml-cuda-decode-CacheGen-pages-into-resident-KV.patch similarity index 100% rename from third_party/llama.cpp/patches/0038-ggml-cuda-decode-CacheGen-pages-into-resident-KV.patch rename to third_party/llama.cpp/patches/0042-ggml-cuda-decode-CacheGen-pages-into-resident-KV.patch diff --git a/third_party/llama.cpp/patches/0039-ggml-metal-align-staged-CacheGen-tiles.patch b/third_party/llama.cpp/patches/0043-ggml-metal-align-staged-CacheGen-tiles.patch similarity index 100% rename from third_party/llama.cpp/patches/0039-ggml-metal-align-staged-CacheGen-tiles.patch rename to third_party/llama.cpp/patches/0043-ggml-metal-align-staged-CacheGen-tiles.patch diff --git a/third_party/llama.cpp/patches/0040-ggml-optimize-CacheGen-arithmetic-decode.patch b/third_party/llama.cpp/patches/0044-ggml-optimize-CacheGen-arithmetic-decode.patch similarity index 100% rename from third_party/llama.cpp/patches/0040-ggml-optimize-CacheGen-arithmetic-decode.patch rename to third_party/llama.cpp/patches/0044-ggml-optimize-CacheGen-arithmetic-decode.patch diff --git a/third_party/llama.cpp/patches/0041-ggml-decode-CacheGen-into-F32-KV-tensors.patch b/third_party/llama.cpp/patches/0045-ggml-decode-CacheGen-into-F32-KV-tensors.patch similarity index 100% rename from third_party/llama.cpp/patches/0041-ggml-decode-CacheGen-into-F32-KV-tensors.patch rename to third_party/llama.cpp/patches/0045-ggml-decode-CacheGen-into-F32-KV-tensors.patch diff --git a/third_party/llama.cpp/patches/0042-ggml-restore-quantized-CacheGen-pages-on-device.patch b/third_party/llama.cpp/patches/0046-ggml-restore-quantized-CacheGen-pages-on-device.patch similarity index 100% rename from third_party/llama.cpp/patches/0042-ggml-restore-quantized-CacheGen-pages-on-device.patch rename to third_party/llama.cpp/patches/0046-ggml-restore-quantized-CacheGen-pages-on-device.patch diff --git a/third_party/llama.cpp/patches/0043-ggml-metal-stage-CacheGen-directly.patch b/third_party/llama.cpp/patches/0047-ggml-metal-stage-CacheGen-directly.patch similarity index 100% rename from third_party/llama.cpp/patches/0043-ggml-metal-stage-CacheGen-directly.patch rename to third_party/llama.cpp/patches/0047-ggml-metal-stage-CacheGen-directly.patch diff --git a/third_party/llama.cpp/patches/0044-ggml-decode-packed-CacheGen-symbols-on-device.patch b/third_party/llama.cpp/patches/0048-ggml-decode-packed-CacheGen-symbols-on-device.patch similarity index 100% rename from third_party/llama.cpp/patches/0044-ggml-decode-packed-CacheGen-symbols-on-device.patch rename to third_party/llama.cpp/patches/0048-ggml-decode-packed-CacheGen-symbols-on-device.patch diff --git a/third_party/llama.cpp/patches/0045-ggml-cuda-stage-CacheGen-payloads-directly.patch b/third_party/llama.cpp/patches/0049-ggml-cuda-stage-CacheGen-payloads-directly.patch similarity index 100% rename from third_party/llama.cpp/patches/0045-ggml-cuda-stage-CacheGen-payloads-directly.patch rename to third_party/llama.cpp/patches/0049-ggml-cuda-stage-CacheGen-payloads-directly.patch diff --git a/third_party/llama.cpp/patches/0046-ggml-cuda-use-native-CacheGen-shuffle-masks.patch b/third_party/llama.cpp/patches/0050-ggml-cuda-use-native-CacheGen-shuffle-masks.patch similarity index 100% rename from third_party/llama.cpp/patches/0046-ggml-cuda-use-native-CacheGen-shuffle-masks.patch rename to third_party/llama.cpp/patches/0050-ggml-cuda-use-native-CacheGen-shuffle-masks.patch From 31ae5677af604fde0a8a30c923e6a6d236be5076 Mon Sep 17 00:00:00 2001 From: scama Date: Tue, 15 Sep 2026 09:31:08 +1000 Subject: [PATCH 11/16] fix(skippy): make rebased CacheGen patches portable --- .../src/inference/skippy/split-certified.json | 2 +- ...ippy-define-CacheGen-page-import-ABI.patch | 14 ++--- ...y-expose-CacheGen-backend-capability.patch | 11 ++-- ...atch-CacheGen-pages-into-resident-KV.patch | 10 ++-- ...code-CacheGen-pages-into-resident-KV.patch | 30 +++++----- ...code-CacheGen-pages-into-resident-KV.patch | 10 ++-- ...ml-metal-align-staged-CacheGen-tiles.patch | 12 ++-- ...-optimize-CacheGen-arithmetic-decode.patch | 4 +- ...-decode-CacheGen-into-F32-KV-tensors.patch | 60 +++++++++---------- ...e-quantized-CacheGen-pages-on-device.patch | 33 +++++----- ...7-ggml-metal-stage-CacheGen-directly.patch | 21 ++++--- ...de-packed-CacheGen-symbols-on-device.patch | 28 ++++----- ...uda-stage-CacheGen-payloads-directly.patch | 27 +++++---- ...da-use-native-CacheGen-shuffle-masks.patch | 4 +- 14 files changed, 136 insertions(+), 130 deletions(-) diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/split-certified.json b/crates/mesh-llm-host-runtime/src/inference/skippy/split-certified.json index c3f963143a..d05628b951 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/split-certified.json +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/split-certified.json @@ -3,7 +3,7 @@ "native_recipe": { "llama_upstream_sha": "3057bb66c86c46d5781e50e85462a760ba7d1feb", "skippy_abi": "0.1.55", - "patch_queue_sha256": "f6e4fe6971508333b246e6204f952405dc24ee21211339e349b51dd90f8304a7" + "patch_queue_sha256": "73a98f4e9725d4018b3dc83fe724af882b0790bb2304ae7b35088869d31fd8b4" }, "models": [ { diff --git a/third_party/llama.cpp/patches/0038-skippy-define-CacheGen-page-import-ABI.patch b/third_party/llama.cpp/patches/0038-skippy-define-CacheGen-page-import-ABI.patch index fd806a0b2a..d2cf540d17 100644 --- a/third_party/llama.cpp/patches/0038-skippy-define-CacheGen-page-import-ABI.patch +++ b/third_party/llama.cpp/patches/0038-skippy-define-CacheGen-page-import-ABI.patch @@ -1,7 +1,7 @@ -From 216d925e721653b7134b46c175e679cc5d0f2eeb Mon Sep 17 00:00:00 2001 +From 7e1830a5596e6b2a15a98f6bb6a950edf6b9f862 Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Fri, 11 Sep 2026 16:02:11 +1000 -Subject: [PATCH] skippy: define CacheGen page import ABI +Subject: [PATCH 38/50] skippy: define CacheGen page import ABI Expose validated record descriptors for direct decode into resident KV storage. @@ -12,7 +12,7 @@ Assisted-by: scama 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/include/skippy/common.h b/include/skippy/common.h -index 65cfa2ce8..be996397a 100644 +index 22a116414..ffef99c77 100644 --- a/include/skippy/common.h +++ b/include/skippy/common.h @@ -46,7 +46,7 @@ extern "C" { @@ -24,13 +24,12 @@ index 65cfa2ce8..be996397a 100644 #if defined(_MSC_VER) #define SKIPPY_DEPRECATED(message) __declspec(deprecated(message)) -@@ -104,6 +104,8 @@ enum skippy_feature { - #define SKIPPY_FEATURE_DEVICE_EVENTS ((uint64_t)1 << 34) +@@ -105,6 +105,8 @@ enum skippy_feature { #define SKIPPY_FEATURE_DIAGNOSTIC_EVENTS ((uint64_t)1 << 35) #define SKIPPY_FEATURE_UNLOAD_EVENTS ((uint64_t)1 << 36) -+ -+#define SKIPPY_FEATURE_CACHEGEN_KV_PAGE (UINT64_C(1) << 37) ++#define SKIPPY_FEATURE_CACHEGEN_KV_PAGE (UINT64_C(1) << 37) ++ enum skippy_status { SKIPPY_STATUS_OK = 0, SKIPPY_STATUS_ERROR = 1, @@ -85,3 +84,4 @@ index aa698c124..7504e573c 100644 struct skippy_session * session, -- 2.54.0 (Apple Git-157) + diff --git a/third_party/llama.cpp/patches/0039-skippy-expose-CacheGen-backend-capability.patch b/third_party/llama.cpp/patches/0039-skippy-expose-CacheGen-backend-capability.patch index bf63211633..df9bfccaff 100644 --- a/third_party/llama.cpp/patches/0039-skippy-expose-CacheGen-backend-capability.patch +++ b/third_party/llama.cpp/patches/0039-skippy-expose-CacheGen-backend-capability.patch @@ -1,7 +1,7 @@ -From c95fb1f1001f0e7bce75676e4a6dc3b82b430d14 Mon Sep 17 00:00:00 2001 +From 0cdb0f1be47bc93dbedd86c2aa683f519edc45fb Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Fri, 11 Sep 2026 16:02:47 +1000 -Subject: [PATCH] skippy: expose CacheGen backend capability +Subject: [PATCH 39/50] skippy: expose CacheGen backend capability Advertise the compressed-page contract and return an explicit unsupported status until a resident backend decoder is available. @@ -12,10 +12,10 @@ Assisted-by: scama 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/src/skippy/abi.cpp b/src/skippy/abi.cpp -index ec2b47444..2a0b1b097 100644 +index 7227dfc48..eca3def02 100644 --- a/src/skippy/abi.cpp +++ b/src/skippy/abi.cpp -@@ -46,7 +46,8 @@ uint64_t skippy_abi_features(void) { +@@ -45,7 +45,8 @@ uint64_t skippy_abi_features(void) { SKIPPY_FEATURE_KV_EVENTS | SKIPPY_FEATURE_DEVICE_EVENTS | SKIPPY_FEATURE_DIAGNOSTIC_EVENTS | @@ -26,7 +26,7 @@ index ec2b47444..2a0b1b097 100644 void skippy_error_free(struct skippy_error * error) { diff --git a/src/skippy/state.cpp b/src/skippy/state.cpp -index 620ee3578..af3a56601 100644 +index 8e91d5086..4e1c64349 100644 --- a/src/skippy/state.cpp +++ b/src/skippy/state.cpp @@ -657,6 +657,29 @@ enum skippy_status skippy_import_kv_page( @@ -61,3 +61,4 @@ index 620ee3578..af3a56601 100644 uint64_t token_start, -- 2.54.0 (Apple Git-157) + diff --git a/third_party/llama.cpp/patches/0040-skippy-dispatch-CacheGen-pages-into-resident-KV.patch b/third_party/llama.cpp/patches/0040-skippy-dispatch-CacheGen-pages-into-resident-KV.patch index 6312aed656..ebb3b84a8f 100644 --- a/third_party/llama.cpp/patches/0040-skippy-dispatch-CacheGen-pages-into-resident-KV.patch +++ b/third_party/llama.cpp/patches/0040-skippy-dispatch-CacheGen-pages-into-resident-KV.patch @@ -1,7 +1,7 @@ -From 562610033432885cdbae80b4fe7d038f2e070ae1 Mon Sep 17 00:00:00 2001 +From bfc5a09290e9700011525698cee33cf28369d4c2 Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Fri, 11 Sep 2026 16:42:03 +1000 -Subject: [PATCH] skippy: dispatch CacheGen pages into resident KV +Subject: [PATCH 40/50] skippy: dispatch CacheGen pages into resident KV Assisted-by: scama --- @@ -50,10 +50,10 @@ index cc3f8cd36..e2460fcda 100644 // Backend registry // diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp -index a7d65f23f..dfd00139b 100644 +index a20892114..05c681fe0 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp -@@ -1958,6 +1958,344 @@ bool llama_kv_cache::stage_import_kv_page( +@@ -1984,6 +1984,344 @@ bool llama_kv_cache::stage_import_kv_page( return src == static_cast(input) + input_bytes; } @@ -429,7 +429,7 @@ index 8368e0ee2..ba95e63dc 100644 // graph_build API // diff --git a/src/skippy/state.cpp b/src/skippy/state.cpp -index af3a56601..dc2203735 100644 +index 4e1c64349..bd69a3770 100644 --- a/src/skippy/state.cpp +++ b/src/skippy/state.cpp @@ -657,27 +657,100 @@ enum skippy_status skippy_import_kv_page( diff --git a/third_party/llama.cpp/patches/0041-ggml-metal-decode-CacheGen-pages-into-resident-KV.patch b/third_party/llama.cpp/patches/0041-ggml-metal-decode-CacheGen-pages-into-resident-KV.patch index df7f68fddc..b3148282c5 100644 --- a/third_party/llama.cpp/patches/0041-ggml-metal-decode-CacheGen-pages-into-resident-KV.patch +++ b/third_party/llama.cpp/patches/0041-ggml-metal-decode-CacheGen-pages-into-resident-KV.patch @@ -1,7 +1,7 @@ -From 32d3370571b5b73f136e29f2b1c05bb137493954 Mon Sep 17 00:00:00 2001 +From 24cd07b947867a3e511941de4ce53836b9c7d0f1 Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Fri, 11 Sep 2026 16:42:09 +1000 -Subject: [PATCH] ggml-metal: decode CacheGen pages into resident KV +Subject: [PATCH 41/50] ggml-metal: decode CacheGen pages into resident KV Assisted-by: scama --- @@ -17,10 +17,10 @@ Assisted-by: scama create mode 100644 tests/test-skippy-cachegen-metal.cpp diff --git a/ggml/src/ggml-metal/CMakeLists.txt b/ggml/src/ggml-metal/CMakeLists.txt -index a661e710a..208c8e888 100644 +index e7afdb695..e886ce822 100644 --- a/ggml/src/ggml-metal/CMakeLists.txt +++ b/ggml/src/ggml-metal/CMakeLists.txt -@@ -30,6 +30,7 @@ set(METALLIB_KERNELS_DEQUANTIZE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/dequantize. +@@ -31,6 +31,7 @@ set(METALLIB_KERNELS_DEQUANTIZE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/dequantize. set(METALLIB_KERNELS_QUANTIZE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/quantize.h") set(METALLIB_KERNEL_SOURCES @@ -29,7 +29,7 @@ index a661e710a..208c8e888 100644 kernels/mul_mv.metal kernels/mul_mm.metal diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h -index 31fc07d44..22e4fb436 100644 +index ced33aadf..bd5848382 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -1,6 +1,7 @@ @@ -40,7 +40,7 @@ index 31fc07d44..22e4fb436 100644 #ifdef __cplusplus extern "C" { -@@ -350,6 +351,12 @@ void ggml_metal_buffer_clear (ggml_metal_buffer_t buf, uint8_t value); +@@ -355,6 +356,12 @@ void ggml_metal_buffer_clear (ggml_metal_buffer_t buf, uint8_t value); // struct ggml_metal_buffer_id ggml_metal_buffer_get_id(ggml_metal_buffer_t buf, const struct ggml_tensor * t); @@ -54,10 +54,10 @@ index 31fc07d44..22e4fb436 100644 } #endif diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m -index d6775211e..c6e22a863 100644 +index 3c03d8670..09f00b75e 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m -@@ -11,6 +11,7 @@ +@@ -12,6 +12,7 @@ #include #include @@ -65,7 +65,7 @@ index d6775211e..c6e22a863 100644 #ifndef TARGET_OS_VISION #define TARGET_OS_VISION 0 -@@ -110,6 +111,7 @@ int ggml_metal_pipeline_max_theads_per_threadgroup(struct ggml_metal_pipeline_wi +@@ -111,6 +112,7 @@ int ggml_metal_pipeline_max_theads_per_threadgroup(struct ggml_metal_pipeline_wi // X(suffix, name): name is both the kernels/.metal basename and the // ggml_metallib__{start,end} embed-symbol stem. #define GGML_METAL_LIBS \ @@ -73,7 +73,7 @@ index d6775211e..c6e22a863 100644 X(FA, fa) \ X(MUL_MV, mul_mv) \ X(MUL_MM, mul_mm) \ -@@ -2641,3 +2643,209 @@ struct ggml_metal_buffer_id ggml_metal_buffer_get_id(ggml_metal_buffer_t buf, co +@@ -2658,3 +2660,209 @@ struct ggml_metal_buffer_id ggml_metal_buffer_get_id(ggml_metal_buffer_t buf, co return res; } @@ -284,10 +284,10 @@ index d6775211e..c6e22a863 100644 + return true; +} diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp -index 3bd6abd06..6bec83efd 100644 +index 4cbec8645..10d005482 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp -@@ -910,6 +910,9 @@ static void * ggml_backend_metal_get_proc_address(ggml_backend_reg_t reg, const +@@ -935,6 +935,9 @@ static void * ggml_backend_metal_get_proc_address(ggml_backend_reg_t reg, const if (strcmp(name, "ggml_backend_get_features") == 0) { return (void *)ggml_backend_metal_get_features; } @@ -402,11 +402,11 @@ index 000000000..82684a71e + } +} diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt -index 5a45037bb..4f57384ca 100644 +index 69dbf219f..dc30b99a2 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt -@@ -279,6 +279,10 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) - set_tests_properties(test-save-load-state PROPERTIES FIXTURES_REQUIRED generate-models) +@@ -281,6 +281,10 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) + llama_build(test-fusion.cpp) endif() +if (GGML_METAL) diff --git a/third_party/llama.cpp/patches/0042-ggml-cuda-decode-CacheGen-pages-into-resident-KV.patch b/third_party/llama.cpp/patches/0042-ggml-cuda-decode-CacheGen-pages-into-resident-KV.patch index 111a06bf73..71524f14c7 100644 --- a/third_party/llama.cpp/patches/0042-ggml-cuda-decode-CacheGen-pages-into-resident-KV.patch +++ b/third_party/llama.cpp/patches/0042-ggml-cuda-decode-CacheGen-pages-into-resident-KV.patch @@ -1,7 +1,7 @@ -From 3391262f8620a3d1a7ddb6a174e75e0c6129c804 Mon Sep 17 00:00:00 2001 +From 62b87b9c385e333765835739e38eead150a47dae Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Fri, 11 Sep 2026 16:58:10 +1000 -Subject: [PATCH] ggml-cuda: decode CacheGen pages into resident KV +Subject: [PATCH 42/50] ggml-cuda: decode CacheGen pages into resident KV Share the F16 decoder across CUDA and HIP, batch launches before synchronization, and write directly into the allocated Skippy cell layout. @@ -164,7 +164,7 @@ index 000000000..d6d216675 + +#endif // !defined(GGML_USE_MUSA) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu -index 38bd4c9a0..e955c4394 100644 +index a23417926..1830c5ac0 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -10,6 +10,7 @@ @@ -183,7 +183,7 @@ index 38bd4c9a0..e955c4394 100644 #include #include #include -@@ -5667,6 +5669,280 @@ static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t +@@ -5669,6 +5671,280 @@ static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t GGML_UNUSED(reg); } @@ -464,7 +464,7 @@ index 38bd4c9a0..e955c4394 100644 static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) { GGML_UNUSED(reg); if (strcmp(name, "ggml_backend_comm_init") == 0) { -@@ -5687,6 +5963,11 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con +@@ -5689,6 +5965,11 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con if (strcmp(name, "ggml_backend_get_features") == 0) { return (void *)ggml_backend_cuda_get_features; } diff --git a/third_party/llama.cpp/patches/0043-ggml-metal-align-staged-CacheGen-tiles.patch b/third_party/llama.cpp/patches/0043-ggml-metal-align-staged-CacheGen-tiles.patch index 697b5c3c13..63de34ed8e 100644 --- a/third_party/llama.cpp/patches/0043-ggml-metal-align-staged-CacheGen-tiles.patch +++ b/third_party/llama.cpp/patches/0043-ggml-metal-align-staged-CacheGen-tiles.patch @@ -1,7 +1,7 @@ -From 717bca3cc3cc9ef8272cb37c63c7c700ab567009 Mon Sep 17 00:00:00 2001 +From 90e967661980eaf7008d17b1582ed1939d05a2f1 Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Fri, 11 Sep 2026 17:43:11 +1000 -Subject: [PATCH] ggml-metal: align staged CacheGen tiles +Subject: [PATCH 43/50] ggml-metal: align staged CacheGen tiles Assisted-by: scama --- @@ -10,10 +10,10 @@ Assisted-by: scama 2 files changed, 39 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m -index c6e22a863..d05b7c32c 100644 +index 09f00b75e..b1a8c2ae9 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m -@@ -2721,7 +2721,7 @@ bool ggml_metal_cachegen_decode_f16( +@@ -2738,7 +2738,7 @@ bool ggml_metal_cachegen_decode_f16( const struct ggml_backend_cachegen_tile * tile = &job->tiles[tile_index]; if (tile->payload == NULL || tile->payload_bytes < 16 || tile->token_count == 0 || tile->token_offset > job->cell_count || tile->token_count > job->cell_count - tile->token_offset || @@ -22,7 +22,7 @@ index c6e22a863..d05b7c32c 100644 prefix_data.length / sizeof(uint32_t) > UINT32_MAX - ((size_t) job->channels + 1)) { [payload_data release]; [tile_data release]; -@@ -2730,6 +2730,17 @@ bool ggml_metal_cachegen_decode_f16( +@@ -2747,6 +2747,17 @@ bool ggml_metal_cachegen_decode_f16( [temporary_buffers release]; return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen tile exceeds bounded geometry"); } @@ -40,7 +40,7 @@ index c6e22a863..d05b7c32c 100644 const uint8_t * payload = (const uint8_t *) tile->payload; const size_t lengths_offset = 16 + (size_t) tile->token_count * sizeof(float) + (size_t) job->channels * 33 * sizeof(uint16_t); -@@ -2743,7 +2754,7 @@ bool ggml_metal_cachegen_decode_f16( +@@ -2760,7 +2771,7 @@ bool ggml_metal_cachegen_decode_f16( return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen tile metadata is truncated"); } struct ggml_metal_cachegen_tile encoded_tile = { diff --git a/third_party/llama.cpp/patches/0044-ggml-optimize-CacheGen-arithmetic-decode.patch b/third_party/llama.cpp/patches/0044-ggml-optimize-CacheGen-arithmetic-decode.patch index 76ae875671..b9dc4a3586 100644 --- a/third_party/llama.cpp/patches/0044-ggml-optimize-CacheGen-arithmetic-decode.patch +++ b/third_party/llama.cpp/patches/0044-ggml-optimize-CacheGen-arithmetic-decode.patch @@ -1,7 +1,7 @@ -From c18e4c75324d89f866236ca2fa631c23801a0830 Mon Sep 17 00:00:00 2001 +From a5b5ef5b7983961e33d136117dd3daab77d739d3 Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Fri, 11 Sep 2026 17:51:30 +1000 -Subject: [PATCH] ggml: optimize CacheGen arithmetic decode +Subject: [PATCH 44/50] ggml: optimize CacheGen arithmetic decode Assisted-by: scama --- diff --git a/third_party/llama.cpp/patches/0045-ggml-decode-CacheGen-into-F32-KV-tensors.patch b/third_party/llama.cpp/patches/0045-ggml-decode-CacheGen-into-F32-KV-tensors.patch index a732fb8818..c5ba9855f9 100644 --- a/third_party/llama.cpp/patches/0045-ggml-decode-CacheGen-into-F32-KV-tensors.patch +++ b/third_party/llama.cpp/patches/0045-ggml-decode-CacheGen-into-F32-KV-tensors.patch @@ -1,7 +1,7 @@ -From 980d2df45910b4d07138a863811911c1a2d8f6d0 Mon Sep 17 00:00:00 2001 +From 8b5e8d04923fd86b78d3265e2e1850d3bc8a8f74 Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Fri, 11 Sep 2026 18:22:56 +1000 -Subject: [PATCH] ggml: decode CacheGen into F32 KV tensors +Subject: [PATCH 45/50] ggml: decode CacheGen into F32 KV tensors Replace the F16-specific backend hook with a typed decoder and write the entropy-decoded half values directly into F16 or F32 resident layouts. @@ -128,10 +128,10 @@ index d6d216675..7118c9aa9 100644 #endif // !defined(GGML_USE_MUSA) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu -index e955c4394..3288653f5 100644 +index 1830c5ac0..a5a7c99ca 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu -@@ -5675,6 +5675,7 @@ struct ggml_cuda_cachegen_staged_job { +@@ -5677,6 +5677,7 @@ struct ggml_cuda_cachegen_staged_job { int physical_device = -1; uint8_t * dst = nullptr; uint32_t channels = 0; @@ -139,7 +139,7 @@ index e955c4394..3288653f5 100644 uint64_t token_stride = 0; uint64_t channel_stride = 0; std::vector payload; -@@ -5749,7 +5750,7 @@ static cudaError_t ggml_cuda_cachegen_release(std::vectorcells == nullptr || job->cell_count == 0 || job->cell_count > UINT32_MAX || job->cell_count > SIZE_MAX / sizeof(uint32_t) || job->channels == 0 || job->channels == UINT32_MAX || @@ -161,7 +161,7 @@ index e955c4394..3288653f5 100644 return ggml_cuda_cachegen_error(error, error_capacity, "invalid CUDA/HIP CacheGen job"); } -@@ -5780,6 +5783,7 @@ static bool ggml_cuda_cachegen_decode_f16( +@@ -5782,6 +5785,7 @@ static bool ggml_cuda_cachegen_decode_f16( staged.physical_device = ggml_cuda_get_physical_device(buffer_context->device); staged.dst = static_cast(job->dst->data); staged.channels = job->channels; @@ -169,7 +169,7 @@ index e955c4394..3288653f5 100644 staged.token_stride = job->token_stride; staged.channel_stride = job->channel_stride; staged.cells.assign(job->cells, job->cells + job->cell_count); -@@ -5789,10 +5793,10 @@ static bool ggml_cuda_cachegen_decode_f16( +@@ -5791,10 +5795,10 @@ static bool ggml_cuda_cachegen_decode_f16( max_cell = std::max(max_cell, static_cast(cell)); } const uint64_t max_channel = static_cast(job->channels) - 1; @@ -183,7 +183,7 @@ index e955c4394..3288653f5 100644 ggml_nbytes(job->dst)) { return ggml_cuda_cachegen_error( error, error_capacity, "CUDA/HIP CacheGen destination geometry is out of bounds"); -@@ -5917,10 +5921,10 @@ static bool ggml_cuda_cachegen_decode_f16( +@@ -5919,10 +5923,10 @@ static bool ggml_cuda_cachegen_decode_f16( cudaMemcpyHostToDevice, cudaStreamPerThread), "CUDA/HIP CacheGen cell upload failed"); GGML_CACHEGEN_CUDA_CALL( @@ -198,7 +198,7 @@ index e955c4394..3288653f5 100644 "CUDA/HIP CacheGen kernel launch failed"); #undef GGML_CACHEGEN_CUDA_CALL } -@@ -5964,8 +5968,8 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con +@@ -5966,8 +5970,8 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con return (void *)ggml_backend_cuda_get_features; } #if !defined(GGML_USE_MUSA) @@ -210,10 +210,10 @@ index e955c4394..3288653f5 100644 #endif // !defined(GGML_USE_MUSA) return nullptr; diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h -index 22e4fb436..0a2cf000c 100644 +index bd5848382..6054f28cc 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h -@@ -351,7 +351,7 @@ void ggml_metal_buffer_clear (ggml_metal_buffer_t buf, uint8_t value); +@@ -356,7 +356,7 @@ void ggml_metal_buffer_clear (ggml_metal_buffer_t buf, uint8_t value); // struct ggml_metal_buffer_id ggml_metal_buffer_get_id(ggml_metal_buffer_t buf, const struct ggml_tensor * t); @@ -223,10 +223,10 @@ index 22e4fb436..0a2cf000c 100644 size_t job_count, char * error, diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m -index d05b7c32c..d4ae1f475 100644 +index b1a8c2ae9..7dc37895b 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m -@@ -2654,6 +2654,7 @@ struct ggml_metal_buffer_id ggml_metal_buffer_get_id(ggml_metal_buffer_t buf, co +@@ -2671,6 +2671,7 @@ struct ggml_metal_buffer_id ggml_metal_buffer_get_id(ggml_metal_buffer_t buf, co struct ggml_metal_cachegen_params { uint32_t channels; uint32_t tile_count; @@ -234,7 +234,7 @@ index d05b7c32c..d4ae1f475 100644 uint64_t token_stride; uint64_t channel_stride; }; -@@ -2665,7 +2666,7 @@ static bool ggml_metal_cachegen_error(char * error, size_t capacity, const char +@@ -2682,7 +2683,7 @@ static bool ggml_metal_cachegen_error(char * error, size_t capacity, const char return false; } @@ -243,7 +243,7 @@ index d05b7c32c..d4ae1f475 100644 const struct ggml_backend_cachegen_job * jobs, size_t job_count, char * error, -@@ -2683,7 +2684,10 @@ bool ggml_metal_cachegen_decode_f16( +@@ -2700,7 +2701,10 @@ bool ggml_metal_cachegen_decode_f16( const struct ggml_backend_cachegen_job * job = &jobs[job_index]; if (job->dst == NULL || job->dst->buffer == NULL || job->tiles == NULL || job->tile_count == 0 || job->tile_count > UINT32_MAX || job->cells == NULL || job->cell_count == 0 || job->channels == 0 || @@ -255,7 +255,7 @@ index d05b7c32c..d4ae1f475 100644 job->cell_count > SIZE_MAX / sizeof(uint32_t)) { [temporary_buffers release]; return ggml_metal_cachegen_error(error, error_capacity, "invalid Metal CacheGen job"); -@@ -2693,9 +2697,10 @@ bool ggml_metal_cachegen_decode_f16( +@@ -2710,9 +2714,10 @@ bool ggml_metal_cachegen_decode_f16( max_cell = MAX(max_cell, job->cells[cell_index]); } const uint64_t max_channel = (uint64_t) job->channels - 1; @@ -269,7 +269,7 @@ index d05b7c32c..d4ae1f475 100644 [temporary_buffers release]; return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen destination geometry is out of bounds"); } -@@ -2820,7 +2825,7 @@ bool ggml_metal_cachegen_decode_f16( +@@ -2837,7 +2842,7 @@ bool ggml_metal_cachegen_decode_f16( const struct ggml_metal_pipeline_with_params pipeline = ggml_metal_library_compile_pipeline(device->library, @@ -278,7 +278,7 @@ index d05b7c32c..d4ae1f475 100644 if (pipeline.pipeline == NULL) { [encoder endEncoding]; [temporary_buffers release]; -@@ -2835,6 +2840,7 @@ bool ggml_metal_cachegen_decode_f16( +@@ -2852,6 +2857,7 @@ bool ggml_metal_cachegen_decode_f16( const struct ggml_metal_cachegen_params params = { job->channels, (uint32_t) job->tile_count, @@ -287,10 +287,10 @@ index d05b7c32c..d4ae1f475 100644 job->channel_stride, }; diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp -index 6bec83efd..a1c190a82 100644 +index 10d005482..cd2c135d2 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp -@@ -910,8 +910,8 @@ static void * ggml_backend_metal_get_proc_address(ggml_backend_reg_t reg, const +@@ -935,8 +935,8 @@ static void * ggml_backend_metal_get_proc_address(ggml_backend_reg_t reg, const if (strcmp(name, "ggml_backend_get_features") == 0) { return (void *)ggml_backend_metal_get_features; } @@ -361,10 +361,10 @@ index 7504e573c..9c78b3698 100644 /** @brief Describes one validated CacheGen record and its logical page destination. */ diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp -index dfd00139b..0671fb296 100644 +index 05c681fe0..38b1b0ab1 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp -@@ -1970,17 +1970,18 @@ static uint32_t skippy_cachegen_read_u32(const uint8_t * bytes) { +@@ -1996,17 +1996,18 @@ static uint32_t skippy_cachegen_read_u32(const uint8_t * bytes) { static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & record, size_t expected_rows, size_t expected_channels, @@ -385,7 +385,7 @@ index dfd00139b..0671fb296 100644 return false; } if (std::memcmp(payload, "LCG1", 4) != 0 || payload[5] != 0 || (payload[4] != 16 && payload[4] != 32) || -@@ -1990,8 +1991,8 @@ static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & r +@@ -2016,8 +2017,8 @@ static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & r return false; } if (expected_rows > std::numeric_limits::max() / expected_channels || @@ -396,7 +396,7 @@ index dfd00139b..0671fb296 100644 error = "CacheGen record decoded geometry is inconsistent"; return false; } -@@ -2086,8 +2087,12 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id +@@ -2112,8 +2113,12 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id if (!stage_import_kv_page(seq_id, desc, &validation_sentinel, desc.payload_bytes, error, true)) { return false; } @@ -411,7 +411,7 @@ index dfd00139b..0671fb296 100644 return false; } -@@ -2104,12 +2109,17 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id +@@ -2130,12 +2135,17 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id size_t logical_offset = output_base; std::vector storage; storage.reserve(selected.size() * 2); @@ -433,7 +433,7 @@ index dfd00139b..0671fb296 100644 entry.channels = static_cast(channels); size_t token_offset = 0; while (token_offset < desc.token_count) { -@@ -2119,13 +2129,14 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id +@@ -2145,13 +2155,14 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id } const auto & record = records[record_index++]; const size_t rows = static_cast(record.token_count); @@ -451,7 +451,7 @@ index dfd00139b..0671fb296 100644 if (error.empty()) { error = "CacheGen record order or destination geometry is invalid"; } -@@ -2141,14 +2152,14 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id +@@ -2167,14 +2178,14 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id }; for (const auto * layer : selected) { @@ -468,7 +468,7 @@ index dfd00139b..0671fb296 100644 return false; } } -@@ -2182,7 +2193,7 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id +@@ -2208,7 +2219,7 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id return false; } @@ -477,7 +477,7 @@ index dfd00139b..0671fb296 100644 for (const auto & entry : storage) { if (entry.dst == nullptr || entry.dst->buffer == nullptr) { error = "CacheGen destination tensor has no backend buffer"; -@@ -2193,12 +2204,12 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id +@@ -2219,12 +2230,12 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id const auto reg = device == nullptr ? nullptr : ggml_backend_dev_backend_reg(device); auto candidate = reg == nullptr ? nullptr : diff --git a/third_party/llama.cpp/patches/0046-ggml-restore-quantized-CacheGen-pages-on-device.patch b/third_party/llama.cpp/patches/0046-ggml-restore-quantized-CacheGen-pages-on-device.patch index 15fb3fad3a..7d992365be 100644 --- a/third_party/llama.cpp/patches/0046-ggml-restore-quantized-CacheGen-pages-on-device.patch +++ b/third_party/llama.cpp/patches/0046-ggml-restore-quantized-CacheGen-pages-on-device.patch @@ -1,7 +1,8 @@ -From 81b1c8f52510ff07b4171020a5892f0a1e5bcead Mon Sep 17 00:00:00 2001 +From a3d5e0772844fe04fbbe55705af946e094fb1141 Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Fri, 11 Sep 2026 18:49:55 +1000 -Subject: [PATCH] feat(skippy): restore quantized CacheGen pages on device +Subject: [PATCH 46/50] feat(skippy): restore quantized CacheGen pages on + device Assisted-by: scama --- @@ -75,10 +76,10 @@ index cc51a634c..fa5683788 100644 } diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu -index 3288653f5..b68649ec4 100644 +index a5a7c99ca..819deb14c 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu -@@ -5764,13 +5764,17 @@ static bool ggml_cuda_cachegen_decode( +@@ -5766,13 +5766,17 @@ static bool ggml_cuda_cachegen_decode( staged_jobs.reserve(job_count); for (size_t job_index = 0; job_index < job_count; ++job_index) { const struct ggml_backend_cachegen_job * job = &jobs[job_index]; @@ -97,7 +98,7 @@ index 3288653f5..b68649ec4 100644 job->token_stride == 0 || job->channel_stride == 0 || job->token_stride % ggml_type_size(job->dst->type) != 0 || job->channel_stride % ggml_type_size(job->dst->type) != 0) { -@@ -5792,7 +5796,8 @@ static bool ggml_cuda_cachegen_decode( +@@ -5794,7 +5798,8 @@ static bool ggml_cuda_cachegen_decode( for (uint32_t cell : staged.cells) { max_cell = std::max(max_cell, static_cast(cell)); } @@ -108,10 +109,10 @@ index 3288653f5..b68649ec4 100644 max_channel > (UINT64_MAX - max_cell * job->token_stride - staged.element_bytes) / job->channel_stride || diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m -index d4ae1f475..0092c558f 100644 +index 7dc37895b..0f7eb1839 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m -@@ -2682,9 +2682,13 @@ bool ggml_metal_cachegen_decode( +@@ -2699,9 +2699,13 @@ bool ggml_metal_cachegen_decode( NSMutableArray * temporary_buffers = [[NSMutableArray alloc] init]; for (size_t job_index = 0; job_index < job_count; ++job_index) { const struct ggml_backend_cachegen_job * job = &jobs[job_index]; @@ -126,7 +127,7 @@ index d4ae1f475..0092c558f 100644 job->token_stride == 0 || job->channel_stride == 0 || job->token_stride % ggml_type_size(job->dst->type) != 0 || job->channel_stride % ggml_type_size(job->dst->type) != 0 || -@@ -2696,7 +2700,8 @@ bool ggml_metal_cachegen_decode( +@@ -2713,7 +2717,8 @@ bool ggml_metal_cachegen_decode( for (size_t cell_index = 0; cell_index < job->cell_count; ++cell_index) { max_cell = MAX(max_cell, job->cells[cell_index]); } @@ -136,7 +137,7 @@ index d4ae1f475..0092c558f 100644 const uint64_t element_bytes = ggml_type_size(job->dst->type); if (max_cell > (UINT64_MAX - element_bytes) / job->token_stride || max_channel > (UINT64_MAX - max_cell * job->token_stride - element_bytes) / job->channel_stride || -@@ -2851,8 +2856,9 @@ bool ggml_metal_cachegen_decode( +@@ -2868,8 +2873,9 @@ bool ggml_metal_cachegen_decode( [encoder setBuffer:temporary_buffers[temporary_buffers.count - 1] offset:0 atIndex:3]; [encoder setBuffer:dst.metal offset:dst.offs atIndex:4]; [encoder setBytes:¶ms length:sizeof(params) atIndex:5]; @@ -226,10 +227,10 @@ index 9c78b3698..981da6b5d 100644 /** @brief Describes one validated CacheGen record and its logical page destination. */ diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp -index 0671fb296..215076924 100644 +index 38b1b0ab1..e6888f096 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp -@@ -1971,6 +1971,7 @@ static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & r +@@ -1997,6 +1997,7 @@ static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & r size_t expected_rows, size_t expected_channels, size_t output_element_bytes, @@ -237,7 +238,7 @@ index 0671fb296..215076924 100644 std::string & error) { constexpr size_t header_bytes = 16; constexpr size_t cdf_entries = 33; -@@ -1991,8 +1992,8 @@ static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & r +@@ -2017,8 +2018,8 @@ static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & r return false; } if (expected_rows > std::numeric_limits::max() / expected_channels || @@ -248,7 +249,7 @@ index 0671fb296..215076924 100644 error = "CacheGen record decoded geometry is inconsistent"; return false; } -@@ -2088,11 +2089,11 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id +@@ -2114,11 +2115,11 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id return false; } const auto cachegen_type_supported = [](uint32_t type) { @@ -262,7 +263,7 @@ index 0671fb296..215076924 100644 return false; } -@@ -2109,17 +2110,25 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id +@@ -2135,17 +2136,25 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id size_t logical_offset = output_base; std::vector storage; storage.reserve(selected.size() * 2); @@ -292,7 +293,7 @@ index 0671fb296..215076924 100644 entry.channels = static_cast(channels); size_t token_offset = 0; while (token_offset < desc.token_count) { -@@ -2129,14 +2138,28 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id +@@ -2155,14 +2164,28 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id } const auto & record = records[record_index++]; const size_t rows = static_cast(record.token_count); @@ -325,7 +326,7 @@ index 0671fb296..215076924 100644 if (error.empty()) { error = "CacheGen record order or destination geometry is invalid"; } -@@ -2152,14 +2175,14 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id +@@ -2178,14 +2201,14 @@ bool llama_kv_cache::stage_import_cachegen_kv_page(llama_seq_id }; for (const auto * layer : selected) { diff --git a/third_party/llama.cpp/patches/0047-ggml-metal-stage-CacheGen-directly.patch b/third_party/llama.cpp/patches/0047-ggml-metal-stage-CacheGen-directly.patch index 1d110dd6e4..ad3a7fb81d 100644 --- a/third_party/llama.cpp/patches/0047-ggml-metal-stage-CacheGen-directly.patch +++ b/third_party/llama.cpp/patches/0047-ggml-metal-stage-CacheGen-directly.patch @@ -1,19 +1,21 @@ -From: scama -Date: Fri, 12 Sep 2026 01:00:00 +1000 -Subject: [PATCH] ggml-metal: stage CacheGen payloads directly +From a347b8f31a7446b61c02559145735105949385cd Mon Sep 17 00:00:00 2001 +From: scama + +Date: Sat, 12 Sep 2026 01:00:00 +1000 +Subject: [PATCH 47/50] ggml-metal: stage CacheGen payloads directly Avoid building each CacheGen job in NSMutableData and then copying the same bytes again into shared Metal buffers. Validate and size the job first, then fill its final MTLBuffers directly. --- - ggml/src/ggml-metal/ggml-metal-device.m | 85 ++++++++++++------------ - 1 file changed, 43 insertions(+), 42 deletions(-) + ggml/src/ggml-metal/ggml-metal-device.m | 84 ++++++++++++------------- + 1 file changed, 42 insertions(+), 42 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m -index 0092c558f..79d40466f 100644 +index 0f7eb1839..da750cf60 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m -@@ -2724,92 +2724,63 @@ bool ggml_metal_cachegen_decode( +@@ -2741,92 +2741,63 @@ bool ggml_metal_cachegen_decode( } } @@ -119,7 +121,7 @@ index 0092c558f..79d40466f 100644 if (payload_buffer == nil || tile_buffer == nil || prefix_buffer == nil || cell_buffer == nil) { [payload_buffer release]; [tile_buffer release]; -@@ -2819,6 +2790,35 @@ bool ggml_metal_cachegen_decode( +@@ -2836,6 +2807,35 @@ bool ggml_metal_cachegen_decode( [temporary_buffers release]; return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen staging allocation failed"); } @@ -156,4 +158,5 @@ index 0092c558f..79d40466f 100644 [temporary_buffers addObject:tile_buffer]; [temporary_buffers addObject:prefix_buffer]; -- -2.52.0 +2.54.0 (Apple Git-157) + diff --git a/third_party/llama.cpp/patches/0048-ggml-decode-packed-CacheGen-symbols-on-device.patch b/third_party/llama.cpp/patches/0048-ggml-decode-packed-CacheGen-symbols-on-device.patch index 299468d419..a1853b99e3 100644 --- a/third_party/llama.cpp/patches/0048-ggml-decode-packed-CacheGen-symbols-on-device.patch +++ b/third_party/llama.cpp/patches/0048-ggml-decode-packed-CacheGen-symbols-on-device.patch @@ -1,7 +1,7 @@ -From 9253c686456c6f90182cd49094c96ec1cb9291c3 Mon Sep 17 00:00:00 2001 +From 7e8c8fb1571509c24b970cfcf172521051fa4d44 Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Sat, 12 Sep 2026 05:42:09 +1000 -Subject: [PATCH] ggml: decode packed CacheGen symbols on device +Subject: [PATCH 48/50] ggml: decode packed CacheGen symbols on device --- ggml/src/ggml-cuda/cachegen.cu | 52 ++++++++++---- @@ -93,10 +93,10 @@ index fa5683788..685dc586a 100644 const uint64_t cdf_high = symbol == 31 ? 65536ull : cdf[symbol + 1]; high = low - 1 + uint32_t((span * cdf_high) >> 16); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu -index b68649ec4..ccfda8a9d 100644 +index 819deb14c..e72a09eb6 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu -@@ -5814,7 +5814,7 @@ static bool ggml_cuda_cachegen_decode( +@@ -5816,7 +5816,7 @@ static bool ggml_cuda_cachegen_decode( for (size_t tile_index = 0; tile_index < job->tile_count; ++tile_index) { const struct ggml_backend_cachegen_tile * tile = &job->tiles[tile_index]; if (tile->payload == nullptr || tile->payload_bytes < 16 || tile->payload_bytes > UINT32_MAX || @@ -105,7 +105,7 @@ index b68649ec4..ccfda8a9d 100644 tile->token_count > job->cell_count - tile->token_offset || staged.payload.size() > UINT32_MAX - 3 || staged.prefixes.size() > UINT32_MAX - (static_cast(job->channels) + 1)) { -@@ -5830,41 +5830,67 @@ static bool ggml_cuda_cachegen_decode( +@@ -5832,41 +5832,67 @@ static bool ggml_cuda_cachegen_decode( staged.payload.resize(aligned_payload_size, 0); const uint8_t * payload = static_cast(tile->payload); @@ -198,10 +198,10 @@ index b68649ec4..ccfda8a9d 100644 staged.tiles.push_back(encoded_tile); staged.payload.insert(staged.payload.end(), payload, payload + tile->payload_bytes); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m -index 79d40466f..2a38e0d4b 100644 +index da750cf60..3d10a182c 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m -@@ -2686,6 +2686,7 @@ bool ggml_metal_cachegen_decode( +@@ -2703,6 +2703,7 @@ bool ggml_metal_cachegen_decode( (job->dst->type == GGML_TYPE_Q8_0 || job->dst->type == GGML_TYPE_Q4_0); if (job->dst == NULL || job->dst->buffer == NULL || job->tiles == NULL || job->tile_count == 0 || job->tile_count > UINT32_MAX || job->cells == NULL || job->cell_count == 0 || job->channels == 0 || @@ -209,7 +209,7 @@ index 79d40466f..2a38e0d4b 100644 (job->dst->type != GGML_TYPE_F16 && job->dst->type != GGML_TYPE_F32 && job->dst->type != GGML_TYPE_Q8_0 && job->dst->type != GGML_TYPE_Q4_0) || (quantized && job->channels % 32 != 0) || -@@ -2729,6 +2730,7 @@ bool ggml_metal_cachegen_decode( +@@ -2746,6 +2747,7 @@ bool ggml_metal_cachegen_decode( for (size_t tile_index = 0; tile_index < job->tile_count; ++tile_index) { const struct ggml_backend_cachegen_tile * tile = &job->tiles[tile_index]; if (tile->payload == NULL || tile->payload_bytes < 16 || tile->token_count == 0 || @@ -217,7 +217,7 @@ index 79d40466f..2a38e0d4b 100644 tile->token_offset > job->cell_count || tile->token_count > job->cell_count - tile->token_offset || staged_payload_bytes > UINT32_MAX - 3 || staged_prefix_count > UINT32_MAX - ((size_t) job->channels + 1)) { -@@ -2743,34 +2745,60 @@ bool ggml_metal_cachegen_decode( +@@ -2760,34 +2762,60 @@ bool ggml_metal_cachegen_decode( return ggml_metal_cachegen_error(error, error_capacity, "Metal CacheGen payload offsets overflow"); } const uint8_t * payload = (const uint8_t *) tile->payload; @@ -298,7 +298,7 @@ index 79d40466f..2a38e0d4b 100644 } id payload_buffer = [device->mtl_device newBufferWithLength:staged_payload_bytes -@@ -2808,14 +2836,16 @@ bool ggml_metal_cachegen_decode( +@@ -2825,14 +2853,16 @@ bool ggml_metal_cachegen_decode( tile->token_count, }; const uint8_t * payload = (const uint8_t *) tile->payload; @@ -404,10 +404,10 @@ index 18c0a0fb4..728899e4b 100644 const ulong cdf_high = symbol == 31 ? 65536ul : cdf[symbol + 1]; high = low - 1 + uint((span * cdf_high) >> 16); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp -index 215076924..dad99f62a 100644 +index e6888f096..43e2539b0 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp -@@ -1985,7 +1985,11 @@ static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & r +@@ -2011,7 +2011,11 @@ static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & r error = "invalid CacheGen record descriptor"; return false; } @@ -420,7 +420,7 @@ index 215076924..dad99f62a 100644 skippy_cachegen_read_u16(payload + 6) != expected_rows || skippy_cachegen_read_u32(payload + 8) != expected_channels) { error = "CacheGen record disagrees with its segment header"; -@@ -1998,6 +2002,35 @@ static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & r +@@ -2024,6 +2028,35 @@ static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & r return false; } const size_t max_bytes = expected_rows * sizeof(float); @@ -456,7 +456,7 @@ index 215076924..dad99f62a 100644 if (expected_channels > std::numeric_limits::max() / (cdf_entries * sizeof(uint16_t)) || expected_channels > std::numeric_limits::max() / sizeof(uint16_t)) { error = "CacheGen segment metadata overflows"; -@@ -2014,20 +2047,10 @@ static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & r +@@ -2040,20 +2073,10 @@ static bool skippy_cachegen_validate_segment(const skippy_cachegen_record_v1 & r const size_t cdf_offset = header_bytes + max_bytes; const size_t length_offset = cdf_offset + cdf_bytes; const size_t stream_offset = length_offset + length_bytes; diff --git a/third_party/llama.cpp/patches/0049-ggml-cuda-stage-CacheGen-payloads-directly.patch b/third_party/llama.cpp/patches/0049-ggml-cuda-stage-CacheGen-payloads-directly.patch index 86da779ccd..f38ec0fc63 100644 --- a/third_party/llama.cpp/patches/0049-ggml-cuda-stage-CacheGen-payloads-directly.patch +++ b/third_party/llama.cpp/patches/0049-ggml-cuda-stage-CacheGen-payloads-directly.patch @@ -1,17 +1,18 @@ -From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 -From: scama +From 44481189f0473e079d70bf157d6cd14f9c6c9194 Mon Sep 17 00:00:00 2001 +From: scama + Date: Sat, 12 Sep 2026 08:45:00 +1000 -Subject: [PATCH] ggml-cuda: stage CacheGen payloads directly +Subject: [PATCH 49/50] ggml-cuda: stage CacheGen payloads directly --- - ggml/src/ggml-cuda/ggml-cuda.cu | 36 +++++++++++++++++++++++++----------- + ggml/src/ggml-cuda/ggml-cuda.cu | 36 +++++++++++++++++++++++---------- 1 file changed, 25 insertions(+), 11 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu -index ccfda8a9d..d32175261 100644 +index e72a09eb6..fc9883acc 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu -@@ -5671,6 +5671,12 @@ static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t +@@ -5673,6 +5673,12 @@ static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t #if !defined(GGML_USE_MUSA) @@ -24,7 +25,7 @@ index ccfda8a9d..d32175261 100644 struct ggml_cuda_cachegen_staged_job { int physical_device = -1; uint8_t * dst = nullptr; -@@ -5678,7 +5684,8 @@ struct ggml_cuda_cachegen_staged_job { +@@ -5680,7 +5686,8 @@ struct ggml_cuda_cachegen_staged_job { uint32_t element_bytes = 0; uint64_t token_stride = 0; uint64_t channel_stride = 0; @@ -34,7 +35,7 @@ index ccfda8a9d..d32175261 100644 std::vector tiles; std::vector prefixes; std::vector cells; -@@ -5811,24 +5818,23 @@ static bool ggml_cuda_cachegen_decode( +@@ -5813,24 +5820,23 @@ static bool ggml_cuda_cachegen_decode( return ggml_cuda_cachegen_error(error, error_capacity, "CUDA/HIP CacheGen tile geometry overflows"); } staged.tiles.reserve(job->tile_count); @@ -62,7 +63,7 @@ index ccfda8a9d..d32175261 100644 const uint8_t * payload = static_cast(tile->payload); const bool packed = memcmp(payload, "LCG2", 4) == 0; const uint32_t bins = payload[4]; -@@ -5845,7 +5851,7 @@ static bool ggml_cuda_cachegen_decode( +@@ -5847,7 +5853,7 @@ static bool ggml_cuda_cachegen_decode( error, error_capacity, "CUDA/HIP CacheGen tile header is inconsistent"); } const ggml_cuda_cachegen_tile encoded_tile = { @@ -71,7 +72,7 @@ index ccfda8a9d..d32175261 100644 static_cast(staged.prefixes.size()), tile->token_offset, tile->token_count, -@@ -5893,7 +5899,12 @@ static bool ggml_cuda_cachegen_decode( +@@ -5895,7 +5901,12 @@ static bool ggml_cuda_cachegen_decode( } } staged.tiles.push_back(encoded_tile); @@ -85,7 +86,7 @@ index ccfda8a9d..d32175261 100644 } } } catch (const std::bad_alloc &) { -@@ -5925,7 +5936,7 @@ static bool ggml_cuda_cachegen_decode( +@@ -5927,7 +5938,7 @@ static bool ggml_cuda_cachegen_decode( return fail((operation), status); \ } \ } while (0) @@ -94,7 +95,7 @@ index ccfda8a9d..d32175261 100644 "CUDA/HIP CacheGen payload allocation failed"); GGML_CACHEGEN_CUDA_CALL(cudaMalloc(reinterpret_cast(&job.tiles_device), job.tiles.size() * sizeof(ggml_cuda_cachegen_tile)), -@@ -5936,9 +5947,12 @@ static bool ggml_cuda_cachegen_decode( +@@ -5938,9 +5949,12 @@ static bool ggml_cuda_cachegen_decode( GGML_CACHEGEN_CUDA_CALL(cudaMalloc(reinterpret_cast(&job.cells_device), job.cells.size() * sizeof(uint32_t)), "CUDA/HIP CacheGen cell allocation failed"); @@ -111,5 +112,5 @@ index ccfda8a9d..d32175261 100644 job.tiles.size() * sizeof(ggml_cuda_cachegen_tile), cudaMemcpyHostToDevice, cudaStreamPerThread), -- -2.51.0 +2.54.0 (Apple Git-157) diff --git a/third_party/llama.cpp/patches/0050-ggml-cuda-use-native-CacheGen-shuffle-masks.patch b/third_party/llama.cpp/patches/0050-ggml-cuda-use-native-CacheGen-shuffle-masks.patch index 27b5efd925..0bd8c08b81 100644 --- a/third_party/llama.cpp/patches/0050-ggml-cuda-use-native-CacheGen-shuffle-masks.patch +++ b/third_party/llama.cpp/patches/0050-ggml-cuda-use-native-CacheGen-shuffle-masks.patch @@ -1,7 +1,7 @@ -From 2014cb675af3629ff1bb6d61ef7885bc32b12ed2 Mon Sep 17 00:00:00 2001 +From 000f422ab67daeab5beade144d13b574d15aa04e Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Sat, 12 Sep 2026 11:49:37 +1000 -Subject: [PATCH] ggml-cuda: use native CacheGen shuffle masks +Subject: [PATCH 50/50] ggml-cuda: use native CacheGen shuffle masks --- ggml/src/ggml-cuda/cachegen.cu | 18 +++++++++++++----- From 52cf804e76fb4616fe2cf74c7c59634dbbae0b25 Mon Sep 17 00:00:00 2001 From: scama Date: Tue, 15 Sep 2026 10:05:52 +1000 Subject: [PATCH 12/16] feat(skippy): remove KV cache policy presets --- crates/mesh-llm-config/src/model.rs | 2 - .../control_behavior/model_fit.rs | 1 - .../src/model/built_in_schema/declarations.rs | 4 - .../src/model/built_in_schema/presentation.rs | 12 +-- .../src/model/built_in_schema/validation.rs | 1 - crates/mesh-llm-config/src/model/profile.rs | 1 - .../src/model/profile/defaults.rs | 4 - .../mesh-llm-config/src/model_validation.rs | 5 -- crates/mesh-llm-config/src/wiring_status.rs | 7 -- .../inference/skippy/resolver/resolution.rs | 85 ++++--------------- .../src/inference/skippy/resolver/support.rs | 32 ------- .../src/inference/skippy/resolver/tests.rs | 59 +++++-------- .../src/inference/skippy/resolver/types.rs | 1 - .../src/plugin/config/tests.rs | 2 - .../config_schema_defaults_ui_reference.json | 7 -- .../fixtures/config_schema_reference.json | 2 +- .../fixtures/skippy_full_surface_valid.toml | 1 - .../e2e/configuration/schema-controls.spec.ts | 13 +-- .../configuration-defaults-runtime.ts | 21 ----- .../app-tabs/configuration-defaults.test.ts | 4 +- .../app-tabs/configuration-defaults.ts | 2 +- .../src/features/app-tabs/types.ts | 2 - .../api/config-adapter-diagnostics.test.ts | 8 +- .../configuration/api/config-adapter-merge.ts | 5 -- .../api/config-adapter-schema-placement.ts | 3 - .../api/config-adapter-schema.test.ts | 12 --- .../api/config-adapter-schema.ts | 23 +---- .../api/config-adapter-status.ts | 8 -- .../api/config-adapter-test-support.ts | 22 ----- .../configuration/components/DefaultsTab.tsx | 2 +- .../configuration/components/TomlView.tsx | 5 +- .../settings/SchemaChoiceControl.tsx | 39 +-------- .../settings/schema-control-utils.ts | 1 - .../lib/build-toml-models.test.ts | 4 +- .../features/configuration/lib/build-toml.ts | 15 +--- .../pages/ConfigurationPage-defaults.test.tsx | 23 ----- .../pages/ConfigurationPage-shell.test.tsx | 1 - crates/skippy-cache/src/identity.rs | 20 ++--- docs/CONFIGURATION_PR8_CLOSEOUT_AUDIT.md | 2 +- docs/USAGE.md | 10 +-- docs/skippy/CONFIGURATION.md | 18 ++-- .../model_fit/runtime-single-stage.toml | 1 - docs/skippy/manual-smoke/manifest.tsv | 1 - website/src/docs/pages/config-defaults.md | 5 +- website/src/docs/pages/config-reference.md | 3 +- 45 files changed, 77 insertions(+), 422 deletions(-) diff --git a/crates/mesh-llm-config/src/model.rs b/crates/mesh-llm-config/src/model.rs index ce11841429..dbe84d69a6 100644 --- a/crates/mesh-llm-config/src/model.rs +++ b/crates/mesh-llm-config/src/model.rs @@ -400,8 +400,6 @@ pub struct ModelFitConfig { #[serde(default)] pub cache_type_v: Option, #[serde(default)] - pub kv_cache_policy: Option, - #[serde(default)] pub kv_offload: Option, #[serde(default)] pub kv_unified: Option, diff --git a/crates/mesh-llm-config/src/model/built_in_schema/control_behavior/model_fit.rs b/crates/mesh-llm-config/src/model/built_in_schema/control_behavior/model_fit.rs index 17e239a5a1..bcb432111b 100644 --- a/crates/mesh-llm-config/src/model/built_in_schema/control_behavior/model_fit.rs +++ b/crates/mesh-llm-config/src/model/built_in_schema/control_behavior/model_fit.rs @@ -26,7 +26,6 @@ pub(super) fn apply_model_fit_behavior( set_static_options(setting); push_constraint(setting, ConfigConstraint::NonEmpty); } - "kv_cache_policy" => set_static_options(setting), "kv_offload" | "kv_unified" | "prompt_cache" | "context_shift" | "swa_full" | "flash_attention" => set_static_options(setting), "cache_ram_mib" => set_numeric(setting, Some(0.0), None, Some(1.0), Some("MiB")), diff --git a/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs b/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs index 39fe9e5ff5..99189a4112 100644 --- a/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs +++ b/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs @@ -341,10 +341,6 @@ fn model_fit_settings( basic_setting(&format!("{prefix}.ubatch"), ConfigValueSchema::Integer), basic_setting(&format!("{prefix}.cache_type_k"), kv_cache_type_schema()), basic_setting(&format!("{prefix}.cache_type_v"), kv_cache_type_schema()), - basic_setting( - &format!("{prefix}.kv_cache_policy"), - string_enum(["auto", "quality", "balanced", "saver"]), - ), basic_setting(&format!("{prefix}.kv_offload"), bool_or_auto_schema()), basic_setting(&format!("{prefix}.kv_unified"), bool_or_auto_schema()), basic_setting( diff --git a/crates/mesh-llm-config/src/model/built_in_schema/presentation.rs b/crates/mesh-llm-config/src/model/built_in_schema/presentation.rs index ebc926dd2e..1f6c69aa30 100644 --- a/crates/mesh-llm-config/src/model/built_in_schema/presentation.rs +++ b/crates/mesh-llm-config/src/model/built_in_schema/presentation.rs @@ -56,7 +56,7 @@ const PROMPT_CACHE_CATEGORY: CategoryPresentation = CategoryPresentation { const MEMORY_CATEGORY: CategoryPresentation = CategoryPresentation { id: "memory", label: "Memory", - summary: "VRAM accounting and KV cache policy", + summary: "VRAM accounting and KV cache precision", order: 20, }; const SPECULATIVE_CATEGORY: CategoryPresentation = CategoryPresentation { @@ -480,19 +480,11 @@ fn runtime_defaults_presentation(rendered: &str) -> Option ) .placeholder("cuda:0 or CUDA0") .hint("text")), - "defaults.model_fit.kv_cache_policy" => Some(sp( - "KV cache policy", - "Select how aggressively KV cache precision is reduced to fit larger contexts.", - MEMORY_CATEGORY, - 10, - ) - .hint("segmented") - .renderer("kv-cache-policy")), "defaults.hardware.safety_margin_gb" => Some(sp( "Memory / safety margin", "Keep this much GPU memory free before placement fit checks pass.", MEMORY_CATEGORY, - 20, + 10, ) .unit("GB") .hint("range")), diff --git a/crates/mesh-llm-config/src/model/built_in_schema/validation.rs b/crates/mesh-llm-config/src/model/built_in_schema/validation.rs index 780538a835..1f16ded463 100644 --- a/crates/mesh-llm-config/src/model/built_in_schema/validation.rs +++ b/crates/mesh-llm-config/src/model/built_in_schema/validation.rs @@ -62,7 +62,6 @@ fn built_in_schema_marks_curated_defaults_user_visible() { for path in [ "defaults.throughput.threads", "defaults.throughput.parallel", - "defaults.model_fit.kv_cache_policy", "defaults.request_defaults.temperature", "defaults.skippy.binary_stage_transport", "defaults.multimodal.mmproj_offload", diff --git a/crates/mesh-llm-config/src/model/profile.rs b/crates/mesh-llm-config/src/model/profile.rs index 635938dd50..dab447d194 100644 --- a/crates/mesh-llm-config/src/model/profile.rs +++ b/crates/mesh-llm-config/src/model/profile.rs @@ -76,7 +76,6 @@ fn write_effective_fit_profile(buffer: &mut Vec, entry: &ModelConfigEntry) { } fn write_fit_cache_profile(buffer: &mut Vec, fit: &ModelFitConfig) { - write_option!(buffer, "kv_cache_policy", fit.kv_cache_policy); write_option!(buffer, "kv_offload", fit.kv_offload); write_option!(buffer, "kv_unified", fit.kv_unified); write_option!(buffer, "cache_ram_mib", fit.cache_ram_mib); diff --git a/crates/mesh-llm-config/src/model/profile/defaults.rs b/crates/mesh-llm-config/src/model/profile/defaults.rs index 2b7003bf19..c08fffe3ae 100644 --- a/crates/mesh-llm-config/src/model/profile/defaults.rs +++ b/crates/mesh-llm-config/src/model/profile/defaults.rs @@ -53,10 +53,6 @@ fn merge_model_fit(effective: &mut ModelConfigEntry, defaults: &ModelConfigDefau .clone() .or_else(|| default_fit.cache_type_v.clone()); fit.flash_attention = fit.flash_attention.or(default_fit.flash_attention); - fit.kv_cache_policy = fit - .kv_cache_policy - .clone() - .or_else(|| default_fit.kv_cache_policy.clone()); fit.kv_offload = fit.kv_offload.clone().or(default_fit.kv_offload.clone()); fit.kv_unified = fit.kv_unified.clone().or(default_fit.kv_unified.clone()); fit.cache_ram_mib = fit.cache_ram_mib.or(default_fit.cache_ram_mib); diff --git a/crates/mesh-llm-config/src/model_validation.rs b/crates/mesh-llm-config/src/model_validation.rs index e18b504e77..6d6aa98be2 100644 --- a/crates/mesh-llm-config/src/model_validation.rs +++ b/crates/mesh-llm-config/src/model_validation.rs @@ -312,11 +312,6 @@ fn validate_model_fit(config: &ModelFitConfig, base_path: &str) -> DiagnosticRes config.cache_type_v.as_deref(), &format!("{base_path}.cache_type_v"), )?; - validate_optional_enum( - config.kv_cache_policy.as_deref(), - &["auto", "quality", "balanced", "saver"], - &format!("{base_path}.kv_cache_policy"), - )?; validate_bool_or_auto( config.kv_offload.as_ref(), &format!("{base_path}.kv_offload"), diff --git a/crates/mesh-llm-config/src/wiring_status.rs b/crates/mesh-llm-config/src/wiring_status.rs index 1bc5fde721..b098092db5 100644 --- a/crates/mesh-llm-config/src/wiring_status.rs +++ b/crates/mesh-llm-config/src/wiring_status.rs @@ -600,13 +600,6 @@ pub const WIRING_MANIFEST: &[WiringEntry] = &[ reason: "", behavior: WiringBehavior::None, }, - WiringEntry { - path: "model_fit.kv_cache_policy", - status: WiringStatus::Wired, - owner: "n/a", - reason: "", - behavior: WiringBehavior::None, - }, WiringEntry { path: "model_fit.kv_offload", status: WiringStatus::Wired, diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs index 2b0803a806..d7bee820cc 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs @@ -6,12 +6,11 @@ use super::super::KvCachePolicy; use super::request_defaults::resolve_request_defaults; use super::speculative::resolve_speculative_config; use super::support::{ - KvMacroDefaults, ThroughputMacroDefaults, bool_or_auto_value, derive_fit_target_mib, - effective_flash_attention, has_explicit_prefill_controls, kv_macro_defaults, parse_gpu_layers, - parse_kv_offload_string, pick_owned, pick_string, pick_string_owned, pick_value, - reject_unsupported_hardware_controls, reject_unsupported_model_fit_controls, - resolve_bool_or_auto, resolve_field_string, resolve_field_value, resolve_prefix_cache, - throughput_macro_defaults, + ThroughputMacroDefaults, bool_or_auto_value, derive_fit_target_mib, effective_flash_attention, + has_explicit_prefill_controls, parse_gpu_layers, parse_kv_offload_string, pick_owned, + pick_string, pick_string_owned, pick_value, reject_unsupported_hardware_controls, + reject_unsupported_model_fit_controls, resolve_bool_or_auto, resolve_field_string, + resolve_field_value, resolve_prefix_cache, throughput_macro_defaults, }; use super::types::{ BUILTIN_BATCH, BUILTIN_CTX_SIZE, BUILTIN_PARALLEL, BUILTIN_PREFILL_CHUNK_SIZE, @@ -221,7 +220,6 @@ fn resolve_model_fit_config( context: &ResolverContext<'_>, kv_policy: KvCachePolicy, ) -> Result { - let kv = resolve_kv_defaults(context, kv_policy); let throughput = resolve_throughput_defaults(context); let ctx_size = pick_value( @@ -255,9 +253,9 @@ fn resolve_model_fit_config( .and_then(|defaults| defaults.ubatch), BUILTIN_UBATCH, ); - let cache_type_k = resolve_cache_type_k(context, &kv, kv_policy); - let cache_type_v = resolve_cache_type_v(context, &kv, kv_policy); - let kv_offload = resolve_kv_offload(context, &kv); + let cache_type_k = resolve_cache_type_k(context, kv_policy); + let cache_type_v = resolve_cache_type_v(context, kv_policy); + let kv_offload = resolve_kv_offload(context); let kv_offload_resolved = parse_kv_offload_string(&kv_offload); let kv_unified = resolve_kv_unified(context)?; let swa_full = pick_owned( @@ -297,7 +295,6 @@ fn resolve_model_fit_config( ubatch, cache_type_k, cache_type_v, - kv_cache_policy: kv.effective_policy, prefix_cache, l2_max_bytes, kv_cache_codec, @@ -326,39 +323,7 @@ fn resolve_kv_unified(context: &ResolverContext<'_>) -> Result> { ) } -struct KvDefaults { - effective_policy: String, - model_macro: Option, - global_macro: Option, -} - -fn resolve_kv_defaults(context: &ResolverContext<'_>, kv_policy: KvCachePolicy) -> KvDefaults { - let model_policy = context - .model_fit - .and_then(|fit| fit.kv_cache_policy.as_deref()); - let global_policy = context - .global_model_fit - .and_then(|fit| fit.kv_cache_policy.as_deref()); - let effective_policy = model_policy.or(global_policy).unwrap_or_else(|| { - if context.publisher_defaults.is_some() { - "publisher" - } else { - "safe_f16" - } - }); - - KvDefaults { - effective_policy: effective_policy.to_string(), - model_macro: model_policy.map(|policy| kv_macro_defaults(policy, kv_policy)), - global_macro: global_policy.map(|policy| kv_macro_defaults(policy, kv_policy)), - } -} - -fn resolve_cache_type_k( - context: &ResolverContext<'_>, - kv: &KvDefaults, - kv_policy: KvCachePolicy, -) -> String { +fn resolve_cache_type_k(context: &ResolverContext<'_>, kv_policy: KvCachePolicy) -> String { if let Some(explicit) = context .model_fit .and_then(|fit| non_auto_string(fit.cache_type_k.as_deref())) @@ -367,24 +332,16 @@ fn resolve_cache_type_k( } resolve_field_string( None, - kv.model_macro - .as_ref() - .and_then(|defaults| defaults.cache_type_k.as_deref()), + None, context .global_model_fit .and_then(|fit| non_auto_string(fit.cache_type_k.as_deref())), - kv.global_macro - .as_ref() - .and_then(|defaults| defaults.cache_type_k.as_deref()), + None, kv_policy.cache_type_k(), ) } -fn resolve_cache_type_v( - context: &ResolverContext<'_>, - kv: &KvDefaults, - kv_policy: KvCachePolicy, -) -> String { +fn resolve_cache_type_v(context: &ResolverContext<'_>, kv_policy: KvCachePolicy) -> String { if let Some(explicit) = context .model_fit .and_then(|fit| non_auto_string(fit.cache_type_v.as_deref())) @@ -393,15 +350,11 @@ fn resolve_cache_type_v( } resolve_field_string( None, - kv.model_macro - .as_ref() - .and_then(|defaults| defaults.cache_type_v.as_deref()), + None, context .global_model_fit .and_then(|fit| non_auto_string(fit.cache_type_v.as_deref())), - kv.global_macro - .as_ref() - .and_then(|defaults| defaults.cache_type_v.as_deref()), + None, kv_policy.cache_type_v(), ) } @@ -410,7 +363,7 @@ fn non_auto_string(value: Option<&str>) -> Option<&str> { value.filter(|item| !item.eq_ignore_ascii_case("auto")) } -fn resolve_kv_offload(context: &ResolverContext<'_>, kv: &KvDefaults) -> String { +fn resolve_kv_offload(context: &ResolverContext<'_>) -> String { let model_kv_offload = context .model_fit .and_then(|fit| fit.kv_offload.as_ref()) @@ -422,13 +375,9 @@ fn resolve_kv_offload(context: &ResolverContext<'_>, kv: &KvDefaults) -> String resolve_field_string( model_kv_offload.as_deref(), - kv.model_macro - .as_ref() - .and_then(|defaults| defaults.kv_offload.as_deref()), + None, global_kv_offload.as_deref(), - kv.global_macro - .as_ref() - .and_then(|defaults| defaults.kv_offload.as_deref()), + None, "auto", ) } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/support.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/support.rs index 036412fdcb..f428e3df7f 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/support.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/support.rs @@ -1,7 +1,6 @@ use anyhow::{Result, bail}; use skippy_protocol::{FlashAttentionType, StageKvCacheMode, StageKvCachePayload}; -use super::super::KvCachePolicy; use super::types::{ BUILTIN_BATCH, BUILTIN_PARALLEL, BUILTIN_UBATCH, ResolvedStageKvCache, ResolvedStageKvCacheTemplate, @@ -179,37 +178,6 @@ pub(super) fn resolve_prefix_cache( )) } -pub(super) struct KvMacroDefaults { - pub(super) cache_type_k: Option, - pub(super) cache_type_v: Option, - pub(super) kv_offload: Option, -} - -pub(super) fn kv_macro_defaults(policy: &str, kv_policy: KvCachePolicy) -> KvMacroDefaults { - match policy { - "quality" => KvMacroDefaults { - cache_type_k: Some("f16".to_string()), - cache_type_v: Some("f16".to_string()), - kv_offload: Some("auto".to_string()), - }, - "saver" => KvMacroDefaults { - cache_type_k: Some("q8_0".to_string()), - cache_type_v: Some("q8_0".to_string()), - kv_offload: Some("true".to_string()), - }, - "auto" | "balanced" => KvMacroDefaults { - cache_type_k: Some(kv_policy.cache_type_k().to_string()), - cache_type_v: Some(kv_policy.cache_type_v().to_string()), - kv_offload: Some("auto".to_string()), - }, - _ => KvMacroDefaults { - cache_type_k: Some(kv_policy.cache_type_k().to_string()), - cache_type_v: Some(kv_policy.cache_type_v().to_string()), - kv_offload: Some("auto".to_string()), - }, - } -} - pub(super) struct ThroughputMacroDefaults { pub(super) batch: Option, pub(super) ubatch: Option, diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs index f2c7735943..d6060e9f9c 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs @@ -68,8 +68,6 @@ fn publisher_kv_default_is_used_below_explicit_user_override() { .unwrap(); assert_eq!(automatic.model_fit.cache_type_k, "q8_0"); assert_eq!(automatic.model_fit.cache_type_v, "q8_0"); - assert_eq!(automatic.model_fit.kv_cache_policy, "publisher"); - let explicit_config = parse_config( r#" [defaults.model_fit] @@ -410,12 +408,9 @@ mlock = true } #[test] -fn resolver_macro_expands_kv_cache_tuning_profile_and_safety_margin() { +fn resolver_expands_throughput_profile_and_safety_margin() { let mesh_config = parse_config( r#" -[defaults.model_fit] -kv_cache_policy = "saver" - [defaults.hardware] safety_margin_gb = 1.5 @@ -436,10 +431,9 @@ tuning_profile = "throughput" }) .unwrap(); - assert_eq!(resolved.model_fit.kv_cache_policy, "saver"); - assert_eq!(resolved.model_fit.cache_type_k, "q8_0"); - assert_eq!(resolved.model_fit.cache_type_v, "q8_0"); - assert_eq!(resolved.model_fit.kv_offload, "true"); + assert_eq!(resolved.model_fit.cache_type_k, "f16"); + assert_eq!(resolved.model_fit.cache_type_v, "f16"); + assert_eq!(resolved.model_fit.kv_offload, "auto"); assert_eq!(resolved.throughput.tuning_profile, "throughput"); assert_eq!(resolved.model_fit.batch, 1024); assert_eq!(resolved.model_fit.ubatch, 1024); @@ -449,11 +443,10 @@ tuning_profile = "throughput" } #[test] -fn resolver_treats_auto_cache_type_as_policy_selected_cache_type() { +fn resolver_treats_auto_cache_type_as_publisher_or_safe_default() { let mesh_config = parse_config( r#" [defaults.model_fit] -kv_cache_policy = "saver" cache_type_k = "auto" cache_type_v = "auto" "#, @@ -471,9 +464,8 @@ cache_type_v = "auto" }) .unwrap(); - assert_eq!(resolved.model_fit.kv_cache_policy, "saver"); - assert_eq!(resolved.model_fit.cache_type_k, "q8_0"); - assert_eq!(resolved.model_fit.cache_type_v, "q8_0"); + assert_eq!(resolved.model_fit.cache_type_k, "f16"); + assert_eq!(resolved.model_fit.cache_type_v, "f16"); } #[test] @@ -482,7 +474,6 @@ fn resolver_treats_auto_cache_type_case_insensitively() { let mesh_config_upper = parse_config( r#" [defaults.model_fit] -kv_cache_policy = "saver" cache_type_k = "AUTO" cache_type_v = "AUTO" "#, @@ -500,15 +491,13 @@ cache_type_v = "AUTO" }) .unwrap(); - assert_eq!(resolved_upper.model_fit.kv_cache_policy, "saver"); - assert_eq!(resolved_upper.model_fit.cache_type_k, "q8_0"); - assert_eq!(resolved_upper.model_fit.cache_type_v, "q8_0"); + assert_eq!(resolved_upper.model_fit.cache_type_k, "f16"); + assert_eq!(resolved_upper.model_fit.cache_type_v, "f16"); // Test mixed-case "Auto" let mesh_config_mixed = parse_config( r#" [defaults.model_fit] -kv_cache_policy = "saver" cache_type_k = "Auto" cache_type_v = "Auto" "#, @@ -526,15 +515,13 @@ cache_type_v = "Auto" }) .unwrap(); - assert_eq!(resolved_mixed.model_fit.kv_cache_policy, "saver"); - assert_eq!(resolved_mixed.model_fit.cache_type_k, "q8_0"); - assert_eq!(resolved_mixed.model_fit.cache_type_v, "q8_0"); + assert_eq!(resolved_mixed.model_fit.cache_type_k, "f16"); + assert_eq!(resolved_mixed.model_fit.cache_type_v, "f16"); // Test mixed-case "AuTo" let mesh_config_mixed2 = parse_config( r#" [defaults.model_fit] -kv_cache_policy = "saver" cache_type_k = "AuTo" cache_type_v = "AuTo" "#, @@ -552,13 +539,12 @@ cache_type_v = "AuTo" }) .unwrap(); - assert_eq!(resolved_mixed2.model_fit.kv_cache_policy, "saver"); - assert_eq!(resolved_mixed2.model_fit.cache_type_k, "q8_0"); - assert_eq!(resolved_mixed2.model_fit.cache_type_v, "q8_0"); + assert_eq!(resolved_mixed2.model_fit.cache_type_k, "f16"); + assert_eq!(resolved_mixed2.model_fit.cache_type_v, "f16"); } #[test] -fn per_model_kv_macro_beats_global_explicit_cache_fields_unless_model_explicit_exists() { +fn per_model_explicit_cache_fields_beat_global_explicit_cache_fields() { let mesh_config = parse_config( r#" [defaults.model_fit] @@ -570,8 +556,9 @@ kv_offload = false model = "Qwen/Qwen3-0.6B:Q4_K_M" [models.model_fit] -kv_cache_policy = "saver" +cache_type_k = "q8_0" cache_type_v = "q4_0" +kv_offload = true "#, ); @@ -587,7 +574,6 @@ cache_type_v = "q4_0" }) .unwrap(); - assert_eq!(resolved.model_fit.kv_cache_policy, "saver"); assert_eq!(resolved.model_fit.cache_type_k, "q8_0"); assert_eq!(resolved.model_fit.cache_type_v, "q4_0"); assert_eq!(resolved.model_fit.kv_offload, "true"); @@ -1529,11 +1515,11 @@ fn resolve_with_config_and_model_path( } #[test] -fn kv_offload_resolved_reaches_model_load_options_via_kv_cache_policy() { +fn explicit_kv_offload_reaches_model_load_options() { let mesh_config = parse_config( r#" [defaults.model_fit] -kv_cache_policy = "saver" +kv_offload = true "#, ); @@ -1942,14 +1928,11 @@ fn safe_f16_default_remains_f16_for_incompatible_quantized_kv_meta() { } #[test] -fn model_name_does_not_override_generic_saver_macro_without_metadata() { +fn model_name_does_not_override_safe_default_without_metadata() { let mesh_config = parse_config( r#" [[models]] model = "meshllm/inkling-UD-Q2_K_XL-layers" - -[models.model_fit] -kv_cache_policy = "saver" "#, ); let resolved = resolve_skippy_config(SkippyConfigResolveRequest { @@ -1964,8 +1947,8 @@ kv_cache_policy = "saver" }) .unwrap(); - assert_eq!(resolved.model_fit.cache_type_k, "q8_0"); - assert_eq!(resolved.model_fit.cache_type_v, "q8_0"); + assert_eq!(resolved.model_fit.cache_type_k, "f16"); + assert_eq!(resolved.model_fit.cache_type_v, "f16"); } #[test] diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs index f54ad22483..7192ec69df 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs @@ -73,7 +73,6 @@ pub(crate) struct ResolvedModelFitConfig { pub(crate) ubatch: u32, pub(crate) cache_type_k: String, pub(crate) cache_type_v: String, - pub(crate) kv_cache_policy: String, pub(crate) prefix_cache: ResolvedStageKvCache, pub(crate) l2_max_bytes: u64, pub(crate) kv_cache_codec: StageKvCacheCodec, diff --git a/crates/mesh-llm-host-runtime/src/plugin/config/tests.rs b/crates/mesh-llm-host-runtime/src/plugin/config/tests.rs index fb8b9d6103..3bde06eb30 100644 --- a/crates/mesh-llm-host-runtime/src/plugin/config/tests.rs +++ b/crates/mesh-llm-host-runtime/src/plugin/config/tests.rs @@ -1097,7 +1097,6 @@ version = 1 [defaults.model_fit] ctx_size = 4096 -kv_cache_policy = "balanced" [defaults.hardware] gpu_layers = 10 @@ -1150,7 +1149,6 @@ assignment = "pinned" ctx_size = 8192 batch = 512 ubatch = 128 -kv_cache_policy = "auto" cache_type_k = "auto" cache_type_v = "auto" kv_offload = "auto" diff --git a/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_defaults_ui_reference.json b/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_defaults_ui_reference.json index cbab8a580b..be2298e258 100644 --- a/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_defaults_ui_reference.json +++ b/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_defaults_ui_reference.json @@ -147,13 +147,6 @@ "kind": "built_in" } }, - { - "canonical_path": "defaults.model_fit.kv_cache_policy", - "support": "supported", - "source": { - "kind": "built_in" - } - }, { "canonical_path": "defaults.model_fit.kv_offload", "support": "supported", diff --git a/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_reference.json b/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_reference.json index 21b833b02e..f7763f2940 100644 --- a/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_reference.json +++ b/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_reference.json @@ -216,7 +216,7 @@ "help": "Set the default prefill batch size.", "category_id": "memory", "category_label": "Memory", - "category_summary": "VRAM accounting and KV cache policy", + "category_summary": "VRAM accounting and KV cache precision", "category_order": 20, "setting_order": 40, "unit": "tokens", diff --git a/crates/mesh-llm-host-runtime/tests/fixtures/skippy_full_surface_valid.toml b/crates/mesh-llm-host-runtime/tests/fixtures/skippy_full_surface_valid.toml index a5e008ba5f..5703ed1a97 100644 --- a/crates/mesh-llm-host-runtime/tests/fixtures/skippy_full_surface_valid.toml +++ b/crates/mesh-llm-host-runtime/tests/fixtures/skippy_full_surface_valid.toml @@ -14,7 +14,6 @@ batch = 512 ubatch = 128 cache_type_k = "auto" cache_type_v = "auto" -kv_cache_policy = "balanced" kv_offload = "auto" kv_unified = "auto" prompt_cache = true diff --git a/crates/mesh-llm-ui/e2e/configuration/schema-controls.spec.ts b/crates/mesh-llm-ui/e2e/configuration/schema-controls.spec.ts index a17e8b3f74..f0183597f3 100644 --- a/crates/mesh-llm-ui/e2e/configuration/schema-controls.spec.ts +++ b/crates/mesh-llm-ui/e2e/configuration/schema-controls.spec.ts @@ -86,14 +86,6 @@ const schemaPayload = { value_schema: { kind: 'integer' }, control_behavior: { numeric: { min: 512, max: 32768, step: 512, unit: 'tokens' } } }), - setting('defaults.model_fit.kv_cache_policy', { - label: 'KV cache policy', - help: 'Select the KV cache profile used for memory planning.', - category_id: 'memory', - category_label: 'Memory', - renderer_id: 'kv-cache-policy', - value_schema: { kind: 'enum', values: ['auto', 'quality', 'balanced', 'saver'] } - }), setting('defaults.hardware.device', { label: 'Pinned GPU device', help: 'Only editable when GPU assignment is pinned.', @@ -245,7 +237,7 @@ const controlStatePayload = { const initialConfig = { defaults: { - model_fit: { ctx_size: 4096, kv_cache_policy: 'balanced' }, + model_fit: { ctx_size: 4096 }, hardware: { device: 'cuda:0' }, speculative: { mode: 'disabled', draft_min_tokens: 4 }, multimodal: { mmproj_path: './existing/mmproj.gguf', mmproj_url: 'https://example.com/mmproj.gguf' } @@ -582,7 +574,6 @@ test.describe('schema-driven configuration controls', () => { const contextSize = page.getByRole('slider', { name: 'Context size' }) await page.getByRole('button', { name: '8K', exact: true }).click() await expect(contextSize).toHaveAttribute('aria-valuenow', '8192') - await page.getByLabel('KV cache policy').getByRole('radio', { name: 'quality' }).click() await page.getByLabel('GPU assignment').getByRole('radio', { name: 'pinned' }).click() await expect(page.getByRole('textbox', { name: 'Pinned GPU device' })).toBeEnabled() await page.getByRole('textbox', { name: 'Pinned GPU device' }).fill('cuda:1') @@ -590,7 +581,6 @@ test.describe('schema-driven configuration controls', () => { await page.getByRole('tab', { name: 'TOML Output' }).click() const toml = page.getByRole('textbox', { name: 'Configuration TOML source' }) await expect(toml).toHaveValue(/ctx_size = 8192/) - await expect(toml).toHaveValue(/kv_cache_policy = "quality"/) await expect(toml).toHaveValue(/assignment = "pinned"/) await expect(toml).toHaveValue(/gpu_id = "cuda:1"/) await expect(page.getByText('Generated TOML validates against mesh-llm config rules.')).toBeVisible() @@ -618,7 +608,6 @@ test.describe('schema-driven configuration controls', () => { expect(applyRequest.endpoint).toBe('local-owner') expect(applyRequest.expected_revision).toBe(7) expect(applyRequest.config.defaults?.model_fit?.ctx_size).toBe(8192) - expect(applyRequest.config.defaults?.model_fit?.kv_cache_policy).toBe('quality') expect(applyRequest.config.gpu?.assignment).toBe('pinned') expect(applyRequest.config.defaults?.hardware?.device).toBe('cuda:1') expect(applyRequest.config.plugin?.[0]?.settings?.endpoint).toBe('https://blackboard.local/api') diff --git a/crates/mesh-llm-ui/src/features/app-tabs/configuration-defaults-runtime.ts b/crates/mesh-llm-ui/src/features/app-tabs/configuration-defaults-runtime.ts index 0dfe146543..cae3afcacb 100644 --- a/crates/mesh-llm-ui/src/features/app-tabs/configuration-defaults-runtime.ts +++ b/crates/mesh-llm-ui/src/features/app-tabs/configuration-defaults-runtime.ts @@ -352,27 +352,6 @@ export const CONFIGURATION_DEFAULT_RUNTIME_SETTINGS = [ inheritedLabel: 'Used only by placements without a model-specific hardware.device', control: { kind: 'text', name: 'device', value: '', placeholder: 'cuda:0 or CUDA0' } }, - { - id: 'kv-cache', - categoryId: 'memory', - tomlSection: 'defaults.model_fit', - tomlKey: 'kv_cache_policy', - icon: 'filter', - label: 'KV cache policy', - description: 'Select how aggressively KV cache precision is reduced to fit larger contexts.', - inheritedLabel: 'Used when the placement has no cache override', - control: { - kind: 'choice', - name: 'kv_cache_policy', - value: 'auto', - options: [ - { value: 'auto', label: 'auto' }, - { value: 'quality', label: 'quality' }, - { value: 'balanced', label: 'balanced' }, - { value: 'saver', label: 'saver' } - ] - } - }, { id: 'memory-margin', categoryId: 'memory', diff --git a/crates/mesh-llm-ui/src/features/app-tabs/configuration-defaults.test.ts b/crates/mesh-llm-ui/src/features/app-tabs/configuration-defaults.test.ts index a60b84fb28..a54852f202 100644 --- a/crates/mesh-llm-ui/src/features/app-tabs/configuration-defaults.test.ts +++ b/crates/mesh-llm-ui/src/features/app-tabs/configuration-defaults.test.ts @@ -59,7 +59,7 @@ describe('CONFIGURATION_DEFAULTS', () => { { id: 'memory', label: 'Memory', - summary: 'KV cache policy and fit headroom.', + summary: 'KV cache precision and fit headroom.', help: 'VRAM accounting and fit headroom' }, { @@ -222,7 +222,7 @@ describe('CONFIGURATION_DEFAULTS', () => { it('keeps all settings keyed to canonical TOML sections without duplicate ids', () => { const settingIds = CONFIGURATION_DEFAULTS.settings.map((setting) => setting.id) - expect(CONFIGURATION_DEFAULTS.settings.length).toBeGreaterThanOrEqual(74) + expect(CONFIGURATION_DEFAULTS.settings.length).toBeGreaterThanOrEqual(73) expect(new Set(settingIds).size).toBe(settingIds.length) expect(CONFIGURATION_DEFAULTS.settings.every((setting) => setting.tomlSection.startsWith('defaults.'))).toBe(true) }) diff --git a/crates/mesh-llm-ui/src/features/app-tabs/configuration-defaults.ts b/crates/mesh-llm-ui/src/features/app-tabs/configuration-defaults.ts index b980aa2bb1..c502cb1751 100644 --- a/crates/mesh-llm-ui/src/features/app-tabs/configuration-defaults.ts +++ b/crates/mesh-llm-ui/src/features/app-tabs/configuration-defaults.ts @@ -72,7 +72,7 @@ export const CONFIGURATION_DEFAULTS = { { id: 'memory', label: 'Memory', - summary: 'KV cache policy and fit headroom.', + summary: 'KV cache precision and fit headroom.', help: 'VRAM accounting and fit headroom' }, { diff --git a/crates/mesh-llm-ui/src/features/app-tabs/types.ts b/crates/mesh-llm-ui/src/features/app-tabs/types.ts index ac25b9a09b..e763a5a205 100644 --- a/crates/mesh-llm-ui/src/features/app-tabs/types.ts +++ b/crates/mesh-llm-ui/src/features/app-tabs/types.ts @@ -263,7 +263,6 @@ export type ConfigAssignModelConfig = { flashAttention?: 'auto' | 'enabled' | 'disabled' cacheTypeK?: string cacheTypeV?: string - kvCachePolicy?: 'auto' | 'quality' | 'balanced' | 'saver' } export type ConfigAssign = { id: string @@ -491,7 +490,6 @@ export type ConfigurationModelPlacementPaths = { gpuLayers: string cacheTypeK?: string cacheTypeV?: string - kvCachePolicy?: string flashAttention?: string mmproj?: string } diff --git a/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-diagnostics.test.ts b/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-diagnostics.test.ts index ea89317db3..fd683540ca 100644 --- a/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-diagnostics.test.ts +++ b/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-diagnostics.test.ts @@ -20,8 +20,7 @@ describe('configuration model merge and diagnostics', () => { model_fit: { ctx_size: 2048, cache_type_k: 'q8_0', - cache_type_v: 'q4_0', - kv_cache_policy: 'balanced' + cache_type_v: 'q4_0' }, hardware: { device: 'cuda:0', gpu_layers: -1 }, keep: 'first' @@ -64,8 +63,7 @@ describe('configuration model merge and diagnostics', () => { draftModelPath: '/models/draft.gguf', flashAttention: 'enabled', cacheTypeK: 'q8_0', - cacheTypeV: 'q5_1', - kvCachePolicy: 'balanced' + cacheTypeV: 'q5_1' } }, { id: 'assign-2', modelId: 'hf://meshllm/dupe@main:Q4_K_M', nodeId: 'self', containerIdx: 1, ctx: 16384 } @@ -85,7 +83,6 @@ describe('configuration model merge and diagnostics', () => { ubatch: 128, cache_type_k: 'q8_0', cache_type_v: 'q5_1', - kv_cache_policy: 'balanced', flash_attention: 'enabled' }, hardware: { @@ -331,7 +328,6 @@ describe('configuration model merge and diagnostics', () => { 'defaults.throughput.parallel', 'defaults.hardware.safety_margin_gb', 'defaults.model_fit.ctx_size', - 'defaults.model_fit.kv_cache_policy', 'defaults.request_defaults.temperature', 'defaults.request_defaults.reasoning_enabled' ]) diff --git a/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-merge.ts b/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-merge.ts index 09807d4e9f..277406b5c6 100644 --- a/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-merge.ts +++ b/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-merge.ts @@ -590,11 +590,6 @@ function writeSelectedModelConfig( placementPaths.cacheTypeV ?? DEFAULT_MODEL_PLACEMENT_PATHS.cacheTypeV!, config?.cacheTypeV ) - writeOptionalModelEntryPath( - entry, - placementPaths.kvCachePolicy ?? DEFAULT_MODEL_PLACEMENT_PATHS.kvCachePolicy!, - config?.kvCachePolicy - ) } export function mergeConfigurationIntoMeshConfig( diff --git a/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-schema-placement.ts b/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-schema-placement.ts index d08b361076..6376b3dd60 100644 --- a/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-schema-placement.ts +++ b/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-schema-placement.ts @@ -3,7 +3,6 @@ import type { RuntimeConfigSchemaEntry, RuntimeConfigSchemaReference } from './c const PATH_RENDERER_FALLBACKS: Record = { 'defaults.throughput.parallel': 'slot-meter', - 'defaults.model_fit.kv_cache_policy': 'kv-cache-policy', 'defaults.model_fit.ctx_size': 'context-slider' } @@ -18,7 +17,6 @@ export const DEFAULT_MODEL_PLACEMENT_PATHS: ConfigurationModelPlacementPaths = { gpuLayers: 'models..hardware.gpu_layers', cacheTypeK: 'models..model_fit.cache_type_k', cacheTypeV: 'models..model_fit.cache_type_v', - kvCachePolicy: 'models..model_fit.kv_cache_policy', flashAttention: 'models..model_fit.flash_attention', mmproj: 'models..multimodal.mmproj' } @@ -42,7 +40,6 @@ export function modelPlacementPathsFromSchema( gpuLayers: pathByRenderer.get('model-placement-gpu-layers') ?? DEFAULT_MODEL_PLACEMENT_PATHS.gpuLayers, cacheTypeK: DEFAULT_MODEL_PLACEMENT_PATHS.cacheTypeK, cacheTypeV: DEFAULT_MODEL_PLACEMENT_PATHS.cacheTypeV, - kvCachePolicy: DEFAULT_MODEL_PLACEMENT_PATHS.kvCachePolicy, flashAttention: DEFAULT_MODEL_PLACEMENT_PATHS.flashAttention, mmproj: DEFAULT_MODEL_PLACEMENT_PATHS.mmproj } diff --git a/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-schema.test.ts b/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-schema.test.ts index dd4b3ffbbd..59e2a587bd 100644 --- a/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-schema.test.ts +++ b/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-schema.test.ts @@ -210,9 +210,6 @@ describe('configuration schema and status adaptation', () => { hardware: { safety_margin_gb: 3.5 }, - model_fit: { - kv_cache_policy: 'quality' - }, request_defaults: { temperature: 0.8, reasoning_enabled: false @@ -229,7 +226,6 @@ describe('configuration schema and status adaptation', () => { expect(values['defaults.throughput.parallel']).toBe('8') expect(values['defaults.hardware.safety_margin_gb']).toBe('3.5') - expect(values['defaults.model_fit.kv_cache_policy']).toBe('quality') expect(values['defaults.request_defaults.temperature']).toBe('0.8') expect(values['defaults.request_defaults.reasoning_enabled']).toBe('off') }) @@ -240,7 +236,6 @@ describe('configuration schema and status adaptation', () => { const reasoningEnabled = defaults.settings.find( (setting) => setting.id === 'defaults.request_defaults.reasoning_enabled' ) - const kvCache = defaults.settings.find((setting) => setting.id === 'defaults.model_fit.kv_cache_policy') const ctxSize = defaults.settings.find((setting) => setting.id === 'defaults.model_fit.ctx_size') expect(temperature).toMatchObject({ @@ -263,13 +258,6 @@ describe('configuration schema and status adaptation', () => { ] }) }) - expect(kvCache).toMatchObject({ - rendererId: 'kv-cache-policy', - control: expect.objectContaining({ - kind: 'choice', - options: expect.arrayContaining([{ value: 'quality', label: 'quality' }]) - }) - }) expect(ctxSize).toMatchObject({ rendererId: 'context-slider', control: expect.objectContaining({ diff --git a/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-schema.ts b/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-schema.ts index 88c4075bec..af4f24b625 100644 --- a/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-schema.ts +++ b/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-schema.ts @@ -169,7 +169,7 @@ const DEFAULTS_CATEGORY_FALLBACKS: Record memory: { id: 'memory', label: 'Memory', - summary: 'VRAM accounting and KV cache policy', + summary: 'VRAM accounting and KV cache precision', help: 'Memory defaults inherited by model placements', tomlSection: 'defaults.model_fit', order: 20 @@ -224,8 +224,6 @@ const DEFAULTS_CATEGORY_FALLBACKS: Record } } -type ChoicePresentation = Extract['presentation'] - function settingIdFromPath(canonicalPath: string) { return canonicalPath } @@ -293,21 +291,6 @@ function controlNameForPath(canonicalPath: string) { return lastPathSegment(canonicalPath) } -function segmentedControl( - name: string, - value: string, - options: readonly string[], - presentation: ChoicePresentation = 'segmented' -): ConfigurationDefaultsControl { - return { - kind: 'choice', - name, - value, - presentation, - options: options.map((option) => ({ value: option, label: option })) - } -} - function bespokeControlForRenderer(entry: RuntimeConfigSchemaEntry): ConfigurationDefaultsControl | undefined { const rendererId = rendererIdForEntry(entry) const name = controlNameForPath(entry.canonical_path) @@ -328,10 +311,6 @@ function bespokeControlForRenderer(entry: RuntimeConfigSchemaEntry): Configurati } } - if (rendererId === 'kv-cache-policy') { - return segmentedControl(name, 'auto', ['auto', 'quality', 'balanced', 'saver']) - } - return undefined } diff --git a/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-status.ts b/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-status.ts index 5a627c9c59..ee99bc1fc2 100644 --- a/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-status.ts +++ b/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-status.ts @@ -249,14 +249,6 @@ function modelConfigFromEntry( placementPaths.cacheTypeV ?? DEFAULT_MODEL_PLACEMENT_PATHS.cacheTypeV! ) - const kvCachePolicy = stringModelEntryValue( - entry, - placementPaths.kvCachePolicy ?? DEFAULT_MODEL_PLACEMENT_PATHS.kvCachePolicy! - ) - if (kvCachePolicy === 'quality' || kvCachePolicy === 'balanced' || kvCachePolicy === 'saver') { - config.kvCachePolicy = kvCachePolicy - } - return Object.keys(config).length ? config : undefined } diff --git a/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-test-support.ts b/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-test-support.ts index 74332557ec..b63d81414f 100644 --- a/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-test-support.ts +++ b/crates/mesh-llm-ui/src/features/configuration/api/config-adapter-test-support.ts @@ -212,28 +212,6 @@ export const SCHEMA_REFERENCE: RuntimeConfigSchemaReference = { renderer_id: 'context-slider' } }, - { - canonical_path: 'defaults.model_fit.kv_cache_policy', - owner: 'built_in', - source: { kind: 'built_in' }, - value_schema: { kind: 'string' }, - support: 'supported', - control_surfaces: ['config_file'], - apply_mode: 'static_on_load', - restart_scope: 'model_reload', - visibility: 'user', - presentation: { - label: 'KV cache policy', - help: 'Select KV cache policy.', - category_id: 'memory', - category_label: 'Memory', - category_summary: 'Memory defaults', - category_order: 20, - setting_order: 20, - control_hint: 'segmented', - renderer_id: 'kv-cache-policy' - } - }, { canonical_path: 'defaults.request_defaults.temperature', owner: 'built_in', diff --git a/crates/mesh-llm-ui/src/features/configuration/components/DefaultsTab.tsx b/crates/mesh-llm-ui/src/features/configuration/components/DefaultsTab.tsx index c7fa73c77f..7deead7018 100644 --- a/crates/mesh-llm-ui/src/features/configuration/components/DefaultsTab.tsx +++ b/crates/mesh-llm-ui/src/features/configuration/components/DefaultsTab.tsx @@ -295,7 +295,7 @@ function sectionSubtitle(category: ConfigurationDefaultsCategory) { if (category.id === 'runtime-policy') return 'Runtime reconciliation behavior' if (category.id === 'network') return 'Owner-control listener settings' if (category.id === 'attestation') return 'Certified-build admission requirements' - if (category.id === 'memory') return 'VRAM accounting and KV cache policy' + if (category.id === 'memory') return 'VRAM accounting and KV cache precision' if (category.id === 'speculative-decoding') return 'Speculative draft policy defaults' if (category.id === 'request-defaults') return 'Request-time sampling and reasoning defaults' return category.help diff --git a/crates/mesh-llm-ui/src/features/configuration/components/TomlView.tsx b/crates/mesh-llm-ui/src/features/configuration/components/TomlView.tsx index 48dc4475fb..0eabb5ac03 100644 --- a/crates/mesh-llm-ui/src/features/configuration/components/TomlView.tsx +++ b/crates/mesh-llm-ui/src/features/configuration/components/TomlView.tsx @@ -412,7 +412,8 @@ function LaunchSummaryPanel({ const localNode = nodes[0] const gpuCount = localNode?.gpus.length ?? 0 const flashAttention = defaultsValues?.['defaults.model_fit.flash_attention'] ?? 'auto' - const kvCache = defaultsValues?.['defaults.model_fit.kv_cache_policy'] ?? 'auto' + const cacheTypeK = defaultsValues?.['defaults.model_fit.cache_type_k'] ?? 'publisher/default' + const cacheTypeV = defaultsValues?.['defaults.model_fit.cache_type_v'] ?? 'publisher/default' const httpBind = launchSummaryConfig?.httpBind ?? '0.0.0.0:9337' const mmap = launchSummaryConfig?.mmap ?? 'off' const rows = [ @@ -420,7 +421,7 @@ function LaunchSummaryPanel({ ['placements:', `${assigns.length} models on ${gpuCount} GPUs`], ['http:', httpBind], ['flash attn:', flashAttention], - ['kv cache:', `${kvCache} (q8_0/q4_0 above 5GB)`], + ['kv cache:', `K ${cacheTypeK} · V ${cacheTypeV}`], ['mmap:', mmap] ] diff --git a/crates/mesh-llm-ui/src/features/configuration/components/settings/SchemaChoiceControl.tsx b/crates/mesh-llm-ui/src/features/configuration/components/settings/SchemaChoiceControl.tsx index 8a9a30082a..af3eb4b310 100644 --- a/crates/mesh-llm-ui/src/features/configuration/components/settings/SchemaChoiceControl.tsx +++ b/crates/mesh-llm-ui/src/features/configuration/components/settings/SchemaChoiceControl.tsx @@ -2,7 +2,6 @@ import { NativeSelect } from '@/components/ui/NativeSelect' import { SegmentedControl } from '@/components/ui/SegmentedControl' import { cn } from '@/lib/cn' import { - effectiveRendererId, resolvedChoiceOptions, type SchemaSettingControlProps } from '@/features/configuration/components/settings/schema-control-utils' @@ -12,7 +11,6 @@ function choiceItemClassName(setting: SchemaSettingControlProps['setting']) { 'min-w-[64px] capitalize', setting.control.kind === 'choice' && setting.control.presentation === 'toggle' && 'min-w-[38px]', setting.canonicalPath?.endsWith('.flash_attention') && 'min-w-[38px]', - effectiveRendererId(setting) === 'kv-cache-policy' && 'min-w-[58px]', setting.canonicalPath === 'defaults.speculative.mode' && 'min-w-[72px]', setting.canonicalPath === 'defaults.speculative.draft_selection_policy' && 'min-w-[86px]', setting.canonicalPath === 'defaults.speculative.pairing_fault' && 'min-w-[104px]', @@ -20,37 +18,6 @@ function choiceItemClassName(setting: SchemaSettingControlProps['setting']) { ) } -function KvPolicyMatrix({ policy }: { readonly policy: string }) { - const rows = [ - { label: '<5GB', detail: 'K F16 · V F16', active: policy === 'auto' || policy === 'quality' }, - { label: '5–50GB', detail: 'K q8_0 · V q4_0', active: policy === 'auto' || policy === 'balanced' }, - { label: '≥50GB', detail: 'K q4_0 · V q4_0', active: policy === 'auto' || policy === 'saver' } - ] - - return ( -
- {rows.map((row) => ( - - {row.label} - {row.detail} - - ))} -
- ) -} - export function SchemaChoiceControl({ ariaDescribedBy, disabled = false, @@ -66,10 +33,7 @@ export function SchemaChoiceControl({ return (
{presentation === 'select' ? ( @@ -97,7 +61,6 @@ export function SchemaChoiceControl({ variant="pill" /> )} - {effectiveRendererId(setting) === 'kv-cache-policy' ? : null} {selectedDescription ? (

{selectedDescription} diff --git a/crates/mesh-llm-ui/src/features/configuration/components/settings/schema-control-utils.ts b/crates/mesh-llm-ui/src/features/configuration/components/settings/schema-control-utils.ts index 2fa49efa99..d5bba4a979 100644 --- a/crates/mesh-llm-ui/src/features/configuration/components/settings/schema-control-utils.ts +++ b/crates/mesh-llm-ui/src/features/configuration/components/settings/schema-control-utils.ts @@ -98,7 +98,6 @@ export function hasSchemaKind( export function effectiveRendererId(setting: ConfigurationDefaultsSetting) { if (setting.rendererId) return setting.rendererId if (setting.id === 'parallel-slots') return 'slot-meter' - if (setting.id === 'kv-cache') return 'kv-cache-policy' if (setting.id === 'ctx-size') return 'context-slider' return undefined } diff --git a/crates/mesh-llm-ui/src/features/configuration/lib/build-toml-models.test.ts b/crates/mesh-llm-ui/src/features/configuration/lib/build-toml-models.test.ts index 5f157e16dc..6013f83a61 100644 --- a/crates/mesh-llm-ui/src/features/configuration/lib/build-toml-models.test.ts +++ b/crates/mesh-llm-ui/src/features/configuration/lib/build-toml-models.test.ts @@ -167,8 +167,7 @@ describe('buildTOML model and plugin serialization', () => { model_fit: { ctx_size: 2048, cache_type_k: 'q8_0', - cache_type_v: 'q4_0', - kv_cache_policy: 'balanced' + cache_type_v: 'q4_0' } }, { @@ -191,7 +190,6 @@ describe('buildTOML model and plugin serialization', () => { expect(toml.match(/\[\[models\]\]/g)).toHaveLength(3) expect(toml).toContain('ctx_size = 131072\ncache_type_k = "q8_0"\ncache_type_v = "q4_0"') - expect(toml).toContain('[models.model_fit]\nkv_cache_policy = "balanced"') expect(toml).toContain('ctx_size = 262144\ncache_type_k = "f16"\ncache_type_v = "f16"') expect(toml).toContain('ctx_size = 65536\ncache_type_v = "q8_0"') expect(toml).not.toContain('[models.model_fit]\nctx_size') diff --git a/crates/mesh-llm-ui/src/features/configuration/lib/build-toml.ts b/crates/mesh-llm-ui/src/features/configuration/lib/build-toml.ts index cbcd0975cc..87517ca3ee 100644 --- a/crates/mesh-llm-ui/src/features/configuration/lib/build-toml.ts +++ b/crates/mesh-llm-ui/src/features/configuration/lib/build-toml.ts @@ -35,8 +35,7 @@ const DEFAULT_MODEL_PLACEMENT_PATHS: ConfigurationModelPlacementPaths = { device: 'models..hardware.device', gpuLayers: 'models..hardware.gpu_layers', cacheTypeK: 'models..model_fit.cache_type_k', - cacheTypeV: 'models..model_fit.cache_type_v', - kvCachePolicy: 'models..model_fit.kv_cache_policy' + cacheTypeV: 'models..model_fit.cache_type_v' } const defaultSectionOrder: readonly ConfigurationTomlSectionId[] = [ @@ -404,8 +403,7 @@ function consumeModelConfigEntry( function modelFitOverridePaths(placementPaths: ConfigurationModelPlacementPaths): Array<{ key: string; path: string }> { return [ { key: 'cache_type_k', path: placementPaths.cacheTypeK ?? DEFAULT_MODEL_PLACEMENT_PATHS.cacheTypeK! }, - { key: 'cache_type_v', path: placementPaths.cacheTypeV ?? DEFAULT_MODEL_PLACEMENT_PATHS.cacheTypeV! }, - { key: 'kv_cache_policy', path: placementPaths.kvCachePolicy ?? DEFAULT_MODEL_PLACEMENT_PATHS.kvCachePolicy! } + { key: 'cache_type_v', path: placementPaths.cacheTypeV ?? DEFAULT_MODEL_PLACEMENT_PATHS.cacheTypeV! } ] } @@ -561,15 +559,6 @@ export function appendSelectedModelConfig( appendModelConfigLine(modelLines, sectionLines, 'models..model_fit.cache_type_v', config.cacheTypeV) emittedKeys.add('cache_type_v') } - if (config.kvCachePolicy) { - appendModelConfigLine( - modelLines, - sectionLines, - 'models..model_fit.kv_cache_policy', - config.kvCachePolicy - ) - emittedKeys.add('kv_cache_policy') - } } export function buildTOML( diff --git a/crates/mesh-llm-ui/src/features/configuration/pages/ConfigurationPage-defaults.test.tsx b/crates/mesh-llm-ui/src/features/configuration/pages/ConfigurationPage-defaults.test.tsx index b36dfebfd4..d5aff354a1 100644 --- a/crates/mesh-llm-ui/src/features/configuration/pages/ConfigurationPage-defaults.test.tsx +++ b/crates/mesh-llm-ui/src/features/configuration/pages/ConfigurationPage-defaults.test.tsx @@ -93,29 +93,6 @@ describe('ConfigurationPage defaults controls', () => { expect(screen.getByText('3.6 GB · 12 × 0.30 GB')).toBeInTheDocument() }) - it('updates KV cache memory tiers from the selected policy', async () => { - const user = userEvent.setup() - - render() - - const policyControl = within(screen.getByRole('radiogroup', { name: 'KV cache policy' })) - const tiers = () => - within(screen.getByRole('group', { name: 'KV cache memory tiers' })) - .getAllByText(/^K /) - .map((node) => node.closest('[data-kv-tier-active]')) - - expect(tiers().map((node) => node?.getAttribute('data-kv-tier-active'))).toEqual(['true', 'true', 'true']) - - await user.click(policyControl.getByRole('radio', { name: 'quality' })) - expect(tiers().map((node) => node?.getAttribute('data-kv-tier-active'))).toEqual(['true', undefined, undefined]) - - await user.click(policyControl.getByRole('radio', { name: 'balanced' })) - expect(tiers().map((node) => node?.getAttribute('data-kv-tier-active'))).toEqual([undefined, 'true', undefined]) - - await user.click(policyControl.getByRole('radio', { name: 'saver' })) - expect(tiers().map((node) => node?.getAttribute('data-kv-tier-active'))).toEqual([undefined, undefined, 'true']) - }) - it('renders speculative decoding defaults and writes them to TOML', async () => { const user = userEvent.setup() diff --git a/crates/mesh-llm-ui/src/features/configuration/pages/ConfigurationPage-shell.test.tsx b/crates/mesh-llm-ui/src/features/configuration/pages/ConfigurationPage-shell.test.tsx index 39d6c30aad..50713bedfd 100644 --- a/crates/mesh-llm-ui/src/features/configuration/pages/ConfigurationPage-shell.test.tsx +++ b/crates/mesh-llm-ui/src/features/configuration/pages/ConfigurationPage-shell.test.tsx @@ -232,7 +232,6 @@ describe('ConfigurationPage shell and feature flags', () => { expect(screen.queryByText('Model Runtime')).not.toBeInTheDocument() expect(screen.getByText('Default GPU device')).toBeInTheDocument() expect(screen.getByText('GPU layers')).toBeInTheDocument() - expect(screen.getByText('KV cache policy')).toBeInTheDocument() expect(screen.getByText('Memory / safety margin')).toBeInTheDocument() expect(screen.getByText('Reasoning format')).toBeInTheDocument() expect(screen.getByText('Temperature')).toBeInTheDocument() diff --git a/crates/skippy-cache/src/identity.rs b/crates/skippy-cache/src/identity.rs index d84d9c20ef..51fc097224 100644 --- a/crates/skippy-cache/src/identity.rs +++ b/crates/skippy-cache/src/identity.rs @@ -63,9 +63,8 @@ fn update_platform_identity(hasher: &mut blake3::Hasher) { /// exported KV page without changing the token sequence. /// /// Identity must cover every input that alters the serialized layout or the -/// numerical content of exported state. Flipping `kv_cache_policy` from -/// `quality` to `saver` rewrites `cache_type_k`/`_v` from `f16` to `q8_0`; -/// without these fields in the hash, incompatible state would share a page id. +/// numerical content of exported state. Changing `cache_type_k`/`_v` from +/// `f16` to `q8_0` must not let incompatible state share a page id. /// /// `NATIVE_KV_DTYPE` is a fixed layout tag and does **not** vary with the /// configured cache types, so it cannot stand in for them. @@ -664,22 +663,21 @@ mod identity_completeness_tests { prefix_hash(config, 0, &[1, 2, 3, 4]) } - /// Changing the KV cache policy rewrites `cache_type_k`/`_v`. These - /// formats must produce distinct page identities to prevent importing - /// cached q8_0 state as f16. + /// Changing `cache_type_k`/`_v` must produce distinct page identities to + /// prevent importing cached q8_0 state as f16. #[test] fn kv_cache_dtype_changes_page_identity() { - let quality = test_config(); - let saver = StageConfig { + let f16 = test_config(); + let quantized = StageConfig { cache_type_k: "q8_0".to_string(), cache_type_v: "q8_0".to_string(), ..test_config() }; - assert_ne!(hash_of(&quality), hash_of(&saver)); + assert_ne!(hash_of(&f16), hash_of(&quantized)); assert_ne!( - prefix_identity(&quality, 0, &[1, 2, 3, 4]).page_id, - prefix_identity(&saver, 0, &[1, 2, 3, 4]).page_id + prefix_identity(&f16, 0, &[1, 2, 3, 4]).page_id, + prefix_identity(&quantized, 0, &[1, 2, 3, 4]).page_id ); } diff --git a/docs/CONFIGURATION_PR8_CLOSEOUT_AUDIT.md b/docs/CONFIGURATION_PR8_CLOSEOUT_AUDIT.md index 99727ab065..c8731d90c8 100644 --- a/docs/CONFIGURATION_PR8_CLOSEOUT_AUDIT.md +++ b/docs/CONFIGURATION_PR8_CLOSEOUT_AUDIT.md @@ -61,7 +61,7 @@ schema row or stale manifest row from passing review: `runtime.model_target_demand_upgrade_max_age_secs`, `advanced.server.alias`, `model`, `hardware.model_path`, `hardware.hf_repo`, `hardware.hf_file`, `model_fit.ctx_size`, `model_fit.batch`, `model_fit.ubatch`, -`model_fit.cache_type_k`, `model_fit.cache_type_v`, `model_fit.kv_cache_policy`, +`model_fit.cache_type_k`, `model_fit.cache_type_v`, `model_fit.kv_offload`, `model_fit.kv_unified`, `model_fit.cache_ram_mib`, `model_fit.cache_idle_slots`, `model_fit.prompt_cache`, `model_fit.prefix_cache.enabled`, `model_fit.prefix_cache.max_entries`, diff --git a/docs/USAGE.md b/docs/USAGE.md index e13044376c..e53f3bb7b0 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -461,11 +461,6 @@ ubatch = 128 # n_ubatch — micro-batch within a batch cache_type_k = "auto" # KV key dtype: auto f16 f32 bf16 q8_0 q4_0 … cache_type_v = "auto" # KV value dtype (same enum) flash_attention = "auto" # auto on off -kv_cache_policy = "balanced" # macro preset: auto quality balanced saver - # quality → f16/f16, no forced RAM cap - # balanced → preserve runtime defaults - # saver → low-memory dtypes + offload - # explicit cache_type_k/v always wins over preset kv_offload = "auto" # bool or "auto" — KV residency / offload policy kv_unified = "auto" # bool or "auto" — unified KV layout (schema-reserved) cache_ram_mib = 0 # host-RAM L2 budget in MiB; 0 = disabled; requires L3 @@ -720,7 +715,6 @@ batch = 1024 ubatch = 256 cache_type_k = "f16" cache_type_v = "f16" -kv_cache_policy = "quality" # overrides global "balanced" flash_attention = "on" prompt_cache = true @@ -831,7 +825,9 @@ fit_target_mib = 20480 [models.model_fit] ctx_size = 8192 -kv_cache_policy = "saver" +cache_type_k = "q8_0" +cache_type_v = "q8_0" +kv_offload = true [models.throughput] parallel = 2 diff --git a/docs/skippy/CONFIGURATION.md b/docs/skippy/CONFIGURATION.md index d967f75e09..85fc476e20 100644 --- a/docs/skippy/CONFIGURATION.md +++ b/docs/skippy/CONFIGURATION.md @@ -38,21 +38,14 @@ protobuf, or lower runtime layers. ranges only execute in staged mode. - Unsupported or deferred rows are intentionally marked `documented-rejected` or excluded from the matrix. Logprobs, backend sampling bags, adaptive sampling, and operational or security exclusions live in that `documented-rejected` bucket. -## Deterministic macro mappings +## Deterministic throughput macro mappings -- `model_fit.kv_cache_policy=quality` prefers `cache_type_k=f16`, - `cache_type_v=f16`, `kv_offload=auto`, and no forced cache RAM cap. -- `model_fit.kv_cache_policy=balanced` preserves runtime and family defaults. -- `model_fit.kv_cache_policy=saver` prefers lower-memory cache dtypes, enables - offload where supported, and must warn if the backend cannot honor that plan. -- `model_fit.kv_cache_policy=auto` lets family or topology policy choose; any - explicit `cache_type_k` or `cache_type_v` value wins. - `throughput.tuning_profile=throughput` biases toward larger batch or ubatch, enables continuous batching when safe, and raises `parallel` if the estimator says memory permits it. - `throughput.tuning_profile=balanced` preserves today’s runtime defaults. - `throughput.tuning_profile=saver` biases toward smaller batch or ubatch, - lower `parallel`, and memory-saving KV defaults unless explicit fields win. + and lower `parallel`. - `safety_margin_gb` is not a saved key here. Later wiring should derive `hardware.fit_target_mib` by subtracting reserved MiB from detected allocatable memory and must never write the resolved MiB value back into TOML. @@ -87,10 +80,9 @@ website configuration reference, with the same `Wiring status`. | 5.1 | Context size | `model_fit.ctx_size` | P0 | `plugin/config.rs` | SkippyModelLoadOptions, StageConfig | single-stage, staged | restart/reload only | llama.cpp or stage runtime default | integer >= 1; size must fit estimated KV memory budget | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | ctx size control for model and stage load | wired | | 5.1 | Batch size | `model_fit.batch` | P0 | `plugin/config.rs` | SkippyModelLoadOptions, StageConfig, RuntimeConfig | single-stage, staged | restart/reload only | runtime family default unless tuning profile overrides | integer >= 1; should be >= ubatch when operator pins both | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | maps to n_batch | wired | | 5.1 | Micro-batch size | `model_fit.ubatch` | P0 | `plugin/config.rs` | SkippyModelLoadOptions, StageConfig, RuntimeConfig | single-stage, staged | restart/reload only | runtime family default unless tuning profile overrides | integer >= 1; should be <= batch to avoid contradictory sizing | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | maps to n_ubatch | wired | -| 5.1 | KV cache K dtype | `model_fit.cache_type_k` | P0 | `plugin/config.rs` | SkippyModelLoadOptions, StageConfig | single-stage, staged | restart/reload only | backend or family default unless kv_cache_policy selects one | enum value accepted by embedded llama runtime; explicit value overrides kv_cache_policy | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | explicit K cache dtype override | wired | -| 5.1 | KV cache V dtype | `model_fit.cache_type_v` | P0 | `plugin/config.rs` | SkippyModelLoadOptions, StageConfig | single-stage, staged | restart/reload only | backend or family default unless kv_cache_policy selects one | enum value accepted by embedded llama runtime; explicit value overrides kv_cache_policy | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | explicit V cache dtype override | wired | -| 5.1 | KV cache policy preset | `model_fit.kv_cache_policy` | P0 | `plugin/config.rs` | policy expander into cache_type_k, cache_type_v, kv_offload, cache_ram_mib | single-stage, staged | restart/reload only | mesh policy default | enum auto, quality, balanced, saver; explicit cache_type_k/v wins over preset expansion | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | quality=f16/f16 with no forced RAM cap; balanced=preserve runtime defaults; saver=prefer lower-memory dtypes plus offload warning if unsupported; auto=family or topology policy decides | wired | -| 5.1 | KV offload | `model_fit.kv_offload` | P1 | `plugin/config.rs` | SkippyModelLoadOptions, StageConfig, native `skippy_runtime_config.kv_offload` tri-state | single-stage, staged | restart/reload only | backend runtime default or kv_cache_policy expansion | boolean or auto; auto preserves llama.cpp's derived `offload_kqv` default | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; capture tests in `crates/skippy-runtime/src/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | controls KV residency and offload policy | wired | +| 5.1 | KV cache K dtype | `model_fit.cache_type_k` | P0 | `plugin/config.rs` | SkippyModelLoadOptions, StageConfig | single-stage, staged | restart/reload only | package-validated publisher KV declaration, then compute dtype, then F16 | enum value accepted by embedded llama runtime; explicit value wins | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | explicit K cache dtype override | wired | +| 5.1 | KV cache V dtype | `model_fit.cache_type_v` | P0 | `plugin/config.rs` | SkippyModelLoadOptions, StageConfig | single-stage, staged | restart/reload only | package-validated publisher KV declaration, then compute dtype, then F16 | enum value accepted by embedded llama runtime; explicit value wins | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | explicit V cache dtype override | wired | +| 5.1 | KV offload | `model_fit.kv_offload` | P1 | `plugin/config.rs` | SkippyModelLoadOptions, StageConfig, native `skippy_runtime_config.kv_offload` tri-state | single-stage, staged | restart/reload only | backend runtime default | boolean or auto; auto preserves llama.cpp's derived `offload_kqv` default | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; capture tests in `crates/skippy-runtime/src/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | controls KV residency and offload policy | wired | | 5.1 | Unified KV cache | `model_fit.kv_unified` | P1 | `plugin/config.rs` | SkippyModelLoadOptions, StageConfig, native `skippy_runtime_config.kv_unified` tri-state | single-stage, staged | restart/reload only | backend runtime default | boolean or auto; recurrent/hybrid architectures still force this true natively regardless of the requested value | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; capture tests in `crates/skippy-runtime/src/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | advanced cache layout flag | wired | | 5.1 | Cache RAM budget | `model_fit.cache_ram_mib` | P1 | `plugin/config.rs`, `skippy-server`, `skippy-cache` | `StageKvCacheConfig.l2_max_bytes` and bounded host-RAM L2 exact-state tier | single-stage, staged | restart/reload only | disabled by default | integer >= 0 MiB; zero or unset disables L2; positive values require prefix caching and active L3 | `#model-fit-context-and-kv-cache` | resolver propagation tests in `crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs`; promotion and admission tests in `skippy-server` and `skippy-cache`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | repeated or high-value L3 fills promote to L2; L2 hits restore before disk and asynchronously rewarm L1 | wired | | 5.1 | Cache idle slots | `model_fit.cache_idle_slots` | P1 | `plugin/config.rs` | SkippyModelLoadOptions, StageConfig, `skippy-server::RuntimeState::max_idle_sessions` idle-pool bound | single-stage, staged | restart/reload only | runtime default (unbounded, capped only by `lane_count`) | integer >= 0; bounds the idle session pool size | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; `crates/skippy-server/src/runtime_state.rs` and `crates/skippy-server/src/runtime_state/lane_lifecycle.rs` unit tests; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | bounds retained idle sessions so `drop_session_timed` discards lanes past the configured cap instead of growing the pool unbounded | wired | diff --git a/docs/skippy/manual-smoke/fixtures/model_fit/runtime-single-stage.toml b/docs/skippy/manual-smoke/fixtures/model_fit/runtime-single-stage.toml index 1304963d8b..fe411487a7 100644 --- a/docs/skippy/manual-smoke/fixtures/model_fit/runtime-single-stage.toml +++ b/docs/skippy/manual-smoke/fixtures/model_fit/runtime-single-stage.toml @@ -6,7 +6,6 @@ batch = 64 ubatch = 16 cache_type_k = "f16" cache_type_v = "f16" -kv_cache_policy = "balanced" kv_offload = "auto" prompt_cache = true swa_full = false diff --git a/docs/skippy/manual-smoke/manifest.tsv b/docs/skippy/manual-smoke/manifest.tsv index d67884d45c..c0c2225901 100644 --- a/docs/skippy/manual-smoke/manifest.tsv +++ b/docs/skippy/manual-smoke/manifest.tsv @@ -4,7 +4,6 @@ model_fit.batch docs/skippy/manual-smoke/fixtures/model_fit/runtime-single-stage model_fit.ubatch docs/skippy/manual-smoke/fixtures/model_fit/runtime-single-stage.toml python3 docs/skippy/manual-smoke/runtime_smoke.py --binary target/debug/mesh-llm --fixture docs/skippy/manual-smoke/fixtures/model_fit/runtime-single-stage.toml --model-path $MESH_LLM_SMOKE_MODEL_PATH --api-port 9437 --console-port 3231 --max-wait 600 unsloth/Qwen3.5-4B-GGUF:UD-Q4_K_XL wait for /api/status and /v1/models, then send /v1/chat/completions Live local runtime smoke succeeded; see model_fit section in evidence for /api/status, /v1/models, chat completion, and cleanup. .sisyphus/evidence/task-11-manual-runtime-smoke.txt PASS model_fit.cache_type_k docs/skippy/manual-smoke/fixtures/model_fit/runtime-single-stage.toml python3 docs/skippy/manual-smoke/runtime_smoke.py --binary target/debug/mesh-llm --fixture docs/skippy/manual-smoke/fixtures/model_fit/runtime-single-stage.toml --model-path $MESH_LLM_SMOKE_MODEL_PATH --api-port 9437 --console-port 3231 --max-wait 600 unsloth/Qwen3.5-4B-GGUF:UD-Q4_K_XL wait for /api/status and /v1/models, then send /v1/chat/completions Live local runtime smoke succeeded; see model_fit section in evidence for /api/status, /v1/models, chat completion, and cleanup. .sisyphus/evidence/task-11-manual-runtime-smoke.txt PASS model_fit.cache_type_v docs/skippy/manual-smoke/fixtures/model_fit/runtime-single-stage.toml python3 docs/skippy/manual-smoke/runtime_smoke.py --binary target/debug/mesh-llm --fixture docs/skippy/manual-smoke/fixtures/model_fit/runtime-single-stage.toml --model-path $MESH_LLM_SMOKE_MODEL_PATH --api-port 9437 --console-port 3231 --max-wait 600 unsloth/Qwen3.5-4B-GGUF:UD-Q4_K_XL wait for /api/status and /v1/models, then send /v1/chat/completions Live local runtime smoke succeeded; see model_fit section in evidence for /api/status, /v1/models, chat completion, and cleanup. .sisyphus/evidence/task-11-manual-runtime-smoke.txt PASS -model_fit.kv_cache_policy docs/skippy/manual-smoke/fixtures/model_fit/runtime-single-stage.toml python3 docs/skippy/manual-smoke/runtime_smoke.py --binary target/debug/mesh-llm --fixture docs/skippy/manual-smoke/fixtures/model_fit/runtime-single-stage.toml --model-path $MESH_LLM_SMOKE_MODEL_PATH --api-port 9437 --console-port 3231 --max-wait 600 unsloth/Qwen3.5-4B-GGUF:UD-Q4_K_XL wait for /api/status and /v1/models, then send /v1/chat/completions Live local runtime smoke succeeded; see model_fit section in evidence for /api/status, /v1/models, chat completion, and cleanup. .sisyphus/evidence/task-11-manual-runtime-smoke.txt PASS model_fit.kv_offload docs/skippy/manual-smoke/fixtures/model_fit/runtime-single-stage.toml python3 docs/skippy/manual-smoke/runtime_smoke.py --binary target/debug/mesh-llm --fixture docs/skippy/manual-smoke/fixtures/model_fit/runtime-single-stage.toml --model-path $MESH_LLM_SMOKE_MODEL_PATH --api-port 9437 --console-port 3231 --max-wait 600 unsloth/Qwen3.5-4B-GGUF:UD-Q4_K_XL wait for /api/status and /v1/models, then send /v1/chat/completions Live local runtime smoke succeeded; see model_fit section in evidence for /api/status, /v1/models, chat completion, and cleanup. .sisyphus/evidence/task-11-manual-runtime-smoke.txt PASS model_fit.prompt_cache docs/skippy/manual-smoke/fixtures/model_fit/runtime-single-stage.toml python3 docs/skippy/manual-smoke/runtime_smoke.py --binary target/debug/mesh-llm --fixture docs/skippy/manual-smoke/fixtures/model_fit/runtime-single-stage.toml --model-path $MESH_LLM_SMOKE_MODEL_PATH --api-port 9437 --console-port 3231 --max-wait 600 unsloth/Qwen3.5-4B-GGUF:UD-Q4_K_XL wait for /api/status and /v1/models, then send /v1/chat/completions Live local runtime smoke succeeded; see model_fit section in evidence for /api/status, /v1/models, chat completion, and cleanup. .sisyphus/evidence/task-11-manual-runtime-smoke.txt PASS model_fit.prefix_cache.enabled docs/skippy/manual-smoke/fixtures/model_fit/runtime-single-stage.toml python3 docs/skippy/manual-smoke/runtime_smoke.py --binary target/debug/mesh-llm --fixture docs/skippy/manual-smoke/fixtures/model_fit/runtime-single-stage.toml --model-path $MESH_LLM_SMOKE_MODEL_PATH --api-port 9437 --console-port 3231 --max-wait 600 unsloth/Qwen3.5-4B-GGUF:UD-Q4_K_XL wait for /api/status and /v1/models, then send /v1/chat/completions Live local runtime smoke succeeded; see model_fit section in evidence for /api/status, /v1/models, chat completion, and cleanup. .sisyphus/evidence/task-11-manual-runtime-smoke.txt PASS diff --git a/website/src/docs/pages/config-defaults.md b/website/src/docs/pages/config-defaults.md index 97f6b63306..ea7a583695 100644 --- a/website/src/docs/pages/config-defaults.md +++ b/website/src/docs/pages/config-defaults.md @@ -14,9 +14,8 @@ Shared default settings applied to every model. Individual model entries can ove ctx_size = 0 # Context window (0 = auto) batch = 0 # Batch size (0 = auto) ubatch = 0 # Micro-batch size (0 = auto) -cache_type_k = "auto" # Key cache dtype (resolved by policy) -cache_type_v = "auto" # Value cache dtype (resolved by policy) -kv_cache_policy = "balanced" # "balanced", "quality", or "saver" +cache_type_k = "auto" # Key cache dtype (publisher metadata, then F16) +cache_type_v = "auto" # Value cache dtype (publisher metadata, then F16) kv_offload = "auto" # KV-cache offload policy prompt_cache = "auto" # Prompt-cache policy flash_attention = "auto" # Flash-attention policy diff --git a/website/src/docs/pages/config-reference.md b/website/src/docs/pages/config-reference.md index 58435cd521..5d12c15989 100644 --- a/website/src/docs/pages/config-reference.md +++ b/website/src/docs/pages/config-reference.md @@ -140,8 +140,7 @@ for the activity policy and privacy boundary. | `model_fit.ctx_size` | integer | `0` = auto | both | model reload | wired | `--ctx-size` on the ad-hoc single-model path | | `model_fit.batch` | integer | `0` = auto (`n_batch`) | both | model reload | wired | none | | `model_fit.ubatch` | integer | `0` = auto (`n_ubatch`); should not exceed `batch` | both | model reload | wired | none | -| `model_fit.cache_type_k`
`model_fit.cache_type_v` | enum (dtype) | `auto` (default), `f16`, `q8_0`, `q4_0`; schema validation also accepts `f32`, `bf16`, `q4_1`, `iq4_nl`, `q5_0`, `q5_1`, but the normal serving runtime rejects them at model load; explicit value overrides `kv_cache_policy` | both | model reload | wired | none | -| `model_fit.kv_cache_policy` | enum | `balanced` (default), `auto`, `quality`, `saver`; expands into cache dtypes | both | model reload | wired | none | +| `model_fit.cache_type_k`
`model_fit.cache_type_v` | enum (dtype) | `auto` (default) follows package-validated publisher KV metadata, then publisher compute dtype, then F16; explicit schema values are `f16`, `f32`, `bf16`, `q8_0`, `q4_0`, `q4_1`, `iq4_nl`, `q5_0`, and `q5_1`, gated by runtime support | both | model reload | wired | none | | `model_fit.kv_offload` | bool-or-`auto` | `auto` | both | model reload | wired | none | | `model_fit.kv_unified` | bool-or-`auto` | `auto` | both | model reload | wired (recurrent/hybrid architectures still force this true natively) | none | | `model_fit.cache_ram_mib` | integer | `0`/unset = host-RAM L2 disabled | both | model reload | wired; requires prefix caching and active L3 | none | From be8058c9d20ab4b320fb6dc89101ffd364dda2ca Mon Sep 17 00:00:00 2001 From: Paul Hogan Date: Tue, 15 Sep 2026 21:09:04 +1000 Subject: [PATCH 13/16] fix(mesh): un-gate upsert_served_model_descriptor for production callers The merge of #1879 into main moved advertise_startup_sources into production runtime code (run_auto), but upsert_served_model_descriptor remained #[cfg(test)], breaking non-test builds the same way it broke main's own Quality/Website CI. Remove the gate; the method is required by model_presentation::advertise_startup_sources. --- crates/mesh-llm-host-runtime/src/mesh/node.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/crates/mesh-llm-host-runtime/src/mesh/node.rs b/crates/mesh-llm-host-runtime/src/mesh/node.rs index 1bad7c7825..7aef15d97d 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/node.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/node.rs @@ -1351,7 +1351,6 @@ impl Node { true } - #[cfg(test)] pub async fn upsert_served_model_descriptor(&self, descriptor: ServedModelDescriptor) { let mut descriptors = self.served_model_descriptors.lock().await; if let Some(existing) = descriptors From 260ca065431fe40fa1a79bf436adb051ca011dcb Mon Sep 17 00:00:00 2001 From: scama Date: Wed, 16 Sep 2026 10:26:31 +1000 Subject: [PATCH 14/16] fix(ci): refresh console print ratchet --- tools/xtask/data/console_print_allowlist.json | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tools/xtask/data/console_print_allowlist.json b/tools/xtask/data/console_print_allowlist.json index 1e3588fd95..2be10f07db 100644 --- a/tools/xtask/data/console_print_allowlist.json +++ b/tools/xtask/data/console_print_allowlist.json @@ -3787,11 +3787,11 @@ ], "crates/skippy-correctness/src/runner/kv_page_growth.rs": [ { - "line": 285, + "line": 289, "macro_name": "println!" }, { - "line": 340, + "line": 344, "macro_name": "println!" } ], @@ -3811,27 +3811,27 @@ "macro_name": "eprintln!" }, { - "line": 933, + "line": 937, "macro_name": "eprintln!" }, { - "line": 942, + "line": 946, "macro_name": "eprintln!" }, { - "line": 954, + "line": 958, "macro_name": "eprintln!" }, { - "line": 957, + "line": 961, "macro_name": "eprintln!" }, { - "line": 961, + "line": 965, "macro_name": "eprintln!" }, { - "line": 1309, + "line": 1313, "macro_name": "eprintln!" } ], From 11cab796992c7c5d88c24f46f075bc62dfc94c71 Mon Sep 17 00:00:00 2001 From: scama Date: Wed, 16 Sep 2026 15:50:55 +1000 Subject: [PATCH 15/16] fix(ffi): export consuming writer for dynamic runtimes --- .../src/inference/skippy/split-certified.json | 2 +- crates/skippy-ffi/src/lib.rs | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/split-certified.json b/crates/mesh-llm-host-runtime/src/inference/skippy/split-certified.json index f921de5353..80cc306245 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/split-certified.json +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/split-certified.json @@ -2,7 +2,7 @@ "schema_version": 1, "native_recipe": { "llama_upstream_sha": "661643e43079a4ee6faab4c1895291767b67ea8d", - "skippy_abi": "0.1.57", + "skippy_abi": "0.1.58", "patch_queue_sha256": "37f66961e57c8cbc0988e02c2a5a677cf0dea28ae6ef719f4fdd4e60d49d9abd" }, "models": [ diff --git a/crates/skippy-ffi/src/lib.rs b/crates/skippy-ffi/src/lib.rs index a22a428912..6e50743c4b 100644 --- a/crates/skippy-ffi/src/lib.rs +++ b/crates/skippy-ffi/src/lib.rs @@ -5,7 +5,7 @@ mod dynamic_library; // without compiling the crate to determine native-runtime compatibility. pub const ABI_VERSION_MAJOR: u32 = 0; pub const ABI_VERSION_MINOR: u32 = 1; -pub const ABI_VERSION_PATCH: u32 = 57; +pub const ABI_VERSION_PATCH: u32 = 58; // Propagate static native archive changes through Cargo dependency metadata so // final binaries are relinked after CMake rebuilds llama.cpp. @@ -138,7 +138,8 @@ pub use dynamic::{ skippy_stage_planner_create_v1, skippy_stage_planner_free, skippy_stage_planner_realize_v1, skippy_token_is_eog, skippy_tokenize, skippy_trim_session, skippy_verify_tokens, skippy_verify_tokens_frame_sampled, skippy_write_gguf_from_parts, - skippy_write_gguf_metadata_from_parts, skippy_write_slice_gguf, + skippy_write_gguf_from_parts_consuming, skippy_write_gguf_metadata_from_parts, + skippy_write_slice_gguf, }; #[cfg(feature = "dynamic-runtime")] From b06a9e01f90630971856f024831fce8e5c025cfd Mon Sep 17 00:00:00 2001 From: Paul Hogan Date: Wed, 16 Sep 2026 17:33:48 +1000 Subject: [PATCH 16/16] fix: satisfy no-console-print gate on kv-cache command output --- crates/mesh-llm-commands/src/kv_cache.rs | 35 ++++++++++++++---------- 1 file changed, 21 insertions(+), 14 deletions(-) diff --git a/crates/mesh-llm-commands/src/kv_cache.rs b/crates/mesh-llm-commands/src/kv_cache.rs index ab7e22630e..463b9da0c2 100644 --- a/crates/mesh-llm-commands/src/kv_cache.rs +++ b/crates/mesh-llm-commands/src/kv_cache.rs @@ -148,8 +148,9 @@ fn confirm_destructive(action: &str, yes: bool) -> Result<()> { if !io::stdin().is_terminal() { bail!("{action} requires --yes when stdin is not interactive"); } - eprint!("{action}? [y/N] "); - io::stderr().flush()?; + let mut err = mesh_llm_events::console_err(); + let _ = write!(err, "{action}? [y/N] "); + let _ = err.flush(); let mut response = String::new(); io::stdin().read_line(&mut response)?; if !matches!(response.trim().to_ascii_lowercase().as_str(), "y" | "yes") { @@ -160,9 +161,11 @@ fn confirm_destructive(action: &str, yes: bool) -> Result<()> { fn print_response(value: &Value, json_output: bool) -> Result<()> { if json_output { - println!("{}", serde_json::to_string(value)?); + let mut machine = mesh_llm_events::machine_out(); + writeln!(machine, "{}", serde_json::to_string(value)?)?; return Ok(()); } + let mut err = mesh_llm_events::console_err(); if let Some(results) = value.get("results").and_then(Value::as_array) { for result in results { let node = result @@ -170,15 +173,16 @@ fn print_response(value: &Value, json_output: bool) -> Result<()> { .and_then(Value::as_str) .unwrap_or("unknown"); if let Some(error) = result.get("error").filter(|value| !value.is_null()) { - println!( + writeln!( + err, "Node {node}: error: {}", error .get("message") .and_then(Value::as_str) .unwrap_or("unknown error") - ); + )?; } else if let Some(freed) = result.get("freed_bytes").and_then(Value::as_u64) { - println!("Node {node}: freed {freed} bytes"); + writeln!(err, "Node {node}: freed {freed} bytes")?; } else { let state = result .get("status") @@ -186,32 +190,35 @@ fn print_response(value: &Value, json_output: bool) -> Result<()> { .and_then(|effective| effective.get("state")) .and_then(Value::as_str) .unwrap_or("unknown"); - println!("Node {node}: {state}"); + writeln!(err, "Node {node}: {state}")?; } } return Ok(()); } if let Some(freed) = value.get("freed_bytes").and_then(Value::as_u64) { - println!("Freed {freed} bytes"); + writeln!(err, "Freed {freed} bytes")?; return Ok(()); } let configured = &value["configured"]; let effective = &value["effective"]; - println!( + writeln!( + err, "Disk prompt cache: {} ({})", effective["state"].as_str().unwrap_or("unknown"), configured["mode"].as_str().unwrap_or("unknown") - ); - println!( + )?; + writeln!( + err, "Root: {}", configured["directory"].as_str().unwrap_or("unknown") - ); + )?; if let Some(usage) = value.get("usage").filter(|usage| !usage.is_null()) { - println!( + writeln!( + err, "Used: {} / {} bytes", usage["used_bytes"].as_u64().unwrap_or(0), usage["budget_bytes"].as_u64().unwrap_or(0) - ); + )?; } Ok(()) }