From 2065ab7df8d9a3b0efe9fd9fd42d9f5ed67923b8 Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Wed, 12 Aug 2026 16:21:44 +0800 Subject: [PATCH 01/16] fix(connector): accept full attention spec subclasses (#429) ## Summary - classify `FullAttentionSpec` subclasses as full-attention cache groups - allow specialized full-attention specs used by the latest Kimi K3 vLLM release to initialize with aligned Mamba groups - cover subclass and MLA inheritance behavior in cache-group layout tests - bump the workspace and Python package version to `0.23.10` for release ## Testing - `cd python && uv run --extra test pytest` (256 passed, 13 deselected) - `cd python && uv run --isolated --no-project --with pytest --with numpy --with 'requests>=2.26.0' pytest` (256 passed, 13 deselected) - pre-commit hooks, including `cargo test --release`, Ruff, and Commitizen checks - `cargo metadata --locked --no-deps --format-version 1` - Cargo workspace, Python package, and Commitizen versions verified at `0.23.10` ## Release After merge, create and push tag `v0.23.10` to trigger the release workflow. ## Performance Work Not a performance change; no benchmark was required. --- Cargo.lock | 16 ++++++++-------- Cargo.toml | 2 +- python/pegaflow/connector/common.py | 9 +++------ python/pyproject.toml | 4 ++-- python/tests/test_cache_group_layout.py | 19 ++++++++++++++----- 5 files changed, 28 insertions(+), 22 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1229644a..5c918629 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1896,7 +1896,7 @@ checksum = "19b17cddbe7ec3f8bc800887bab5e717348c95ea2ca0b1bf0837fb964dc67099" [[package]] name = "pegaflow-common" -version = "0.23.9" +version = "0.23.10" dependencies = [ "colored", "libc", @@ -1906,7 +1906,7 @@ dependencies = [ [[package]] name = "pegaflow-core" -version = "0.23.9" +version = "0.23.10" dependencies = [ "ahash", "bytesize", @@ -1939,7 +1939,7 @@ dependencies = [ [[package]] name = "pegaflow-metaserver" -version = "0.23.9" +version = "0.23.10" dependencies = [ "axum", "clap", @@ -1959,7 +1959,7 @@ dependencies = [ [[package]] name = "pegaflow-pd-wire" -version = "0.23.9" +version = "0.23.10" dependencies = [ "serde", "serde_json", @@ -1967,7 +1967,7 @@ dependencies = [ [[package]] name = "pegaflow-proto" -version = "0.23.9" +version = "0.23.10" dependencies = [ "prost", "tonic", @@ -1977,7 +1977,7 @@ dependencies = [ [[package]] name = "pegaflow-py" -version = "0.23.9" +version = "0.23.10" dependencies = [ "log", "mea", @@ -1995,7 +1995,7 @@ dependencies = [ [[package]] name = "pegaflow-server" -version = "0.23.9" +version = "0.23.10" dependencies = [ "axum", "clap", @@ -2027,7 +2027,7 @@ dependencies = [ [[package]] name = "pegaflow-transfer" -version = "0.23.9" +version = "0.23.10" dependencies = [ "anyhow", "bincode", diff --git a/Cargo.toml b/Cargo.toml index d74f9323..c5ae3a6c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,7 @@ members = [ resolver = "2" [workspace.package] -version = "0.23.9" +version = "0.23.10" edition = "2024" license = "Apache-2.0" diff --git a/python/pegaflow/connector/common.py b/python/pegaflow/connector/common.py index 859fd309..d08e8ea9 100644 --- a/python/pegaflow/connector/common.py +++ b/python/pegaflow/connector/common.py @@ -280,15 +280,12 @@ def from_config(cls, kv_cache_config) -> "CacheGroupLayout": "or uniformly grouped MLA layers" ) else: - if any( - type(spec) is not FullAttentionSpec and not isinstance(spec, MambaSpec) - for spec in specs - ): + if any(not isinstance(spec, (FullAttentionSpec, MambaSpec)) for spec in specs): raise RuntimeError( "PegaFlow HMA supports only FullAttention and Mamba cache groups" ) - has_full_attention = any(type(spec) is FullAttentionSpec for spec in specs) + has_full_attention = any(isinstance(spec, FullAttentionSpec) for spec in specs) has_mamba = any(isinstance(spec, MambaSpec) for spec in specs) if not has_full_attention: raise RuntimeError( @@ -316,7 +313,7 @@ def from_config(cls, kv_cache_config) -> "CacheGroupLayout": ( index for index, group in enumerate(groups) - if type(group.kv_cache_spec) is FullAttentionSpec + if isinstance(group.kv_cache_spec, FullAttentionSpec) ), None, ) diff --git a/python/pyproject.toml b/python/pyproject.toml index 5a3db29b..a3cb17c6 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "pegaflow-llm" -version = "0.23.9" +version = "0.23.10" description = "High-performance key-value storage engine with Python bindings" readme = "README.md" requires-python = ">=3.10" @@ -114,6 +114,6 @@ profile = "black" [tool.commitizen] name = "cz_conventional_commits" -version = "0.23.9" +version = "0.23.10" version_files = ["pyproject.toml:^version", "../Cargo.toml:^version"] tag_format = "v$version" diff --git a/python/tests/test_cache_group_layout.py b/python/tests/test_cache_group_layout.py index 15451279..d26253af 100644 --- a/python/tests/test_cache_group_layout.py +++ b/python/tests/test_cache_group_layout.py @@ -49,9 +49,16 @@ def _mla(block_size=16, head_size=128): return spec -def test_accepts_full_attention_with_aligned_mamba(): +class SpecializedFullAttentionSpec(FullAttentionSpec): + pass + + +@pytest.mark.parametrize("spec_type", [FullAttentionSpec, SpecializedFullAttentionSpec]) +def test_accepts_full_attention_with_aligned_mamba(spec_type): + attention = spec_type() + attention.block_size = 528 config = _config( - _group("attention", _full_attention(block_size=528)), + _group("attention", attention), _group("recurrent", _mamba(block_size=528)), ) @@ -170,7 +177,7 @@ def test_rejects_full_attention_with_sliding_window(): CacheGroupLayout.from_config(config) -def test_rejects_mla_with_mamba(): +def test_accepts_mla_with_mamba(): mla = MLAAttentionSpec() mla.block_size = 16 config = _config( @@ -178,8 +185,10 @@ def test_rejects_mla_with_mamba(): _group("recurrent", _mamba()), ) - with pytest.raises(RuntimeError, match="only FullAttention and Mamba"): - CacheGroupLayout.from_config(config) + layout = CacheGroupLayout.from_config(config) + + assert layout.hash_group_index == 0 + assert layout.has_recurrent_state def test_rejects_non_align_mamba_mode(): From 6e31cd39c71651b45ceb8456530084caf771cbac Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Thu, 13 Aug 2026 15:47:41 +0800 Subject: [PATCH 02/16] fix(connector): discard drifted TP prefetch results (#432) ## Summary - keep polling an in-flight prefetch with its original query identity when local prefix hits change - let TP shard validation check the stale result against the query that created it - release stale per-shard leases before issuing the updated query ## Root cause The scheduler already snapshots `computed_blocks` and query hashes when `query_prefetch` returns Loading. After TP shard support was added, each shard validated a later Ready result against the current, shorter query before the scheduler could compare it with the saved snapshot. A valid stale result therefore raised `hits > queried blocks` instead of being discarded. ## Tests - `cd python && uv run --extra test pytest -q` (257 passed, 13 deselected) - `cd python && uv run --extra test pytest -q tests/test_tp_shards.py tests/test_combine_hashes.py` (66 passed) - commit hooks, including `cargo test --release`, ruff, format, typos, and Commitizen --- python/pegaflow/connector/scheduler.py | 11 ++++-- python/tests/test_tp_shards.py | 46 ++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/python/pegaflow/connector/scheduler.py b/python/pegaflow/connector/scheduler.py index a83e72e5..2e162869 100644 --- a/python/pegaflow/connector/scheduler.py +++ b/python/pegaflow/connector/scheduler.py @@ -237,9 +237,16 @@ def get_num_new_matched_tokens( self._release_pending_query_probe(req_id) probe = None - # No reusable Ready result. Ask backend. + # A Loading task is keyed by req_id server-side. If the current query + # drifted, finish polling with the original identity so the TP layer can + # validate the stale Ready before we release it below. + backend_query_hashes = query_hashes + if probe is not None and not probe.matches(computed_blocks, query_hashes, tail_tokens): + backend_query_hashes = probe.query_hashes + + # No reusable Ready result. Ask backend. lookup_start = time.perf_counter() - ready = self._count_available_block_prefix(query_hashes, req_id) + ready = self._count_available_block_prefix(backend_query_hashes, req_id) lookup_us = (time.perf_counter() - lookup_start) * 1e6 # Backend is still loading. Keep the original snapshot. diff --git a/python/tests/test_tp_shards.py b/python/tests/test_tp_shards.py index 4ef02ed0..063d73a0 100644 --- a/python/tests/test_tp_shards.py +++ b/python/tests/test_tp_shards.py @@ -238,6 +238,52 @@ def test_scheduler_releases_ready_shards_when_another_shard_is_loading(): first.release.assert_called_once_with(b"first") +def test_scheduler_discards_drifted_prefetch_before_querying_new_hashes(): + first = MagicMock() + second = MagicMock() + first.query_prefetch.side_effect = [ + QueryLoading(), + QueryReady(4, b"first-old"), + QueryLoading(), + ] + second.query_prefetch.return_value = QueryReady(4, b"second-old") + scheduler = SchedulerConnector(_context(), engine_clients=(first, second)) + request = SimpleNamespace( + request_id="request", + block_hashes=[b"h0", b"h1", b"h2", b"h3"], + num_tokens=64, + ) + + assert scheduler.get_num_new_matched_tokens(request, 0) == (None, False) + assert scheduler.get_num_new_matched_tokens(request, 32) == (None, False) + assert scheduler.get_num_new_matched_tokens(request, 32) == (None, False) + + original_hashes = request.block_hashes + current_hashes = request.block_hashes[2:] + assert first.query_prefetch.call_args_list == [ + call( + "instance", + original_hashes, + req_id="request", + wait_for_full_prefix=False, + ), + call( + "instance", + original_hashes, + req_id="request", + wait_for_full_prefix=False, + ), + call( + "instance", + current_hashes, + req_id="request", + wait_for_full_prefix=False, + ), + ] + first.release.assert_called_once_with(b"first-old") + second.release.assert_called_once_with(b"second-old") + + @pytest.mark.parametrize( "invalid_ready", [ From 9b7d31689e4967b5cc639200e63dbebb3a886ebe Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Fri, 14 Aug 2026 13:41:18 +0800 Subject: [PATCH 03/16] feat: hybrid KV cache group save/load for vLLM models with mixed attention and recurrent layers (#433) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What this changes End-to-end PegaFlow support for vLLM v1 hybrid KV cache models (full attention layers mixed with Mamba/KDA/GDN recurrent layers), landed on top of the group-aware plumbing introduced in recent connector work and proved out on GB300 (SM100) with Kimi-K3. Before this, every engine-side keying decision assumed a single cache group: group 0's content hashes addressed all segments, so recurrent groups were never saved, and a warm restore for a hybrid model could only answer attention hash queries. This patch teaches the whole stack about cache groups: ### Engine (`feat(core)`) - `pegaflow-common::block::group_hash` encodes a block content hash with its vLLM cache-group id. Group 0 stays byte-identical to the raw hash, so all existing single-group deployments produce exactly the same block keys as before. - `RegisterContextRequest` now carries `layer_group_ids`; the instance registry seals a per-group layer set and slot topology and refuses to serve instances whose workers disagree about group layout. - Read cache, write path, layup and offload all carry explicit group ids; slot derivation dedupes per group instead of per layer. - `QueryRequest.group_id` selects semantics: group 0 keeps the existing prefix query, other groups answer sparse-checkpoint membership queries. Membership results do not feed remote RDMA fetch in this release (see Limitations). - gRPC service routes save fan-out and per-(slot, hash) reports so the connector can place recurrent checkpoints where its membership leases require. ### Connector (`feat(connector)`) - Every group's block hashes are derived from the same prefix hash chain (group id + checkpoint-offset/private-salt emptiness mixed in) so a restore is only counted as a hit when every group's data can come back together. Group 0 provides prompt-layout blocks; recurrent groups provide sparse checkpoints at per-model-config positions, including the final HMA blocks (`fix(connector): save final HMA cache blocks`). - `_reconcile_hybrid` combines per-group answers into the longest fully restorable prefix; a partial reconcile is dropped entirely rather than restored with stale recurrent context. - Save/load intents keep per-group slot positions and write recurrent groups under their own hash namespaces inside the same request flow. - TP shard topology now accepts an explicit `pegaflow.tp_shard_endpoints` list so rank-to-server mapping no longer depends on every server being reachable at the same host:port pattern; without it the previous deterministic default layout is preserved. ## Verification - New coverage: `pegaflow-core/tests/hybrid_groups.rs`, instance-slot sealing tests for multi-group topologies, and `python/tests/test_hybrid_reconcile.py` for the reconcile/drop/final- block contracts. All existing engine, connector and fault-tolerance gates pass unmodified (`prek run` green: fmt, clippy, release-profile tests, ruff, typos). - Cluster validation on GB300 (SM100, aarch64): Kimi-K3 (MLA + KDA), 6 `xingming-k3-debug` LWS groups of TP8x2 with per-tray `pegaflow-server` shards (`550gb` hugepage pools each). Under Kimi production prefill traffic, port 8001 shows 24/24-25/25 layout block hits after a cold pass (TTFT 4.9s -> 1.7s for a ~39K-token prompt), and pooled saves/loads run at multi-TB/hour overnight with `pegaflow_cache_block_hits`-driven throughput scaling as expected. ## Limitations - The vLLM connector's reconcile keeps using the local membership form: its sparse queries stay resident-cache-only, so hybrid models still recompute after local eviction of recurrent checkpoints. Engine-side the limitation is closed: `wait_for_full_prefix` on a group > 0 query selects an all-or-nothing set fetch that routes misses through the same SSD prefetch and MetaServer + RDMA machinery as prefix queries (the want-set either completes or reports short, which callers treat as a miss) — the form a prefill/decode handoff consumer needs. Wiring the connector's warm-restore path onto it is follow-up work. ## Test plan - `cargo test --workspace --no-default-features --features cuda-13,rdma` - `cd python && uv run --extra test pytest` - `python/tests/test_vllm_e2e_correctness.py -m e2e` (hybrid model, GB300 cluster): cold -> warm verify loop with `hit_blocks` matching the full layout block count. --------- Co-authored-by: xiaguan --- Cargo.lock | 16 +- Cargo.toml | 2 +- pegaflow-common/src/block.rs | 48 +++ pegaflow-common/src/lib.rs | 2 +- pegaflow-core/src/instance.rs | 138 +++++++- pegaflow-core/src/instance/tests.rs | 142 ++++++++ pegaflow-core/src/lib.rs | 163 ++++++++- pegaflow-core/src/offload.rs | 141 +++++--- pegaflow-core/src/storage/mod.rs | 17 +- pegaflow-core/src/storage/read_cache.rs | 16 + pegaflow-core/tests/common/harness.rs | 29 ++ pegaflow-core/tests/common/helpers.rs | 7 +- pegaflow-core/tests/hybrid_groups.rs | 323 ++++++++++++++++++ pegaflow-proto/proto/engine.proto | 31 +- pegaflow-server/src/service.rs | 165 ++++++--- pegaflow-server/tests/common/mod.rs | 1 + .../tests/http_cleanup_hang_repro.rs | 1 + python/pegaflow/connector/common.py | 102 +++++- python/pegaflow/connector/scheduler.py | 220 +++++++++++- python/pegaflow/connector/tp_shards.py | 59 +++- python/pegaflow/connector/worker.py | 73 +++- python/pegaflow/pegaflow.pyi | 10 +- python/pyproject.toml | 4 +- python/src/lib.rs | 26 +- python/tests/test_combine_hashes.py | 46 +++ .../tests/test_connector_fault_tolerance.py | 4 +- python/tests/test_connector_save_lifecycle.py | 44 +++ python/tests/test_hybrid_reconcile.py | 128 +++++++ 28 files changed, 1812 insertions(+), 146 deletions(-) create mode 100644 pegaflow-core/tests/hybrid_groups.rs create mode 100644 python/tests/test_hybrid_reconcile.py diff --git a/Cargo.lock b/Cargo.lock index 5c918629..c36ecde6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1896,7 +1896,7 @@ checksum = "19b17cddbe7ec3f8bc800887bab5e717348c95ea2ca0b1bf0837fb964dc67099" [[package]] name = "pegaflow-common" -version = "0.23.10" +version = "0.23.11" dependencies = [ "colored", "libc", @@ -1906,7 +1906,7 @@ dependencies = [ [[package]] name = "pegaflow-core" -version = "0.23.10" +version = "0.23.11" dependencies = [ "ahash", "bytesize", @@ -1939,7 +1939,7 @@ dependencies = [ [[package]] name = "pegaflow-metaserver" -version = "0.23.10" +version = "0.23.11" dependencies = [ "axum", "clap", @@ -1959,7 +1959,7 @@ dependencies = [ [[package]] name = "pegaflow-pd-wire" -version = "0.23.10" +version = "0.23.11" dependencies = [ "serde", "serde_json", @@ -1967,7 +1967,7 @@ dependencies = [ [[package]] name = "pegaflow-proto" -version = "0.23.10" +version = "0.23.11" dependencies = [ "prost", "tonic", @@ -1977,7 +1977,7 @@ dependencies = [ [[package]] name = "pegaflow-py" -version = "0.23.10" +version = "0.23.11" dependencies = [ "log", "mea", @@ -1995,7 +1995,7 @@ dependencies = [ [[package]] name = "pegaflow-server" -version = "0.23.10" +version = "0.23.11" dependencies = [ "axum", "clap", @@ -2027,7 +2027,7 @@ dependencies = [ [[package]] name = "pegaflow-transfer" -version = "0.23.10" +version = "0.23.11" dependencies = [ "anyhow", "bincode", diff --git a/Cargo.toml b/Cargo.toml index c5ae3a6c..f3c502f1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,7 @@ members = [ resolver = "2" [workspace.package] -version = "0.23.10" +version = "0.23.11" edition = "2024" license = "Apache-2.0" diff --git a/pegaflow-common/src/block.rs b/pegaflow-common/src/block.rs index 6ff795b6..e899ba53 100644 --- a/pegaflow-common/src/block.rs +++ b/pegaflow-common/src/block.rs @@ -23,3 +23,51 @@ impl BlockKey { (self.namespace.capacity() + self.hash.capacity() + 48) as u64 } } + +/// Encode a raw content hash with a hybrid-cache group id. +/// +/// Group 0 keeps the raw hash byte-for-byte so existing single-group caches +/// (and every current connector) stay bit-identical. Groups >= 1 append the +/// big-endian group id, which cannot collide with a raw content hash because +/// every real hash family is fixed-length. +pub fn group_hash(hash: &[u8], group_id: u32) -> Vec { + if group_id == 0 { + return hash.to_vec(); + } + let mut encoded = Vec::with_capacity(hash.len() + 4); + encoded.extend_from_slice(hash); + encoded.extend_from_slice(&group_id.to_be_bytes()); + encoded +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn group_zero_hash_is_bit_identical_to_raw() { + // Backward-compat contract: existing single-group deployments must not + // observe any key change. + let hash = b"sha256-content-hash"; + assert_eq!(group_hash(hash, 0), hash.to_vec()); + assert_eq!(group_hash(b"", 0), Vec::::new()); + } + + #[test] + fn nonzero_group_appends_big_endian_group_id() { + let hash = [0xAA, 0xBB]; + assert_eq!(group_hash(&hash, 1), vec![0xAA, 0xBB, 0, 0, 0, 1]); + assert_eq!(group_hash(&hash, 0x01020304), vec![0xAA, 0xBB, 1, 2, 3, 4]); + } + + #[test] + fn distinct_groups_do_not_share_keys() { + // The same content hash in two groups must be two different keys, and + // an encoded key must never equal a raw hash of any length (length + // differs, so this holds even across hash families). + let hash = [1, 2, 3, 4]; + assert_ne!(group_hash(&hash, 0), group_hash(&hash, 1)); + assert_ne!(group_hash(&hash, 1), group_hash(&hash, 2)); + assert_ne!(group_hash(&hash, 1).len(), hash.len()); + } +} diff --git a/pegaflow-common/src/lib.rs b/pegaflow-common/src/lib.rs index abe43f30..2438e444 100644 --- a/pegaflow-common/src/lib.rs +++ b/pegaflow-common/src/lib.rs @@ -5,7 +5,7 @@ pub mod logging; #[cfg(target_os = "linux")] pub mod numa; -pub use block::BlockKey; +pub use block::{BlockKey, group_hash}; #[cfg(target_os = "linux")] pub use numa::{ NumaNode, NumaTopology, format_cpu_list, pin_thread_to_numa_node, query_pages_numa, diff --git a/pegaflow-core/src/instance.rs b/pegaflow-core/src/instance.rs index 263d0462..52c6da26 100644 --- a/pegaflow-core/src/instance.rs +++ b/pegaflow-core/src/instance.rs @@ -49,11 +49,25 @@ struct RegistrationState { /// /// Layer ids are the rank of the layer name in sorted order, so every /// instance that registers the same layer set derives the same ids — the -/// property block slot layout depends on. +/// property block slot layout depends on. On top of that, each layer belongs +/// to a hybrid-cache *storage group* (attention vs. recurrent state, mirroring +/// vLLM's `KVCacheGroupSpec`): slots are dense *within* a group, giving each +/// group its own seal domain. A group seals a block once every +/// `(group layer, tp_rank)` slot of that block is saved, regardless of other +/// groups — that is what lets a recurrent-state group seal the final block +/// even though attention groups save every block. #[derive(Debug)] pub(crate) struct LayerTopology { name_to_id: HashMap, tp_size: usize, + /// Storage group id per layer, indexed by layer_id. All zeros for + /// single-group (classic) instances. + layer_group: Vec, + /// Layer count per storage group, indexed by group id. + group_layer_count: Vec, + /// Dense rank of the layer within its storage group (0..group_layer_count), + /// indexed by layer_id. Within-group slots are `rank * tp_size + tp_rank`. + layer_group_rank: Vec, /// Page-first layout when `Some`: each block's layers collapse into /// contiguous per-shard pages, so `total_slots = num_shards`. `None` is the /// legacy layer-first layout (`total_slots = num_layers * tp_size`). @@ -167,6 +181,11 @@ impl LayerTopology { /// Total number of storage slots per block: `num_shards` page-first, else /// `num_layers * tp_size` (layer-first). + /// + /// This is the union across groups and remains the right answer for + /// single-group instances; multi-group callers must use + /// [`Self::group_total_slots`] — blocks seal per group, so one global + /// denominator would never fill. pub(crate) fn total_slots(&self) -> usize { match &self.page_layout { Some(p) => p.shard_page_sizes.len(), @@ -174,6 +193,34 @@ impl LayerTopology { } } + /// Number of distinct storage groups. Always >= 1. + pub(crate) fn num_groups(&self) -> usize { + self.group_layer_count.len() + } + + /// Storage group id of a layer. + pub(crate) fn group_of_layer(&self, layer_id: usize) -> u32 { + self.layer_group[layer_id] + } + + /// Number of storage slots per block within one group. A block of that + /// group seals when exactly this many slots are saved. + pub(crate) fn group_total_slots(&self, group_id: u32) -> Result { + let group_idx = group_id as usize; + if group_idx >= self.num_groups() { + return Err(EngineError::InvalidArgument(format!( + "storage group {group_id} out of range ({} groups)", + self.num_groups() + ))); + } + Ok(match &self.page_layout { + // Page-first is seal-validated to be single-group: slots follow + // shards, not the layer grid. + Some(p) => p.shard_page_sizes.len(), + None => self.group_layer_count[group_idx] * self.tp_size, + }) + } + /// Page-first only: contiguous page size in bytes of `shard` (one slot). pub(crate) fn shard_page_size(&self, shard: usize) -> Option { self.page_layout.as_ref().map(|p| p.shard_page_sizes[shard]) @@ -198,7 +245,9 @@ impl LayerTopology { /// /// Page-first collapses the layer dimension into per-shard pages, so the /// slot is the layer's shard; the layer's position is its page offset - /// ([`Self::page_placement`]). Layer-first keeps `[layer][tp_rank]`. + /// ([`Self::page_placement`]). Layer-first slots are dense *within* the + /// layer's storage group: `group_rank * tp_size + tp_rank`. Single-group + /// topologies degenerate to the historical `[layer][tp_rank]` grid. pub(crate) fn slot_index(&self, layer_id: usize, tp_rank: usize) -> Result { if layer_id >= self.num_layers() { return Err(EngineError::InvalidArgument(format!( @@ -215,7 +264,7 @@ impl LayerTopology { } Ok(match &self.page_layout { Some(p) => p.layer_shard[layer_id], - None => layer_id * self.tp_size + tp_rank, + None => self.layer_group_rank[layer_id] * self.tp_size + tp_rank, }) } } @@ -243,6 +292,9 @@ pub struct GpuContext { /// KV cache layouts by layer name. kv_caches: HashMap, + /// Hybrid-cache storage group id by layer name; absent = group 0. + layer_groups: HashMap, + /// CUDA context handle (kept alive for the lifetime of this context). _cuda_ctx: Arc, @@ -260,6 +312,10 @@ impl GpuContext { /// # Errors /// Returns `EngineError::CudaInit` if CUDA context creation or worker /// pool initialization fails. + #[allow( + clippy::too_many_arguments, + reason = "GPU context construction mirrors one registration payload" + )] fn new( cuda_ctx: Arc, device_id: i32, @@ -268,6 +324,7 @@ impl GpuContext { numa_node: NumaNode, transfer_mode: TransferMode, kv_caches: HashMap, + layer_groups: HashMap, ) -> Result { let worker_pool = GpuWorkerPool::spawn(device_id, numa_node, transfer_mode)?; @@ -277,6 +334,7 @@ impl GpuContext { pp_rank, preferred_numa: numa_node, kv_caches, + layer_groups, _cuda_ctx: cuda_ctx, worker_pool, }) @@ -311,6 +369,12 @@ impl GpuContext { pub(crate) fn worker_pool(&self) -> &GpuWorkerPool { &self.worker_pool } + + /// Hybrid-cache storage group of a layer; unregistered layers default to + /// group 0, preserving single-group behavior. + pub(crate) fn group_of_layer(&self, layer_name: &str) -> u32 { + self.layer_groups.get(layer_name).copied().unwrap_or(0) + } } pub(crate) struct GpuRegistration { @@ -320,6 +384,8 @@ pub(crate) struct GpuRegistration { pub(crate) numa_node: NumaNode, pub(crate) transfer_mode: TransferMode, pub(crate) kv_caches: HashMap, + /// Hybrid-cache storage group id by layer name; absent = group 0. + pub(crate) layer_groups: HashMap, } /// Instance context for a model inference process. @@ -483,6 +549,61 @@ impl InstanceContext { .map(|(id, name)| (name.clone(), id)) .collect(); + // Storage groups must agree across devices (same name ⇒ same group, + // checked while scanning) and be dense `0..N` so slot spaces are + // complete; a group nobody registers would never seal. + let mut group_by_name: HashMap<&str, u32> = HashMap::new(); + for gpu in gpus() { + for name in gpu.kv_caches.keys() { + let group = gpu.group_of_layer(name); + match group_by_name.insert(name, group) { + None => {} + Some(existing) if existing == group => {} + Some(existing) => { + return Err(EngineError::InvalidArgument(format!( + "layer {name} registered in storage group {existing} on one device \ + but group {group} on device {}", + gpu.device_id(), + ))); + } + } + } + } + let num_groups = 1 + group_by_name.values().copied().max().unwrap_or(0) as usize; + if group_by_name + .values() + .copied() + .collect::>() + .len() + != num_groups + { + return Err(EngineError::InvalidArgument(format!( + "storage groups must be dense 0..{num_groups}: an unregistered group \ + would never seal its blocks", + ))); + } + + let mut layer_group = vec![0u32; names.len()]; + let mut layer_group_rank = vec![0usize; names.len()]; + let mut group_layer_count = vec![0usize; num_groups]; + for name in &names { + let layer_id = name_to_id[name]; + let group = group_by_name[name.as_str()]; + layer_group[layer_id] = group; + layer_group_rank[layer_id] = group_layer_count[group as usize]; + group_layer_count[group as usize] += 1; + } + + // Page-first packs every layer of a block into one page slot per + // shard; splitting a block across storage groups contradicts that + // layout, so reject instead of silently picking one. + if self.page_first && num_groups > 1 { + return Err(EngineError::InvalidArgument(format!( + "page-first storage does not support multiple storage groups \ + ({num_groups} registered)", + ))); + } + // Validate registration completeness on the full `[layer][tp_rank]` // grid, independent of how slots are stored: every (layer, tp_rank) // pair needs an owner regardless of page-first collapsing. @@ -543,6 +664,9 @@ impl InstanceContext { Ok(LayerTopology { name_to_id, tp_size: self.tp_size, + layer_group, + group_layer_count, + layer_group_rank, page_layout, }) } @@ -555,6 +679,10 @@ impl InstanceContext { /// # Errors /// Returns `EngineError::InvalidArgument` for negative device IDs, /// or `EngineError::CudaInit` if CUDA context creation fails. + #[allow( + clippy::too_many_arguments, + reason = "GPU context construction mirrors one registration payload" + )] fn build_gpu_context( &self, device_id: i32, @@ -563,6 +691,7 @@ impl InstanceContext { numa_node: NumaNode, transfer_mode: TransferMode, kv_caches: HashMap, + layer_groups: HashMap, ) -> Result, EngineError> { if device_id < 0 { return Err(EngineError::InvalidArgument(format!( @@ -582,6 +711,7 @@ impl InstanceContext { numa_node, transfer_mode, kv_caches, + layer_groups, )?)) } @@ -635,6 +765,7 @@ impl InstanceContext { numa_node, transfer_mode, kv_caches, + layer_groups, } = registration; if tp_rank >= self.tp_size { @@ -661,6 +792,7 @@ impl InstanceContext { numa_node, transfer_mode, kv_caches, + layer_groups, )?; { diff --git a/pegaflow-core/src/instance/tests.rs b/pegaflow-core/src/instance/tests.rs index 6358dc93..9513ab21 100644 --- a/pegaflow-core/src/instance/tests.rs +++ b/pegaflow-core/src/instance/tests.rs @@ -22,6 +22,23 @@ fn gpu_registration(device_id: i32, tp_rank: usize, layers: &[&str]) -> GpuRegis gpu_registration_with_segment_bytes(device_id, tp_rank, layers, 1024) } +/// Build a `GpuRegistration` with explicit hybrid-cache group ids per layer. +/// `(name, group)` pairs; group 0 is the attention default, higher groups are +/// e.g. recurrent-state checkpoints in hybrid (mamba) models. +fn gpu_registration_with_groups( + device_id: i32, + tp_rank: usize, + layers_and_groups: &[(&str, u32)], +) -> GpuRegistration { + let layers: Vec<&str> = layers_and_groups.iter().map(|(name, _)| *name).collect(); + let mut registration = gpu_registration(device_id, tp_rank, &layers); + registration.layer_groups = layers_and_groups + .iter() + .map(|(name, group)| ((*name).to_string(), *group)) + .collect(); + registration +} + fn gpu_registration_with_segment_bytes( device_id: i32, tp_rank: usize, @@ -48,6 +65,7 @@ fn gpu_registration_with_segment_bytes( numa_node: NumaNode::UNKNOWN, transfer_mode: TransferMode::Direct, kv_caches, + layer_groups: HashMap::new(), } } @@ -448,3 +466,127 @@ fn seal_rejects_inconsistent_layer_geometry() { .expect_err("same name with different geometry must fail the seal"); assert!(err.to_string().contains("inconsistent geometry")); } + +/// Hybrid (mamba-style) topology: attention layers in group 0, recurrent-state +/// layers in group 1. Each group gets its own dense slot space, which is what +/// lets a group seal a block independently of the others — the precondition +/// for "recurrent saves only the final block" to ever become visible. +/// Commits device 0. +#[test] +fn hybrid_topology_seals_per_group_slot_spaces() { + // tp_size=1 keeps this single-device; the group-rank reset and per-group + // totals are what distinguish groups (tp_rank math is unchanged). + let instance = InstanceContext::new("hybrid".into(), "hybrid-ns".into(), 1, 1, false).unwrap(); + instance + .register_new_gpu(gpu_registration_with_groups( + 0, + 0, + &[ + ("attn_a", 0), + ("attn_b", 0), + ("recurrent_c", 1), + ("recurrent_d", 1), + ], + )) + .expect("register hybrid gpu"); + + let topology = instance.sealed_topology().expect("sealed"); + assert_eq!(topology.num_layers(), 4); + assert_eq!(topology.num_groups(), 2); + + // Sorted-name ids: attn_a=0, attn_b=1, recurrent_c=2, recurrent_d=3. + assert_eq!(topology.group_of_layer(0), 0); + assert_eq!(topology.group_of_layer(1), 0); + assert_eq!(topology.group_of_layer(2), 1); + assert_eq!(topology.group_of_layer(3), 1); + + // Per-group seal domains: group_slots = group_layers * tp_size. + assert_eq!(topology.group_total_slots(0).unwrap(), 2); + assert_eq!(topology.group_total_slots(1).unwrap(), 2); + // Union across groups stays the historical denominator. + assert_eq!(topology.total_slots(), 4); + + // Layer-first slots are dense WITHIN the group: group rank, not global id. + assert_eq!(topology.slot_index(0, 0).unwrap(), 0); // attn_a, group rank 0 + assert_eq!(topology.slot_index(1, 0).unwrap(), 1); // attn_b, group rank 1 + assert_eq!(topology.slot_index(2, 0).unwrap(), 0); // recurrent_c restarts at 0 + assert_eq!(topology.slot_index(3, 0).unwrap(), 1); +} + +/// With every layer in the default group 0 the within-group rank equals the +/// global layer id, so all existing slot math is bit-identical. This is the +/// backward-compat guard for current connectors. Commits device 0. +#[test] +fn default_groups_keep_global_slot_layout() { + let instance = + InstanceContext::new("classic".into(), "classic-ns".into(), 1, 1, false).unwrap(); + instance + .register_new_gpu(gpu_registration(0, 0, &["layer_a", "layer_b", "layer_c"])) + .expect("register classic gpu"); + + let topology = instance.sealed_topology().expect("sealed"); + assert_eq!(topology.num_groups(), 1); + assert_eq!(topology.group_total_slots(0).unwrap(), 3); + for layer_id in 0..3 { + assert_eq!(topology.group_of_layer(layer_id), 0); + assert_eq!(topology.slot_index(layer_id, 0).unwrap(), layer_id); + } + assert!(topology.group_total_slots(1).is_err()); +} + +/// A layer must live in the same storage group on every worker; disagreeing +/// devices would seal blocks with different slot counts per device silently. +/// Needs 2 CUDA devices. +#[test] +fn seal_rejects_inconsistent_layer_group_across_devices() { + if !has_cuda_devices(2) { + eprintln!( + "skipping seal_rejects_inconsistent_layer_group_across_devices: needs >= 2 CUDA devices" + ); + return; + } + + let instance = + InstanceContext::new("grp-conflict".into(), "grp-ns".into(), 1, 2, false).unwrap(); + instance + .register_new_gpu(gpu_registration_with_groups(0, 0, &[("layer_0", 0)])) + .expect("first worker"); + + let err = instance + .register_new_gpu(gpu_registration_with_groups(1, 0, &[("layer_0", 1)])) + .expect_err("same layer in different groups must fail the seal"); + assert!(err.to_string().contains("storage group"), "{err}"); +} + +/// Group ids must be dense `0..N`: a gap means a group nobody registers, whose +/// blocks would sit inflight forever waiting for slots that cannot arrive. +/// Commits device 0. +#[test] +fn seal_rejects_sparse_group_ids() { + let instance = InstanceContext::new("sparse".into(), "sparse-ns".into(), 1, 1, false).unwrap(); + let err = instance + .register_new_gpu(gpu_registration_with_groups( + 0, + 0, + &[("layer_a", 0), ("layer_b", 2)], + )) + .expect_err("group 1 with no layers must fail the seal"); + assert!(err.to_string().contains("dense"), "{err}"); +} + +/// Page-first storage packs all layers of a block into one page per shard; +/// splitting blocks across groups contradicts that. Reject instead of +/// silently picking one layout. Commits device 0. +#[test] +fn page_first_rejects_multiple_groups() { + let instance = + InstanceContext::new("page-hybrid".into(), "page-hybrid-ns".into(), 1, 1, true).unwrap(); + let err = instance + .register_new_gpu(gpu_registration_with_groups( + 0, + 0, + &[("layer_a", 0), ("layer_b", 1)], + )) + .expect_err("page-first with two storage groups must be rejected"); + assert!(err.to_string().contains("page-first"), "{err}"); +} diff --git a/pegaflow-core/src/lib.rs b/pegaflow-core/src/lib.rs index 0dc1a3eb..be7020c4 100644 --- a/pegaflow-core/src/lib.rs +++ b/pegaflow-core/src/lib.rs @@ -44,7 +44,7 @@ pub use internode::{ use layout::KVCacheLayout; pub use lease::QueryLeaseId; pub use pegaflow_common::NumaNode; -use pegaflow_common::NumaTopology; +use pegaflow_common::{NumaTopology, group_hash}; pub use pinned_pool::PinnedAllocation; pub use seal_offload::SlotMeta; pub use storage::{DEFAULT_RDMA_QPS_PER_PEER, MemoryCacheCleanupStats, StorageConfig}; @@ -281,6 +281,7 @@ impl PegaEngine { kv_stride_bytes_list, segments_list, None, + None, transfer_mode, page_first, ) @@ -315,6 +316,7 @@ impl PegaEngine { kv_stride_bytes_list: &[usize], segments_list: &[usize], block_stride_bytes_list: Option<&[usize]>, + layer_group_ids: Option<&[u32]>, transfer_mode: TransferMode, page_first: bool, ) -> Result<(), EngineError> { @@ -346,6 +348,14 @@ impl PegaEngine { strides.len() ))); } + if let Some(group_ids) = layer_group_ids + && group_ids.len() != batch_size + { + return Err(EngineError::InvalidArgument(format!( + "registration metadata length mismatch: layer_names={batch_size}, layer_group_ids={}", + group_ids.len() + ))); + } let mut kv_caches = HashMap::with_capacity(batch_size); for i in 0..batch_size { @@ -384,6 +394,15 @@ impl PegaEngine { } } + let layer_groups: HashMap = match layer_group_ids { + Some(group_ids) => layer_names + .iter() + .cloned() + .zip(group_ids.iter().copied()) + .collect(), + None => HashMap::new(), + }; + // Get or create instance let instance = self.get_or_create_instance(instance_id, namespace, tp_size, world_size, page_first)?; @@ -410,6 +429,7 @@ impl PegaEngine { numa_node, transfer_mode, kv_caches, + layer_groups, })?; info!( @@ -507,6 +527,84 @@ impl PegaEngine { Ok(status) } + /// All-or-nothing membership fetch over one hybrid-cache storage group, + /// eligible for the same SSD prefetch and MetaServer + RDMA remote fetch + /// as prefix queries. + /// + /// Where [`Self::query_group_membership`] answers from the resident read + /// cache only, this treats `block_hashes` as an exact want-set: misses are + /// pulled from remote tiers, and the query stays `Loading` until the whole + /// set is fetchable (the prefix machinery over an explicit key list *is* + /// a set fetch once the full length is required). Use it when partial + /// state is useless — e.g. restoring a recurrent-state checkpoint on a + /// prefill/decode handoff, where the peer that saved the set holds every + /// member. A `Ready` result with fewer blocks than requested means the + /// set could not be completed anywhere; callers treat that as a miss. + pub async fn query_group_membership_with_fetch( + &self, + instance_id: &str, + req_id: &str, + group_id: u32, + block_hashes: &[Vec], + ) -> Result { + let instance = self.get_instance(instance_id)?; + let topology = instance.sealed_topology()?; + // Same contract as the local membership query: an unknown group is a + // bug in the caller, not an all-miss answer. + topology.group_total_slots(group_id)?; + + let namespace = instance.namespace(); + let encoded: Vec> = block_hashes + .iter() + .map(|hash| group_hash(hash, group_id)) + .collect(); + + let status = self + .storage + .check_prefix_and_prefetch(req_id, namespace, &encoded, true) + .await; + + if let PrefetchStatus::Ready { blocks, missing } = &status { + let metrics = core_metrics(); + metrics.cache_block_hits.add(blocks.len() as u64, &[]); + if *missing > 0 { + metrics.cache_block_misses.add(*missing as u64, &[]); + } + } + + Ok(status) + } + + /// Position-aligned membership query over one hybrid-cache storage group. + /// + /// Unlike prefix queries, every position reports independently: entry `i` + /// is the sealed block for `block_hashes[i]` in `group_id`, or `None` on + /// miss. Sparse hit patterns are the point — callers (e.g. the vLLM + /// connector's hybrid reconcile) pick the rightmost hit themselves. + /// Group 0 keeps raw-hash keys, so classic instances observe no change. + /// + /// The returned blocks hold plain `Arc` refs, not leases: pin what you + /// need via [`Self::create_query_lease`]. + pub fn query_group_membership( + &self, + instance_id: &str, + group_id: u32, + block_hashes: &[Vec], + ) -> Result>>, EngineError> { + let instance = self.get_instance(instance_id)?; + let topology = instance.sealed_topology()?; + // Validate the group against the sealed topology; an unknown group is + // a contract bug, not an answer of "all miss". + topology.group_total_slots(group_id)?; + + let namespace = instance.namespace(); + let encoded: Vec> = block_hashes + .iter() + .map(|hash| group_hash(hash, group_id)) + .collect(); + Ok(self.storage.get_membership(namespace, &encoded)) + } + /// Create an opaque lease that owns query-ready blocks. pub fn create_query_lease( &self, @@ -642,6 +740,32 @@ impl PegaEngine { )); } + // Resolve each load group's storage group and require homogeneity: a + // group's blocks seal against exactly one slot space. An empty load + // group owns no storage group; targets pointing at it load into + // nothing (it exists only to keep group indices aligned with the + // connector's cache-group layout). + let mut storage_group_of_load_group: Vec> = + Vec::with_capacity(layer_groups.len()); + for layer_names in layer_groups { + let mut group: Option = None; + for layer_name in layer_names { + let layer_id = topology.layer_id(layer_name)?; + let layer_group = topology.group_of_layer(layer_id); + match group { + None => group = Some(layer_group), + Some(existing) if existing == layer_group => {} + Some(existing) => { + return Err(EngineError::InvalidArgument(format!( + "load group mixes storage groups {existing} and {layer_group} \ + (layer {layer_name})", + ))); + } + } + } + storage_group_of_load_group.push(group); + } + // Consume query leases reserved for this load. trace_scope!("load.cache_lookup", _s); let mut block_targets_by_group = vec![Vec::new(); layer_groups.len()]; @@ -668,6 +792,28 @@ impl PegaEngine { group_index ))); } + // A stored block must carry exactly the slot layout of the + // storage group this target group loads into. A mismatch + // means the namespace is shared by instances with different + // layer sets (e.g. MTP enabled vs disabled) — loading would + // silently leave layers uninitialized, so fail loudly and + // let vLLM recompute. Empty load groups own no storage + // group and skip the check entirely. + if let Some(storage_group) = storage_group_of_load_group[group_index] { + let expected_slots = topology.group_total_slots(storage_group)?; + for (source_index, destination) in lease_block_targets.iter().enumerate() { + if destination.is_some() + && blocks[source_index].slots().len() != expected_slots + { + return Err(EngineError::InvalidArgument(format!( + "stored block has {} slots but storage group {storage_group} of \ + instance {instance_id} expects {expected_slots}: \ + namespace is shared by incompatible KV layouts", + blocks[source_index].slots().len(), + ))); + } + } + } block_targets_by_group[group_index].extend( lease_block_targets.iter().enumerate().filter_map( |(source_index, destination)| { @@ -676,21 +822,6 @@ impl PegaEngine { ), ); } - // A stored block must carry exactly this instance's slot layout. - // A mismatch means the namespace is shared by instances with - // different layer sets (e.g. MTP enabled vs disabled) — loading - // would silently leave layers uninitialized, so fail loudly and - // let vLLM recompute. - for block in &blocks { - if block.slots().len() != topology.total_slots() { - return Err(EngineError::InvalidArgument(format!( - "stored block has {} slots but instance {instance_id} expects {}: \ - namespace is shared by incompatible KV layouts", - block.slots().len(), - topology.total_slots() - ))); - } - } block_cache.extend(blocks); } trace_drop!(_s); diff --git a/pegaflow-core/src/offload.rs b/pegaflow-core/src/offload.rs index 7cf938a4..43106a2f 100644 --- a/pegaflow-core/src/offload.rs +++ b/pegaflow-core/src/offload.rs @@ -6,7 +6,7 @@ // is deferred to the storage insert worker via `RawSaveBatch`. // ============================================================================ -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::num::NonZeroU64; use std::sync::Arc; @@ -21,7 +21,7 @@ use crate::layout::KVCacheLayout; use crate::metrics::core_metrics; use crate::pinned_pool::PinnedAllocation; use crate::{EngineError, PegaEngine}; -use pegaflow_common::NumaNode; +use pegaflow_common::{NumaNode, group_hash}; // ============================================================================ // Types sent to the insert worker (deferred Phase 4) @@ -56,6 +56,8 @@ struct LayerContext { layer_name: String, layout: KVCacheLayout, slot_id: usize, + /// Hybrid-cache storage group of this layer; keyed and sealed per group. + group: u32, /// Blocks to save: (block_idx, hash). Filtered in Phase 1. blocks_to_save: Vec<(usize, Vec)>, /// Host-side block images, parallel to `blocks_to_save`. @@ -138,6 +140,35 @@ fn build_hashed_insert_entries(namespace: String, layers: Vec) -> .collect() } +/// Drop `(group, hash)` candidates whose group-encoded key already exists in +/// the read cache. Filtering is per group because groups key the same content +/// hash independently (e.g. attention block vs. recurrent checkpoint). +fn filter_new_hashes_per_group( + storage: &crate::storage::StorageEngine, + namespace: &str, + candidates: HashSet<(u32, Vec)>, +) -> HashSet<(u32, Vec)> { + let mut by_group: HashMap>> = HashMap::new(); + for (group, hash) in &candidates { + by_group.entry(*group).or_default().push(hash.clone()); + } + + let mut survivors = HashSet::with_capacity(candidates.len()); + for (group, raw_hashes) in by_group { + let mut encoded: HashSet> = raw_hashes + .iter() + .map(|hash| group_hash(hash, group)) + .collect(); + storage.filter_hashes_not_in_cache_inplace(namespace, &mut encoded); + for raw in raw_hashes { + if encoded.contains(&group_hash(&raw, group)) { + survivors.insert((group, raw)); + } + } + } + survivors +} + // ============================================================================ // Save methods on PegaEngine (moved from lib.rs) // ============================================================================ @@ -275,7 +306,6 @@ impl PegaEngine { let instance = self.get_instance(instance_id)?; let topology = instance.sealed_topology()?; let namespace = instance.namespace().to_string(); - let total_slots = topology.total_slots(); // ── Phase 0: Resolve per-layer metadata and build valid_blocks ── trace_scope!("save.resolve_metadata", _s); @@ -283,7 +313,9 @@ impl PegaEngine { let gpu_context = instance.get_gpu_for_save_group(device_id, tp_rank, pp_rank)?; let mut layer_contexts: Vec = Vec::with_capacity(saves.len()); - let mut hashes_to_save: HashSet> = HashSet::new(); + // Dedupe candidates per (group, hash): the same content hash may exist + // in the attention group while the recurrent group still needs it. + let mut hashes_to_save: HashSet<(u32, Vec)> = HashSet::new(); for LayerSave { layer_name, @@ -309,6 +341,7 @@ impl PegaEngine { })?; let slot_id = topology.slot_index(layer_id, tp_rank)?; + let group = topology.group_of_layer(layer_id); let num_blocks = layout.num_blocks(); @@ -331,13 +364,14 @@ impl PegaEngine { } for (_, hash) in &blocks_to_save { - hashes_to_save.insert(hash.clone()); + hashes_to_save.insert((group, hash.clone())); } layer_contexts.push(LayerContext { layer_name, layout, slot_id, + group, blocks_to_save, raw_blocks: Vec::new(), }); @@ -356,20 +390,23 @@ impl PegaEngine { trace_scope!("save.hash_filter", _s); - // Single in-place cache filter for all unique hashes - self.storage - .filter_hashes_not_in_cache_inplace(&namespace, &mut hashes_to_save); + let hashes_to_save = filter_new_hashes_per_group( + &self.storage, + &namespace, + std::mem::take(&mut hashes_to_save), + ); if hashes_to_save.is_empty() { trace_drop!(_s); return Ok(()); } - // Per-layer filter: keep only blocks whose hash needs saving + // Per-layer filter: keep only blocks whose (group, hash) needs saving let mut total_blocks_to_save = 0usize; for ctx in &mut layer_contexts { + let group = ctx.group; ctx.blocks_to_save - .retain(|(_, hash)| hashes_to_save.contains(hash.as_slice())); + .retain(|(_, hash)| hashes_to_save.contains(&(group, hash.clone()))); total_blocks_to_save += ctx.blocks_to_save.len(); } // Remove layers with no blocks to save @@ -550,10 +587,15 @@ impl PegaEngine { batch_start.elapsed().as_secs_f64() * 1000.0 ); - // Build RawSaveBatch and send to insert worker (fire-and-forget). - // Page-first collapses a shard's layers into one page slot per block, so - // a whole shard is a single RawSaveLayer regardless of its layer count. - let raw_layers: Vec = if page_first { + // Build per-group RawSaveBatches and send to the insert worker + // (fire-and-forget). Each group seals blocks against its own slot + // count and stores them under group-encoded hashes (group 0 keeps raw + // hashes, so single-group instances produce byte-identical keys). + if page_first { + // Page-first collapses a shard's layers into one page slot per + // block, so a whole shard is a single RawSaveLayer regardless of + // its layer count. Sealing rejects multi-group page-first, so the + // whole batch is group 0. let slot_id = layer_contexts[0].slot_id; let page_size = topology .shard_page_size(slot_id) @@ -570,37 +612,48 @@ impl PegaEngine { RawBlock::single_segment(Segment::new(ptr, page_size, page)) }) .collect(); - vec![RawSaveLayer { - slot_id, - padded_block_size: page_size, - blocks, - block_hashes, - }] + self.storage.send_raw_insert(RawSaveBatch { + namespace, + total_slots: topology.total_slots(), + numa_node: save_numa_node, + layers: vec![RawSaveLayer { + slot_id, + padded_block_size: page_size, + blocks, + block_hashes, + }], + }); } else { - layer_contexts - .into_iter() - .map(|ctx| { - let block_hashes: Vec> = ctx - .blocks_to_save - .into_iter() - .map(|(_, hash)| hash) - .collect(); - RawSaveLayer { - slot_id: ctx.slot_id, - padded_block_size: ctx.layout.padded_block_bytes(), - blocks: ctx.raw_blocks, - block_hashes, - } - }) - .collect() - }; - - self.storage.send_raw_insert(RawSaveBatch { - namespace, - total_slots, - numa_node: save_numa_node, - layers: raw_layers, - }); + let mut by_group: std::collections::BTreeMap> = + Default::default(); + for ctx in layer_contexts { + by_group.entry(ctx.group).or_default().push(ctx); + } + for (group, contexts) in by_group { + let raw_layers: Vec = contexts + .into_iter() + .map(|ctx| { + let block_hashes: Vec> = ctx + .blocks_to_save + .into_iter() + .map(|(_, hash)| group_hash(&hash, group)) + .collect(); + RawSaveLayer { + slot_id: ctx.slot_id, + padded_block_size: ctx.layout.padded_block_bytes(), + blocks: ctx.raw_blocks, + block_hashes, + } + }) + .collect(); + self.storage.send_raw_insert(RawSaveBatch { + namespace: namespace.clone(), + total_slots: topology.group_total_slots(group)?, + numa_node: save_numa_node, + layers: raw_layers, + }); + } + } Ok(()) } diff --git a/pegaflow-core/src/storage/mod.rs b/pegaflow-core/src/storage/mod.rs index 857ee8f2..632a5f57 100644 --- a/pegaflow-core/src/storage/mod.rs +++ b/pegaflow-core/src/storage/mod.rs @@ -5,8 +5,6 @@ pub(crate) mod transfer_lock; mod write_path; use bytesize::ByteSize; -#[cfg(not(feature = "rdma"))] -use log::warn; use log::{debug, info, warn}; use std::collections::HashSet; use std::num::NonZeroU64; @@ -380,6 +378,21 @@ impl StorageEngine { } } + /// Position-aligned membership lookup in the resident read cache: entry + /// `i` is the sealed block for `hashes[i]`, or `None` on miss. Hashes must + /// already carry any group encoding (see `group_hash`). + pub(crate) fn get_membership( + &self, + namespace: &str, + hashes: &[Vec], + ) -> Vec>> { + let keys: Vec = hashes + .iter() + .map(|hash| BlockKey::new(namespace.to_string(), hash.clone())) + .collect(); + self.read_cache.get_blocks_aligned(&keys) + } + /// Evict all blocks from the resident in-memory read cache. /// /// This preserves backing-store copies. Blocks with diff --git a/pegaflow-core/src/storage/read_cache.rs b/pegaflow-core/src/storage/read_cache.rs index 1e3a6a01..325f5e05 100644 --- a/pegaflow-core/src/storage/read_cache.rs +++ b/pegaflow-core/src/storage/read_cache.rs @@ -140,6 +140,22 @@ impl ReadCache { found } + /// Position-aligned membership: entry `i` is the block for `keys[i]`, or + /// `None` on miss. Unlike [`Self::get_prefix_blocks`] this never stops at + /// the first gap — hybrid-cache checkpoint groups (recurrent state) have + /// sparse hit patterns by design, where the caller picks the rightmost + /// hit instead of a prefix. + pub(super) fn get_blocks_aligned(&self, keys: &[BlockKey]) -> Vec>> { + let mut inner = self.inner.lock(); + keys.iter() + .map(|key| { + inner.cache.get(key).inspect(|_| { + refresh_recency(&mut inner, key); + }) + }) + .collect() + } + pub(super) fn remove_lru_batch(&self, batch_size: usize) -> Vec<(BlockKey, Arc)> { let removed = { let mut inner = self.inner.lock(); diff --git a/pegaflow-core/tests/common/harness.rs b/pegaflow-core/tests/common/harness.rs index dfa2ef13..779b59a8 100644 --- a/pegaflow-core/tests/common/harness.rs +++ b/pegaflow-core/tests/common/harness.rs @@ -98,6 +98,11 @@ impl TestGpuData { assert_eq!(self.gpu.copy_to_host(), expected, "GPU data mismatch"); } + /// The deterministic host-side pattern backing this GPU buffer. + pub fn expected_bytes(&self) -> &[u8] { + &self.expected + } + pub fn ptr(&self) -> u64 { self.gpu.as_u64() } @@ -131,6 +136,7 @@ struct LayerSpec { block_size: usize, kv_stride: usize, segments: usize, + group: u32, } pub struct TestEnvBuilder { @@ -172,6 +178,27 @@ impl TestEnvBuilder { block_size, kv_stride: 0, segments: 1, + group: 0, + }); + self + } + + /// Add a contiguous layer in an explicit hybrid-cache storage group + /// (e.g. group 1 = recurrent-state checkpoints in mamba-style models). + pub fn grouped_layer( + mut self, + name: &'static str, + num_blocks: usize, + block_size: usize, + group: u32, + ) -> Self { + self.layers.push(LayerSpec { + name, + num_blocks, + block_size, + kv_stride: 0, + segments: 1, + group, }); self } @@ -198,6 +225,7 @@ impl TestEnvBuilder { block_size: segment_size, kv_stride, segments: 2, + group: 0, }); self } @@ -261,6 +289,7 @@ impl TestEnvBuilder { block_size: spec.block_size, kv_stride: spec.kv_stride, segments: spec.segments, + group: spec.group, }) .collect(); diff --git a/pegaflow-core/tests/common/helpers.rs b/pegaflow-core/tests/common/helpers.rs index 3b047bbd..3f24920a 100644 --- a/pegaflow-core/tests/common/helpers.rs +++ b/pegaflow-core/tests/common/helpers.rs @@ -21,6 +21,8 @@ pub struct LayerInfo { pub block_size: usize, pub kv_stride: usize, pub segments: usize, + /// Hybrid-cache storage group id; 0 is the default attention group. + pub group: u32, } #[allow( @@ -47,9 +49,10 @@ pub fn register_layers( let block_sizes: Vec = layers.iter().map(|l| l.block_size).collect(); let kv_strides: Vec = layers.iter().map(|l| l.kv_stride).collect(); let segments: Vec = layers.iter().map(|l| l.segments).collect(); + let group_ids: Vec = layers.iter().map(|l| l.group).collect(); engine - .register_context_layer_batch( + .register_context_layer_batch_strided( instance_id, namespace, device_id, @@ -64,6 +67,8 @@ pub fn register_layers( &block_sizes, &kv_strides, &segments, + None, + Some(&group_ids), transfer_mode, page_first, ) diff --git a/pegaflow-core/tests/hybrid_groups.rs b/pegaflow-core/tests/hybrid_groups.rs new file mode 100644 index 00000000..7fca17ad --- /dev/null +++ b/pegaflow-core/tests/hybrid_groups.rs @@ -0,0 +1,323 @@ +//! Hybrid-cache (mamba/HMA) storage groups. +//! +//! Attention groups save every block; the recurrent group saves only +//! checkpoint blocks (the state *after* a block's tokens). Each group seals +//! blocks against its own slot space, so a final-block-only recurrent save +//! becomes visible without waiting for attention's per-block cadence — and +//! one lease per group loads both halves back into the GPU. + +mod common; + +use common::*; +use pegaflow_core::*; + +/// The HMA bull's-eye: recurrent group seals a tail-only checkpoint save, +/// queries isolate groups by hash, and a two-lease load restores both the +/// attention prefix and the recurrent state block. +#[tokio::test] +async fn recurrent_group_seals_final_block_save() { + if !has_cuda_devices(1) { + eprintln!("skipping recurrent_group_seals_final_block_save: needs >= 1 CUDA device"); + return; + } + + const ATTN_BLOCK: usize = 1024; + const RECUR_BLOCK: usize = 2048; + + let env = TestEnvBuilder::new("hybrid-hma", "hybrid-hma-ns") + .grouped_layer("attn_0", 8, ATTN_BLOCK, 0) + .grouped_layer("attn_1", 8, ATTN_BLOCK, 0) + .grouped_layer("recurrent_state", 8, RECUR_BLOCK, 1) + .pool_size(64 << 20) + .build(); + + // Same hash namespace for both groups on purpose: group encoding, not the + // connector, must keep attention and recurrent blocks apart. + let hashes = make_block_hashes(8, 10); + let prefix_hashes = hashes[0..3].to_vec(); + + // Attention saves blocks 0..2 for both layers; the recurrent group saves + // ONLY the block-2 checkpoint (state after block 2's tokens), mirroring + // the connector's final-save. The recurrent checkpoint physically lives + // in GPU block 5 — state blocks are per-request slots, not prefix cells. + env.engine + .batch_save_kv_blocks_from_ipc( + &env.instance_id, + 0, + 0, + 0, + vec![ + LayerSave { + layer_name: "attn_0".into(), + block_ids: vec![0, 1, 2], + block_hashes: prefix_hashes.clone(), + }, + LayerSave { + layer_name: "attn_1".into(), + block_ids: vec![0, 1, 2], + block_hashes: prefix_hashes.clone(), + }, + ], + ) + .await + .expect("save attention prefix"); + env.engine + .batch_save_kv_blocks_from_ipc( + &env.instance_id, + 0, + 0, + 0, + vec![LayerSave { + layer_name: "recurrent_state".into(), + block_ids: vec![5], + block_hashes: vec![hashes[2].clone()], + }], + ) + .await + .expect("save recurrent checkpoint"); + env.engine.flush_saves().await; + + // Membership is per group: group 0 holds all three prefix blocks... + let attn_hits = env + .engine + .query_group_membership(&env.instance_id, 0, &prefix_hashes) + .expect("query group 0 membership"); + assert_eq!( + attn_hits.iter().map(|b| b.is_some()).collect::>(), + vec![true, true, true] + ); + + // ...while group 1 holds ONLY the block-2 checkpoint, even though the + // identical hash bytes for blocks 0/1 exist in group 0 (isolation). + let recur_hits = env + .engine + .query_group_membership(&env.instance_id, 1, &prefix_hashes) + .expect("query group 1 membership"); + assert_eq!( + recur_hits.iter().map(|b| b.is_some()).collect::>(), + vec![false, false, true], + "recurrent group must seal the tail-only checkpoint" + ); + + // The classic prefix query (group 0) is untouched by the recurrent save. + match env.query(&prefix_hashes).await { + PrefetchStatus::Ready { blocks, missing } => { + assert_eq!(blocks.len(), 3); + assert_eq!(missing, 0); + } + other => panic!("expected Ready, got {other:?}"), + } + + // Reconcile (connector-side): rightmost recurrent checkpoint within the + // attention prefix → hit = 2 + 1 = 3 blocks. + let attn_lease = match env.query(&prefix_hashes).await { + PrefetchStatus::Ready { blocks, .. } => env + .engine + .create_query_lease(&env.instance_id, blocks) + .expect("attention lease"), + other => panic!("expected Ready, got {other:?}"), + }; + let recur_block = recur_hits[2] + .as_ref() + .expect("rightmost recurrent checkpoint") + .clone(); + let recur_lease = env + .engine + .create_query_lease(&env.instance_id, vec![recur_block]) + .expect("recurrent lease"); + + // Snapshot expected bytes, then wipe GPU memory so the load proves the + // roundtrip instead of re-reading its own source. + let attn_expected: Vec> = env.layers[0..2] + .iter() + .map(|l| l.data.expected_bytes().to_vec()) + .collect(); + let recur_expected = env.layers[2].data.expected_bytes().to_vec(); + for layer in &env.layers { + layer.data.zero_gpu(); + } + + // One load, two leases, two layer groups: attention restores the [0,3) + // prefix; the recurrent group has exactly one physical target (the + // request's live state slot, block 7) and None elsewhere. + let layer_groups: Vec> = vec![vec!["attn_0", "attn_1"], vec!["recurrent_state"]]; + let load_state = LoadState::new().expect("create LoadState"); + let shm_name = load_state.shm_name().to_string(); + env.engine + .batch_load_kv_blocks_multi_layer( + &env.instance_id, + 0, + 0, + &shm_name, + &layer_groups, + &[ + ( + attn_lease, + vec![vec![Some(0), Some(1), Some(2)], vec![None, None, None]], + ), + (recur_lease, vec![vec![None], vec![Some(7)]]), + ], + ) + .expect("submit hybrid load"); + wait_for_load(&load_state, LOAD_WAIT_TIMEOUT).await; + + // Attention blocks 0..2 restored, the rest still zero. + for (layer_idx, expected) in attn_expected.iter().enumerate() { + let mut want = vec![0u8; expected.len()]; + want[..3 * ATTN_BLOCK].copy_from_slice(&expected[..3 * ATTN_BLOCK]); + env.layers[layer_idx].data.assert_gpu_matches(&want); + } + // Recurrent: only block 7 holds the checkpoint copied from block 5. + let mut want = vec![0u8; recur_expected.len()]; + want[7 * RECUR_BLOCK..8 * RECUR_BLOCK] + .copy_from_slice(&recur_expected[5 * RECUR_BLOCK..6 * RECUR_BLOCK]); + env.layers[2].data.assert_gpu_matches(&want); +} + +/// Group 0 must behave bit-identically to classic single-group storage when +/// nothing declares a nonzero group: prefix query, per-hash isolation, and +/// slot counts all follow the historical layout. +#[tokio::test] +async fn single_group_instances_are_unaffected() { + if !has_cuda_devices(1) { + eprintln!("skipping single_group_instances_are_unaffected: needs >= 1 CUDA device"); + return; + } + + let env = TestEnvBuilder::new("hybrid-compat", "hybrid-compat-ns") + .layer("layer_0", 4, 1024) + .layer("layer_1", 4, 1024) + .build(); + + let hashes = make_block_hashes(4, 7); + env.save_and_wait(&hashes[0..2]).await; + + let hits = env + .engine + .query_group_membership(&env.instance_id, 0, &hashes[0..4]) + .expect("membership"); + assert_eq!( + hits.iter().map(|b| b.is_some()).collect::>(), + vec![true, true, false, false] + ); + + // A nonsensical group on a single-group instance is a hard error, not an + // all-miss answer. + assert!( + env.engine + .query_group_membership(&env.instance_id, 7, &hashes[0..1]) + .is_err() + ); +} + +/// The prefill/decode handoff shape: `query_group_membership_with_fetch` +/// treats the hash list as an exact want-set with all-or-nothing semantics. +/// A fully saved recurrent checkpoint set answers Ready and complete; a +/// want-set with any absent member (no remote tier configured here) comes +/// back short, which callers must read as a miss. Group isolation holds: a +/// hash saved only under the attention group never satisfies the recurrent +/// want-set. +#[tokio::test] +async fn recurrent_membership_fetch_is_all_or_nothing() { + if !has_cuda_devices(1) { + eprintln!("skipping recurrent_membership_fetch_is_all_or_nothing: needs >= 1 CUDA device"); + return; + } + + const ATTN_BLOCK: usize = 1024; + const RECUR_BLOCK: usize = 2048; + + let env = TestEnvBuilder::new("hybrid-pd", "hybrid-pd-ns") + .grouped_layer("attn_0", 8, ATTN_BLOCK, 0) + .grouped_layer("recurrent_state", 8, RECUR_BLOCK, 1) + .pool_size(64 << 20) + .build(); + + let hashes = make_block_hashes(8, 21); + + // Attention saves blocks 0..2; the recurrent group saves the block-1 and + // block-2 checkpoints (a two-member set, like a multi-component state). + env.engine + .batch_save_kv_blocks_from_ipc( + &env.instance_id, + 0, + 0, + 0, + vec![LayerSave { + layer_name: "attn_0".into(), + block_ids: vec![0, 1, 2], + block_hashes: hashes[0..3].to_vec(), + }], + ) + .await + .expect("save attention prefix"); + env.engine + .batch_save_kv_blocks_from_ipc( + &env.instance_id, + 0, + 0, + 0, + vec![LayerSave { + layer_name: "recurrent_state".into(), + block_ids: vec![4, 5], + block_hashes: vec![hashes[1].clone(), hashes[2].clone()], + }], + ) + .await + .expect("save recurrent checkpoints"); + env.engine.flush_saves().await; + + // The exact want-set resolves Ready and complete, in requested order. + let want = vec![hashes[1].clone(), hashes[2].clone()]; + match env + .engine + .query_group_membership_with_fetch(&env.instance_id, "pd-req-full", 1, &want) + .await + .expect("want-set query") + { + PrefetchStatus::Ready { blocks, missing } => { + assert_eq!(blocks.len(), 2, "full want-set must resolve completely"); + assert_eq!(missing, 0); + let _lease = env + .engine + .create_query_lease(&env.instance_id, blocks) + .expect("lease over the fetched set"); + } + other => panic!("expected Ready, got {other:?}"), + } + + // A want-set with an absent member comes back short (no remote tier is + // configured in this harness): the caller must treat that as a miss. + let short = vec![hashes[1].clone(), hashes[5].clone(), hashes[2].clone()]; + match env + .engine + .query_group_membership_with_fetch(&env.instance_id, "pd-req-short", 1, &short) + .await + .expect("short want-set query") + { + PrefetchStatus::Ready { blocks, missing } => { + assert!( + blocks.len() < short.len(), + "an incomplete want-set must not report as complete" + ); + assert!(missing > 0); + } + other => panic!("expected Ready, got {other:?}"), + } + + // Group isolation: hashes[0] exists in the attention group only, so the + // recurrent want-set containing it cannot complete. + let cross = vec![hashes[0].clone()]; + match env + .engine + .query_group_membership_with_fetch(&env.instance_id, "pd-req-cross", 1, &cross) + .await + .expect("cross-group query") + { + PrefetchStatus::Ready { blocks, missing } => { + assert_eq!(blocks.len(), 0); + assert_eq!(missing, 1); + } + other => panic!("expected Ready, got {other:?}"), + } +} diff --git a/pegaflow-proto/proto/engine.proto b/pegaflow-proto/proto/engine.proto index 3a039433..e594377a 100644 --- a/pegaflow-proto/proto/engine.proto +++ b/pegaflow-proto/proto/engine.proto @@ -46,6 +46,13 @@ message RegisterContextRequest { // Set by the connector for MLA; the engine derives per-layer page offsets // from the registered layouts. See spec.md. bool page_first = 18; + // Hybrid-cache storage group per layer, parallel to layer_names (e.g. + // vLLM's kv-cache group index: 0 = full attention, 1 = recurrent state). + // Empty means every layer belongs to group 0, preserving single-group + // behavior. Groups key blocks independently and seal against their own + // slot space, so a sparse (checkpoint-only) group can seal without + // attention's per-block cadence. Page-first stays single-group. + repeated uint32 layer_group_ids = 19; } message RegisterContextResponse { @@ -110,8 +117,20 @@ message QueryRequest { string req_id = 3; // Keep the query in Loading until the full remaining prefix is fetchable // from a remote node via MetaServer + RDMA. Does not observe saves landing - // in the local engine; no effect when RDMA is not configured. + // in the local engine; no effect when RDMA is not configured. On groups > 0 + // this switches the membership query to an all-or-nothing set fetch: the + // hash list is an exact want-set, misses are pulled from SSD / remote + // peers, and a Ready answer shorter than the request means the set could + // not be completed anywhere: callers treat it as a miss, and it carries no + // lease (the short hit count is reported for observability only). bool wait_for_full_prefix = 4; + // Hybrid-cache storage group to query. 0 (default) is the attention group + // and uses prefix semantics (hit count + lease over the contiguous prefix, + // SSD/RDMA prefetch eligible). Groups > 0 use membership semantics: every + // position reports independently (QueryReady.hit_positions), hits are not + // prefix-clamped, and the lease pins exactly the hit blocks; local-only + // unless wait_for_full_prefix requests the remote-fetching set form. + uint32 group_id = 5; } message QueryResponse { @@ -124,8 +143,18 @@ message QueryResponse { message QueryLoading {} message QueryReady { + // Prefix query (group_id = 0): number of leading hit blocks. + // Membership query (group_id > 0): number of hit positions in total. uint64 num_hit_blocks = 1; + // Lease over the hit blocks: the prefix for group 0, or exactly the blocks + // at hit_positions (in that order) for membership queries. Empty on 0 hits. bytes lease = 2; + // Membership queries only: indices into QueryRequest.block_hashes whose + // block is cached, in ascending order. The lease's i-th block corresponds + // to query position hit_positions[i]. Position-aligned so hybrid callers + // can pick e.g. the rightmost recurrent checkpoint within the attention + // prefix without a second round trip. + repeated uint32 hit_positions = 3; } message ReleaseRequest { diff --git a/pegaflow-server/src/service.rs b/pegaflow-server/src/service.rs index b7d8a3c5..c34f1d7e 100644 --- a/pegaflow-server/src/service.rs +++ b/pegaflow-server/src/service.rs @@ -301,7 +301,12 @@ impl Engine for GrpcEngineService { } // Call engine batch registration - if let Err(err) = self.engine.register_context_layer_batch( + let layer_group_ids: Option<&[u32]> = if req.layer_group_ids.is_empty() { + None + } else { + Some(&req.layer_group_ids) + }; + if let Err(err) = self.engine.register_context_layer_batch_strided( &req.instance_id, &req.namespace, req.device_id, @@ -316,6 +321,8 @@ impl Engine for GrpcEngineService { &bytes_per_block_list, &kv_stride_bytes_list, &segments_list, + None, + layer_group_ids, transfer_mode, req.page_first, ) { @@ -553,46 +560,122 @@ impl Engine for GrpcEngineService { req.block_hashes.len() ); - // SSD prefetch-aware query - let status = self - .engine - .count_prefix_hit_blocks_with_prefetch( - &req.instance_id, - &req.req_id, - &req.block_hashes, - req.wait_for_full_prefix, - ) - .await - .map_err(Self::map_engine_error)?; - - let outcome = match status { - PrefetchStatus::Ready { blocks, missing } => { - let hit = blocks.len(); - if let Ok(mut t) = self.hll_tracker.lock() { - t.record_hashes(&req.block_hashes); + let outcome = if req.group_id > 0 && req.wait_for_full_prefix { + // All-or-nothing membership fetch: the hash list is an exact + // want-set and misses are pulled from SSD / remote peers, so + // the query may report Loading before it resolves. A Ready + // answer shorter than the request means the set could not be + // completed anywhere and the caller must treat it as a miss. + let status = self + .engine + .query_group_membership_with_fetch( + &req.instance_id, + &req.req_id, + req.group_id, + &req.block_hashes, + ) + .await + .map_err(Self::map_engine_error)?; + match status { + PrefetchStatus::Ready { blocks, .. } => { + // All-or-nothing: a short answer is a miss by + // contract, so it must not pin the partial blocks — + // a lease would hold them non-evictable for the full + // lease TTL with no consumer. The short hit count is + // still reported for observability; only a complete + // want-set carries a lease. + let complete = blocks.len() == req.block_hashes.len(); + let positions: Vec = (0..blocks.len() as u32).collect(); + let lease = if complete && !blocks.is_empty() { + self.engine + .create_query_lease(&req.instance_id, blocks) + .map_err(Self::map_engine_error)? + .to_bytes() + .to_vec() + } else { + Vec::new() + }; + query_response::Outcome::Ready(QueryReady { + num_hit_blocks: positions.len() as u64, + lease, + hit_positions: positions, + }) } - let lease = if hit == 0 { - Vec::new() - } else { - self.engine - .create_query_lease(&req.instance_id, blocks) - .map_err(Self::map_engine_error)? - .to_bytes() - .to_vec() - }; - debug!( - "RPC [query_prefetch] ready: instance_id={} hit={} missing={} lease={}", - req.instance_id, - hit, - missing, - !lease.is_empty() - ); - query_response::Outcome::Ready(QueryReady { - num_hit_blocks: hit as u64, - lease, - }) + PrefetchStatus::Loading => query_response::Outcome::Loading(QueryLoading {}), + } + } else if req.group_id > 0 { + // Membership query (hybrid-cache checkpoint groups): every + // position reports independently and the lease pins exactly + // the hit blocks, in hit_positions order. + let hits = self + .engine + .query_group_membership(&req.instance_id, req.group_id, &req.block_hashes) + .map_err(Self::map_engine_error)?; + let mut positions = Vec::new(); + let mut blocks = Vec::new(); + for (pos, block) in hits.into_iter().enumerate() { + if let Some(block) = block { + positions.push(pos as u32); + blocks.push(block); + } + } + let lease = if blocks.is_empty() { + Vec::new() + } else { + self.engine + .create_query_lease(&req.instance_id, blocks) + .map_err(Self::map_engine_error)? + .to_bytes() + .to_vec() + }; + query_response::Outcome::Ready(QueryReady { + num_hit_blocks: positions.len() as u64, + lease, + hit_positions: positions, + }) + } else { + // SSD prefetch-aware query + let status = self + .engine + .count_prefix_hit_blocks_with_prefetch( + &req.instance_id, + &req.req_id, + &req.block_hashes, + req.wait_for_full_prefix, + ) + .await + .map_err(Self::map_engine_error)?; + + match status { + PrefetchStatus::Ready { blocks, missing } => { + let hit = blocks.len(); + if let Ok(mut t) = self.hll_tracker.lock() { + t.record_hashes(&req.block_hashes); + } + let lease = if hit == 0 { + Vec::new() + } else { + self.engine + .create_query_lease(&req.instance_id, blocks) + .map_err(Self::map_engine_error)? + .to_bytes() + .to_vec() + }; + debug!( + "RPC [query_prefetch] ready: instance_id={} hit={} missing={} lease={}", + req.instance_id, + hit, + missing, + !lease.is_empty() + ); + query_response::Outcome::Ready(QueryReady { + num_hit_blocks: hit as u64, + lease, + hit_positions: Vec::new(), + }) + } + PrefetchStatus::Loading => query_response::Outcome::Loading(QueryLoading {}), } - PrefetchStatus::Loading => query_response::Outcome::Loading(QueryLoading {}), }; Ok(Response::new(QueryResponse { @@ -993,6 +1076,7 @@ mod tests { block_hashes: Vec::new(), req_id: String::new(), wait_for_full_prefix: false, + group_id: 0, }) .expect_err("empty req_id must be rejected before engine lookup"); @@ -1007,6 +1091,7 @@ mod tests { block_hashes: Vec::new(), req_id: "request".to_string(), wait_for_full_prefix: false, + group_id: 0, }) .expect("empty block_hashes are a valid zero-hit query"); } @@ -1030,6 +1115,7 @@ mod tests { pp_rank: 0, transfer_mode: ProtoTransferMode::Direct as i32, page_first: false, + layer_group_ids: Vec::new(), }) .expect_err("tp_rank outside tp_size must be rejected at RPC boundary"); @@ -1056,6 +1142,7 @@ mod tests { pp_rank: 0, transfer_mode: ProtoTransferMode::Direct as i32, page_first: false, + layer_group_ids: Vec::new(), }) .expect_err("client/server version mismatch must be rejected before registration"); diff --git a/pegaflow-server/tests/common/mod.rs b/pegaflow-server/tests/common/mod.rs index 78c752dc..f4c669db 100644 --- a/pegaflow-server/tests/common/mod.rs +++ b/pegaflow-server/tests/common/mod.rs @@ -262,6 +262,7 @@ impl MockVllmRpcHarness { block_hashes: hashes.to_vec(), req_id: req_id.to_string(), wait_for_full_prefix: false, + group_id: 0, }; match self.scheduler.query_prefetch(request.clone()).await { Ok(response) => Ok(RpcExchange { diff --git a/pegaflow-server/tests/http_cleanup_hang_repro.rs b/pegaflow-server/tests/http_cleanup_hang_repro.rs index cdf3ad23..f0b5a976 100644 --- a/pegaflow-server/tests/http_cleanup_hang_repro.rs +++ b/pegaflow-server/tests/http_cleanup_hang_repro.rs @@ -248,6 +248,7 @@ fn wedge_register_request(i: usize) -> RegisterContextRequest { segments: vec![1], pp_rank: 0, transfer_mode: TransferMode::Direct as i32, + layer_group_ids: Vec::new(), page_first: false, } } diff --git a/python/pegaflow/connector/common.py b/python/pegaflow/connector/common.py index d08e8ea9..11e8b6a7 100644 --- a/python/pegaflow/connector/common.py +++ b/python/pegaflow/connector/common.py @@ -224,6 +224,65 @@ class LoadIntent: block_ids_by_group: tuple[tuple[int | None, ...], ...] leases: tuple[bytes, ...] num_tokens: int + # Hybrid-cache loads carry one membership lease per recurrent storage + # group (pinned checkpoints in hit-positions order) on top of the + # attention prefix leases. See RecurrentLoadHold. + recurrent_hold: "RecurrentLoadHold | None" = None + + +@dataclass(frozen=True) +class RecurrentLoadHold: + """Pinned recurrent checkpoints for one hybrid external load. + + Indexed by ``sorted(recurrent_group_indices)`` on the outside and TP + shard on the inside: ``leases[g][shard]`` is the membership lease over + group ``g``'s hit blocks; ``hit_positions[g][shard]`` lists each leased + block's position in the scheduler's query hash list (lease order). + ``checkpoint`` is the chosen query position — the mamba state stored + there covers all tokens through the end of that block (vLLM convention: + state block ``i`` ends at token ``(i + 1) * block_size``), so the + resumable prefix is ``checkpoint + 1`` blocks. + """ + + leases: tuple[tuple[bytes, ...], ...] + hit_positions: tuple[tuple[tuple[int, ...], ...], ...] + checkpoint: int + + +def reconcile_hybrid_hit( + attention_hit_blocks: int, + recurrent_hits: tuple[tuple[tuple[int, ...], ...], ...], +) -> tuple[int, int | None, frozenset[int]]: + """Combine per-group query results into one hybrid hit. + + ``attention_hit_blocks`` is the (already shard-minimized) attention prefix + length in blocks. ``recurrent_hits[g][s]`` lists the query positions whose + checkpoint block is cached in recurrent group ``g`` on shard ``s``. + + HMA needs the whole prefix resumable: every recurrent group must hold a + checkpoint state inside the attention prefix (attention KV alone cannot + skip mamba's sequential prefill), and that state must exist on every TP + shard. Returns ``(hit_blocks, checkpoint, usable)`` where ``hit_blocks`` + is ``checkpoint + 1`` — the checkpoint covers tokens through the end of + its own block — and ``usable`` is every legal boundary position (for + re-derivation when the token budget later shrinks the hit). ``(0, None, + frozenset())`` means no usable boundary: recompute from scratch. + """ + if attention_hit_blocks <= 0 or not recurrent_hits: + return 0, None, frozenset() + # A checkpoint position is usable only inside the attention prefix AND + # present in every recurrent group on every shard. + usable: set[int] | None = None + for group_hits in recurrent_hits: + for shard_hits in group_hits: + in_prefix = {p for p in shard_hits if p < attention_hit_blocks} + usable = in_prefix if usable is None else usable & in_prefix + if not usable: + return 0, None, frozenset() + if not usable: + return 0, None, frozenset() + checkpoint = max(usable) + return checkpoint + 1, checkpoint, frozenset(usable) @dataclass(frozen=True) @@ -236,13 +295,20 @@ class SaveIntent: @dataclass(frozen=True) class CacheGroupLayout: - """Stable vLLM cache-group order shared by scheduler and worker.""" + """Stable vLLM cache-group order shared by scheduler and worker. + + `storage_group_ids` maps each connector cache group onto the engine's + hybrid storage groups: every attention-like group shares storage group 0 + (prefix cadence, raw hash keys), while each recurrent group gets its own + id starting at 1 (membership semantics, group-encoded keys). + """ layer_names: tuple[tuple[str, ...], ...] hash_group_index: int has_recurrent_state: bool recurrent_group_indices: frozenset[int] recurrent_layer_names: frozenset[str] + storage_group_ids: tuple[int, ...] = (0,) @classmethod def from_config(cls, kv_cache_config) -> "CacheGroupLayout": @@ -323,21 +389,34 @@ def from_config(cls, kv_cache_config) -> "CacheGroupLayout": "PegaFlow requires a dense FullAttention cache group for block hashes" ) + recurrent_group_indices = frozenset( + index + for index, group in enumerate(groups) + if isinstance(group.kv_cache_spec, MambaSpec) + ) + # Attention-like groups all share storage group 0 (they advance in + # per-block prefix cadence); each recurrent group gets a dense id + # from 1. Engine keys are raw only for group 0, so this keeps every + # existing single-group cache layout bit-identical. + storage_group_ids = tuple( + 0 + if index not in recurrent_group_indices + else 1 + sum(1 for other in recurrent_group_indices if other < index) + for index in range(len(groups)) + ) + return cls( layer_names=tuple(tuple(group.layer_names) for group in groups), hash_group_index=hash_group_index, has_recurrent_state=any(isinstance(group.kv_cache_spec, MambaSpec) for group in groups), - recurrent_group_indices=frozenset( - index - for index, group in enumerate(groups) - if isinstance(group.kv_cache_spec, MambaSpec) - ), + recurrent_group_indices=recurrent_group_indices, recurrent_layer_names=frozenset( layer_name for group in groups if isinstance(group.kv_cache_spec, MambaSpec) for layer_name in group.layer_names ), + storage_group_ids=storage_group_ids, ) @property @@ -353,6 +432,10 @@ def layer_to_group(self) -> dict[str, int]: result[name] = group_index return result + def storage_group_of(self, group_index: int) -> int: + """Engine storage group id for a connector cache group index.""" + return self.storage_group_ids[group_index] + class PegaConnectorMetadata(KVConnectorMetadata): """Metadata passed from scheduler to worker for KV cache operations.""" @@ -361,17 +444,20 @@ def __init__( self, load_intents: dict[str, LoadIntent] | None = None, save_intents: dict[str, SaveIntent] | None = None, + ready_save_intents: dict[str, SaveIntent] | None = None, preempted_req_ids: set[str] | None = None, ): super().__init__() # Maps request_id -> intent self.load_intents: dict[str, LoadIntent] = load_intents or {} self.save_intents: dict[str, SaveIntent] = save_intents or {} + self.ready_save_intents: dict[str, SaveIntent] = ready_save_intents or {} self.preempted_req_ids: set[str] = preempted_req_ids or set() def __repr__(self) -> str: return ( - f"PegaConnectorMetadata(loads={len(self.load_intents)}, saves={len(self.save_intents)})" + f"PegaConnectorMetadata(loads={len(self.load_intents)}, " + f"saves={len(self.save_intents)}, ready_saves={len(self.ready_save_intents)})" ) @@ -515,12 +601,14 @@ def resolve_transfer_backend(is_mla: bool, override: str | None) -> str: "PegaConnectorMetadata", "PegaKVConnectorStats", "PegaPromMetrics", + "RecurrentLoadHold", "SaveIntent", "TpShardTopology", "derive_namespace", "detect_mla", "logger", "parse_env_int", + "reconcile_hybrid_hit", "resolve_instance_id", "resolve_transfer_backend", ] diff --git a/python/pegaflow/connector/scheduler.py b/python/pegaflow/connector/scheduler.py index 2e162869..3db3cda6 100644 --- a/python/pegaflow/connector/scheduler.py +++ b/python/pegaflow/connector/scheduler.py @@ -5,7 +5,7 @@ import os import time from collections.abc import Iterable -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import TYPE_CHECKING from pegaflow.connector.common import ( @@ -14,8 +14,10 @@ LoadIntent, PegaConnectorMetadata, PegaKVConnectorStats, + RecurrentLoadHold, SaveIntent, logger, + reconcile_hybrid_hit, ) from pegaflow.connector.connector_metrics import PrefetchTracker from pegaflow.connector.tp_shards import ShardedQueryReady, TpShardQueryClient @@ -53,6 +55,13 @@ class _QueryProbe: # ``None`` means the backend is still loading. hit_blocks: int | None = None leases: tuple[bytes, ...] = () + # Hybrid (HMA): pinned recurrent checkpoints from the membership queries, + # set together with `leases` when the hybrid reconcile found a boundary. + recurrent_hold: RecurrentLoadHold | None = None + # Sorted query positions a hybrid hit may end at (intersection over all + # recurrent groups and shards, below the attention prefix). Used to + # re-derive a legal boundary when the token budget shrinks the hit. + usable_positions: frozenset[int] = frozenset() @property def is_ready(self) -> bool: @@ -79,6 +88,8 @@ def mark_ready(self, ready: ShardedQueryReady) -> None: ) self.hit_blocks = hit_blocks self.leases = ready.leases + self.recurrent_hold = ready.recurrent_hold + self.usable_positions = frozenset(ready.usable_positions) def require_hit_blocks(self) -> int: if self.hit_blocks is None: @@ -179,6 +190,7 @@ def __init__( self._requests: dict[str, Request] = {} # Completion tracking + self._deferred_save_intents: dict[str, SaveIntent] = {} self._pending_saves: set[str] = set() self._held_requests: set[str] = set() @@ -273,6 +285,9 @@ def get_num_new_matched_tokens( len(query_hashes), ) self._release_leases(ready.leases, req_id) + if ready.recurrent_hold is not None: + for group_index, group_leases in enumerate(ready.recurrent_hold.leases): + self._tp_shard_client.release(group_leases, f"{req_id}:g{group_index}") self._pending_query_probes.pop(req_id, None) return (None, False) @@ -320,6 +335,23 @@ def _finish_cache_lookup( locally_computed_tokens = computed_blocks * vbs hit_tokens = min(hit_tokens, max(0, num_tokens - locally_computed_tokens - 1)) + if probe.recurrent_hold is not None: + # A mamba checkpoint is valid only at its own block boundary. If + # the token budget cut inside the reconciled span, fall back to + # the best earlier boundary; if none survives, drop the hit (a + # partial mamba resume cannot exist). + boundary_blocks = hit_tokens // vbs + usable = [p for p in probe.usable_positions if p < boundary_blocks] + if not usable: + if self._pending_query_probes.get(req_id) is probe: + self._release_pending_query_probe(req_id) + return (0, False) + checkpoint = max(usable) + hit_blocks = checkpoint + 1 + hit_tokens = hit_blocks * vbs + probe.hit_blocks = hit_blocks + probe.recurrent_hold = replace(probe.recurrent_hold, checkpoint=checkpoint) + # Cacheable tails contain at least two tokens, so recomputing the final # prompt token cannot remove the last leased block from the load. loaded_blocks = (hit_tokens + vbs - 1) // vbs @@ -418,6 +450,9 @@ def update_state_after_alloc( block_ids_by_group=load_block_ids_by_group, leases=pending_probe.leases if pending_probe is not None else (), num_tokens=num_external_tokens, + recurrent_hold=( + pending_probe.recurrent_hold if pending_probe is not None else None + ), ) if pending_probe is not None: query_hashes, tail_tokens = self._build_query(request, num_computed_blocks) @@ -449,6 +484,8 @@ def update_state_after_alloc( def build_connector_meta(self, scheduler_output: "SchedulerOutput") -> PegaConnectorMetadata: # Collect all save intents that became available this scheduler step. + ready_save_intents = self._deferred_save_intents + self._deferred_save_intents = {} potential_saves: dict[str, SaveIntent] = {} load_intents = self._pending_load_intents @@ -552,6 +589,7 @@ def build_connector_meta(self, scheduler_output: "SchedulerOutput") -> PegaConne return PegaConnectorMetadata( load_intents=load_intents, save_intents=save_intents, + ready_save_intents=ready_save_intents, preempted_req_ids=scheduler_output.preempted_req_ids or None, ) @@ -750,6 +788,54 @@ def _local_cached_block_ids( block_ids.append(block.block_id) return tuple(tuple(block_ids) for block_ids in block_ids_by_group) + def _consume_finished_hma_save( + self, + req_id: str, + block_ids: tuple[list[int], ...], + written: int, + ) -> SaveIntent | None: + block_hashes = self._block_hashes.get(req_id) + if block_hashes is None: + return None + + start_block_idx = self._next_stored_block_idx.get( + req_id, self._block_index_offsets.get(req_id, 0) + ) + saveable_block_idx = min(len(block_hashes), written // self._ctx.virtual_block_size) + if saveable_block_idx <= start_block_idx: + return None + + groups = self._copy_block_ids_by_group(block_ids) + available = tuple(len(group) for group in groups) + required = tuple( + saveable_block_idx + 1 + if group_index in self._cache_groups.recurrent_group_indices + else saveable_block_idx + for group_index in range(self._cache_groups.group_count) + ) + if any(length < minimum for length, minimum in zip(available, required, strict=True)): + raise RuntimeError( + f"req {req_id} final HMA block table is shorter than its hash prefix: " + f"saveable={saveable_block_idx} available_by_group={available} " + f"required_by_group={required}" + ) + + num_new_blocks = saveable_block_idx - start_block_idx + save_block_ids_by_group = [] + for group_index, group in enumerate(groups): + if group_index in self._cache_groups.recurrent_group_indices: + save_block_ids_by_group.append( + (0,) * (num_new_blocks - 1) + (group[saveable_block_idx],) + ) + else: + save_block_ids_by_group.append(group[start_block_idx:saveable_block_idx]) + + self._next_stored_block_idx[req_id] = saveable_block_idx + return SaveIntent( + block_ids_by_group=tuple(save_block_ids_by_group), + block_hashes=block_hashes[start_block_idx:saveable_block_idx], + ) + def _copy_block_ids_by_group(self, block_ids) -> tuple[tuple[int, ...], ...]: groups = tuple(tuple(group) for group in block_ids) if len(groups) != self._cache_groups.group_count: @@ -789,7 +875,7 @@ def update_connector_output(self, connector_output: "KVConnectorOutput") -> None logger.debug("[PegaKVConnector] Request %s save completed", req_id) # Clean up if request already finished - if req_id in self._held_requests: + if req_id in self._held_requests and req_id not in self._deferred_save_intents: self._cleanup_request(req_id) self._held_requests.discard(req_id) @@ -800,6 +886,17 @@ def request_finished( ) -> tuple[bool, dict | None]: req_id = request.request_id + if self._cache_groups.has_recurrent_state and req_id in self._block_hashes: + self._block_hashes[req_id] = tuple(request.block_hashes) + final_save = self._consume_finished_hma_save( + req_id, + block_ids, + request.num_computed_tokens, + ) + if final_save is not None: + self._deferred_save_intents[req_id] = final_save + self._pending_saves.add(req_id) + # Check if there are pending saves for this request if req_id in self._pending_saves: self._held_requests.add(req_id) @@ -823,6 +920,7 @@ def _cleanup_request(self, req_id: str) -> None: self._allocated_blocks.pop(req_id, None) self._scheduled_tokens.pop(req_id, None) self._next_stored_block_idx.pop(req_id, None) + self._deferred_save_intents.pop(req_id, None) self._pending_saves.discard(req_id) self._tail_saved.discard(req_id) @@ -868,8 +966,113 @@ def _count_available_block_prefix( self._prefetch_tracker.pending_prefetches, ) + if self._cache_groups.has_recurrent_state: + return self._reconcile_hybrid(block_hash_list, ready, req_id) return ready + def _reconcile_hybrid( + self, + block_hashes: list[bytes], + ready: ShardedQueryReady, + req_id: str, + ) -> ShardedQueryReady: + """Gate an attention-prefix hit on a usable recurrent boundary. + + HMA can resume only where every recurrent group cached its state on + every shard: attention KV alone cannot skip mamba's sequential + prefill. On success returns the reduced hit with the membership + leases attached; on failure every lease acquired is released and the + result degrades to a plain miss. + """ + if ready.num_hit_blocks == 0: + return ready + + group_ids = tuple( + self._cache_groups.storage_group_ids[index] + for index in sorted(self._cache_groups.recurrent_group_indices) + ) + per_group: list[list[tuple[tuple[int, ...], bytes]]] = [] + try: + for group_id in group_ids: + per_group.append( + self._tp_shard_client.query_group_membership( + self._ctx.instance_id, + block_hashes, + f"{req_id}:g{group_id}", + group_id, + ) + ) + hybrid_hit, checkpoint, usable = reconcile_hybrid_hit( + ready.num_hit_blocks, + tuple( + tuple(positions for positions, _ in group_results) + for group_results in per_group + ), + ) + except Exception: + self._release_leases(ready.leases, req_id) + for group_id, group_results in zip(group_ids, per_group, strict=False): + self._tp_shard_client.release( + tuple(lease for _, lease in group_results), f"{req_id}:g{group_id}" + ) + raise + + if hybrid_hit == 0 or checkpoint is None: + self._release_leases(ready.leases, req_id) + for group_id, group_results in zip(group_ids, per_group, strict=True): + self._tp_shard_client.release( + tuple(lease for _, lease in group_results), f"{req_id}:g{group_id}" + ) + logger.info( + "[PegaKVConnector] req=%s HMA attention prefix of %d blocks has no " + "common recurrent checkpoint; recomputing instead", + req_id, + ready.num_hit_blocks, + ) + return ShardedQueryReady(0, tuple(b"" for _ in ready.leases)) + + if hybrid_hit < ready.num_hit_blocks: + # The hit shrank behind the prefix lease: re-lease the exact + # shortened attention prefix so lease count and load agree. + exact = self._tp_shard_client.query( + self._ctx.instance_id, + block_hashes[:hybrid_hit], + f"{req_id}:hma-exact-{hybrid_hit}", + False, + ) + self._release_leases(ready.leases, req_id) + if exact is None or exact.num_hit_blocks != hybrid_hit: + if exact is not None: + self._release_leases(exact.leases, req_id) + for group_id, group_results in zip(group_ids, per_group, strict=True): + self._tp_shard_client.release( + tuple(lease for _, lease in group_results), f"{req_id}:g{group_id}" + ) + logger.warning( + "[PegaKVConnector] req=%s could not re-lease the reconciled " + "%d-block HMA prefix; recomputing instead", + req_id, + hybrid_hit, + ) + return ShardedQueryReady(0, tuple(b"" for _ in ready.leases)) + ready = exact + + return ShardedQueryReady( + hybrid_hit, + ready.leases, + RecurrentLoadHold( + leases=tuple( + tuple(lease for _, lease in group_results) for group_results in per_group + ), + hit_positions=tuple( + tuple(positions for positions, _ in group_results) + for group_results in per_group + ), + checkpoint=checkpoint, + ), + usable_positions=tuple(sorted(usable)), + ) + def _cancel_prefetch_tracking(self, req_id: str) -> None: """Drop in-flight prefetch metrics when polling stops before QueryReady.""" if req_id not in self._prefetch_start_times: @@ -914,10 +1117,17 @@ def _release_pending_query_probe(self, req_id: str) -> bool: return self._release_query_probe(req_id, probe) def _release_query_probe(self, req_id: str, probe: _QueryProbe) -> bool: - if not probe.leases or not any(probe.leases): + released = True + if probe.leases and any(probe.leases): + released = self._release_leases(probe.leases, req_id) + else: self._cancel_prefetch_tracking(req_id) - return True # nothing leased server-side (still loading, or zero-hit Ready) - return self._release_leases(probe.leases, req_id) + hold = probe.recurrent_hold + if hold is not None: + for group_index, group_leases in enumerate(hold.leases): + if not self._tp_shard_client.release(group_leases, f"{req_id}:g{group_index}"): + released = False + return released def _release_leases(self, leases: tuple[bytes, ...], req_id: str) -> bool: return self._tp_shard_client.release(leases, req_id) diff --git a/python/pegaflow/connector/tp_shards.py b/python/pegaflow/connector/tp_shards.py index a8871819..d29649c0 100644 --- a/python/pegaflow/connector/tp_shards.py +++ b/python/pegaflow/connector/tp_shards.py @@ -2,7 +2,7 @@ from dataclasses import dataclass -from pegaflow.connector.common import logger +from pegaflow.connector.common import RecurrentLoadHold, logger from pegaflow.pegaflow import EngineRpcClient, QueryLoading, QueryReady @@ -10,6 +10,13 @@ class ShardedQueryReady: num_hit_blocks: int leases: tuple[bytes, ...] + # HMA only: per recurrent group, per shard membership leases and their + # hit positions (see RecurrentLoadHold for the wire/load contract). + recurrent_hold: RecurrentLoadHold | None = None + # HMA only: sorted query positions the hit may legally end at (every + # recurrent group holds a checkpoint there on every shard, below the + # attention prefix). Drives boundary re-derivation under token clamps. + usable_positions: tuple[int, ...] = () class TpShardQueryClient: @@ -83,6 +90,56 @@ def query( return ShardedQueryReady(common_blocks, tuple(leases)) + def query_group_membership( + self, + instance_id: str, + block_hashes: list[bytes], + req_id: str, + group_id: int, + ) -> list[tuple[tuple[int, ...], bytes]]: + """Per-shard membership query over one hybrid storage group. + + Returns ``(hit_positions, lease)`` per shard; the lease pins exactly + the hit blocks in positions order. Membership queries are local-only, + so every shard answers Ready immediately (never Loading). + """ + results: list[tuple[tuple[int, ...], bytes]] = [] + try: + for shard_index, client in enumerate(self._clients): + result = client.query_prefetch( + instance_id, + block_hashes, + req_id=req_id, + group_id=group_id, + ) + if isinstance(result, QueryLoading): + raise RuntimeError( + f"TP shard {shard_index} membership query for group {group_id} " + "returned Loading; membership queries are local-only" + ) + if not isinstance(result, QueryReady): + raise TypeError(f"query_prefetch returned unexpected outcome {type(result)!r}") + positions = tuple(result.hit_positions) + if len(positions) != result.num_hit_blocks: + raise RuntimeError( + f"TP shard {shard_index} reported {result.num_hit_blocks} hits " + f"but returned {len(positions)} positions" + ) + if any(position >= len(block_hashes) for position in positions): + raise RuntimeError( + f"TP shard {shard_index} returned hit positions outside " + f"a {len(block_hashes)}-hash query" + ) + if positions and not result.lease: + raise RuntimeError( + f"TP shard {shard_index} returned {len(positions)} hits without a lease" + ) + results.append((positions, result.lease)) + except Exception: + self.release(tuple(lease for _, lease in results), req_id) + raise + return results + def release(self, leases: tuple[bytes, ...], req_id: str) -> bool: released = True for client, lease in zip(self._clients, leases, strict=False): diff --git a/python/pegaflow/connector/worker.py b/python/pegaflow/connector/worker.py index 23d15a51..dd5d50d6 100644 --- a/python/pegaflow/connector/worker.py +++ b/python/pegaflow/connector/worker.py @@ -17,6 +17,7 @@ ConnectorContext, PegaConnectorMetadata, PegaKVConnectorStats, + SaveIntent, logger, parse_env_int, ) @@ -360,6 +361,11 @@ def register_kv_caches(self, kv_caches: dict[str, Any]): registration.num_blocks, ) + layer_group_ids = [ + self._cache_groups.storage_group_ids[self._layer_to_group.get(name, 0)] + for name in layer_names + ] + ok, message = self._ctx.engine_client.register_context_batch( self._ctx.instance_id, self._ctx.namespace, @@ -376,6 +382,7 @@ def register_kv_caches(self, kv_caches: dict[str, Any]): layer_segments, self._ctx.transfer_backend, self._page_first, + layer_group_ids=layer_group_ids, ) if not ok: @@ -547,6 +554,11 @@ def start_load_kv( ) -> None: self._current_metadata = metadata + if metadata.ready_save_intents: + if not self._cache_groups.has_recurrent_state: + raise RuntimeError("ready save intents are only valid for HMA") + self._process_save_batch([self._make_save_task(metadata.ready_save_intents)]) + if not metadata.load_intents: return @@ -564,9 +576,46 @@ def start_load_kv( f"expected {self._ctx.tp_shard_count}" ) block_ids_by_group = [list(group) for group in load_intent.block_ids_by_group] + hold = load_intent.recurrent_hold + recurrent_groups = sorted(self._cache_groups.recurrent_group_indices) + if hold is not None: + # The attention lease pins group-0 (dense prefix) blocks only; + # recurrent destinations travel with the membership leases. + for group_index in recurrent_groups: + block_ids_by_group[group_index] = [None] * len(block_ids_by_group[group_index]) for block_ids in block_ids_by_group: all_block_ids.extend(block_id for block_id in block_ids if block_id is not None) loads.append((load_intent.leases[self._ctx.tp_shard_index], block_ids_by_group)) + if hold is not None: + shard = self._ctx.tp_shard_index + for slot, group_index in enumerate(recurrent_groups): + positions = hold.hit_positions[slot][shard] + if hold.checkpoint not in positions: + raise RuntimeError( + f"req {req_id}: recurrent group {group_index} shard {shard} " + f"lease has no checkpoint at query position {hold.checkpoint}" + ) + destination = next( + ( + block_id + for block_id in reversed(load_intent.block_ids_by_group[group_index]) + if block_id is not None + ), + None, + ) + if destination is None: + raise RuntimeError( + f"req {req_id}: no recurrent destination block in group " + f"{group_index} for checkpoint {hold.checkpoint}" + ) + # The membership lease pins `[hit_positions]` blocks; only + # the chosen checkpoint has a physical destination. + vectors: list[list[int | None]] = [ + [None] * len(positions) for _ in range(self._cache_groups.group_count) + ] + vectors[group_index][positions.index(hold.checkpoint)] = destination + loads.append((hold.leases[slot][shard], vectors)) + all_block_ids.append(destination) request_ids.append(req_id) if not all_block_ids: @@ -721,7 +770,17 @@ def wait_for_save(self) -> None: if metadata is None or not metadata.save_intents: return - request_ids = list(metadata.save_intents.keys()) + task = self._make_save_task(metadata.save_intents) + if self._cache_groups.has_recurrent_state: + # Align-mode recurrent states can reuse their only live block on + # the next scheduler step. Finish D2H before returning so the + # saved boundary state cannot be overwritten underneath the copy. + self._process_save_batch([task]) + else: + self._save_queue.put(task) + + def _make_save_task(self, save_intents: dict[str, SaveIntent]) -> SaveTask: + request_ids = list(save_intents) with self._save_completion_lock: for req_id in request_ids: @@ -731,14 +790,10 @@ def wait_for_save(self) -> None: self._save_completion_events[req_id] = threading.Event() self._req_pending_save_tasks[req_id] = pending_tasks + 1 - task = SaveTask(metadata=metadata, request_ids=request_ids) - if self._cache_groups.has_recurrent_state: - # Align-mode recurrent states can reuse their only live block on - # the next scheduler step. Finish D2H before returning so the - # saved boundary state cannot be overwritten underneath the copy. - self._process_save_batch([task]) - else: - self._save_queue.put(task) + return SaveTask( + metadata=PegaConnectorMetadata(save_intents=save_intents), + request_ids=request_ids, + ) def _save_worker(self) -> None: logger.debug("[PegaKVConnector] Save worker thread started") diff --git a/python/pegaflow/pegaflow.pyi b/python/pegaflow/pegaflow.pyi index a81a200b..b188fd2e 100644 --- a/python/pegaflow/pegaflow.pyi +++ b/python/pegaflow/pegaflow.pyi @@ -58,7 +58,13 @@ class QueryLoading: class QueryReady: num_hit_blocks: int lease: bytes - def __init__(self, num_hit_blocks: int, lease: bytes) -> None: ... + hit_positions: list[int] + def __init__( + self, + num_hit_blocks: int, + lease: bytes, + hit_positions: list[int] = ..., + ) -> None: ... class EngineRpcClient: """gRPC client for remote PegaEngine server communication. @@ -111,6 +117,7 @@ class EngineRpcClient: segments_list: list[int], transfer_backend: str, page_first: bool, + layer_group_ids: list[int] | None = None, ) -> tuple[bool, str]: """Register all KV cache layers on a GPU with a single RPC call. @@ -216,6 +223,7 @@ class EngineRpcClient: block_hashes: list[bytes], req_id: str, wait_for_full_prefix: bool = False, + group_id: int = 0, ) -> QueryLoading | QueryReady: """Query prefix cache hits with SSD prefetch support. diff --git a/python/pyproject.toml b/python/pyproject.toml index a3cb17c6..68406f8e 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "pegaflow-llm" -version = "0.23.10" +version = "0.23.11" description = "High-performance key-value storage engine with Python bindings" readme = "README.md" requires-python = ">=3.10" @@ -114,6 +114,6 @@ profile = "black" [tool.commitizen] name = "cz_conventional_commits" -version = "0.23.10" +version = "0.23.11" version_files = ["pyproject.toml:^version", "../Cargo.toml:^version"] tag_format = "v$version" diff --git a/python/src/lib.rs b/python/src/lib.rs index 2bb71ddd..c97d54b3 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -133,15 +133,22 @@ struct QueryReady { #[pyo3(get)] num_hit_blocks: usize, lease: PyQueryLease, + /// Membership queries (group_id > 0) only: indices into the queried + /// block_hashes whose block is cached; lease block i corresponds to + /// query position hit_positions[i]. Empty for prefix queries. + #[pyo3(get)] + hit_positions: Vec, } #[pymethods] impl QueryReady { #[new] - fn new(num_hit_blocks: usize, lease: PyQueryLease) -> Self { + #[pyo3(signature = (num_hit_blocks, lease, hit_positions=Vec::new()))] + fn new(num_hit_blocks: usize, lease: PyQueryLease, hit_positions: Vec) -> Self { Self { num_hit_blocks, lease, + hit_positions, } } @@ -152,8 +159,9 @@ impl QueryReady { fn __repr__(&self) -> String { format!( - "QueryReady(num_hit_blocks={}, has_lease={})", + "QueryReady(num_hit_blocks={}, hits={:?}, has_lease={})", self.num_hit_blocks, + self.hit_positions, !self.lease.0.is_empty() ) } @@ -276,7 +284,11 @@ impl EngineRpcClient { clippy::too_many_arguments, reason = "PyO3 binding mirrors the public batch registration call shape" )] - #[pyo3(signature = (instance_id, namespace, tp_rank, pp_rank, tp_size, world_size, device_id, layer_names, wrapper_bytes_list, num_blocks_list, bytes_per_block_list, kv_stride_bytes_list, segments_list, transfer_backend, page_first))] + #[pyo3(signature = (instance_id, namespace, tp_rank, pp_rank, tp_size, world_size, device_id, layer_names, wrapper_bytes_list, num_blocks_list, bytes_per_block_list, kv_stride_bytes_list, segments_list, transfer_backend, page_first, layer_group_ids=None))] + #[allow( + clippy::too_many_arguments, + reason = "PyO3 binding mirrors the public batch registration call shape" + )] fn register_context_batch( &self, py: Python<'_>, @@ -295,6 +307,7 @@ impl EngineRpcClient { segments_list: Vec, transfer_backend: &str, page_first: bool, + layer_group_ids: Option>, ) -> PyResult<(bool, String)> { let transfer_mode = match transfer_backend { "direct" => TransferMode::Direct, @@ -324,6 +337,7 @@ impl EngineRpcClient { pp_rank, transfer_mode: transfer_mode as i32, page_first, + layer_group_ids: layer_group_ids.unwrap_or_default(), }) .await?; Ok(resp.into_inner()) @@ -462,7 +476,7 @@ impl EngineRpcClient { /// /// Returns: /// QueryLoading while backing fetch is in progress, otherwise QueryReady. - #[pyo3(signature = (instance_id, block_hashes, req_id, wait_for_full_prefix=false))] + #[pyo3(signature = (instance_id, block_hashes, req_id, wait_for_full_prefix=false, group_id=0))] fn query_prefetch( &self, py: Python<'_>, @@ -470,6 +484,7 @@ impl EngineRpcClient { block_hashes: Vec>, req_id: String, wait_for_full_prefix: bool, + group_id: u32, ) -> PyResult> { let result = py.detach(|| { self.rt_handle.block_on(async { @@ -480,6 +495,7 @@ impl EngineRpcClient { block_hashes, req_id, wait_for_full_prefix, + group_id, }) .await .map(|resp| resp.into_inner()) @@ -495,6 +511,7 @@ impl EngineRpcClient { QueryReady { num_hit_blocks: 0, lease: PyQueryLease(Vec::new()), + hit_positions: Vec::new(), }, ) .map(|obj| obj.into_any()) @@ -512,6 +529,7 @@ impl EngineRpcClient { QueryReady { num_hit_blocks: u64_to_usize(ready.num_hit_blocks, "num_hit_blocks")?, lease: PyQueryLease(ready.lease), + hit_positions: ready.hit_positions, }, ) .map(|obj| obj.into_any()), diff --git a/python/tests/test_combine_hashes.py b/python/tests/test_combine_hashes.py index f9003ed1..d538d1ad 100644 --- a/python/tests/test_combine_hashes.py +++ b/python/tests/test_combine_hashes.py @@ -104,6 +104,52 @@ def test_recurrent_save_waits_until_vllm_commits_all_groups(): assert scheduler._next_stored_block_idx["r1"] == 0 +def test_recurrent_final_step_uses_finished_block_table(): + scheduler = _make_recurrent_scheduler() + scheduler._get_local_cached_blocks = lambda _block_hash, _group_ids: None + + assert scheduler._consume_full_block_saves("r1", written=32, computed_before_step=0) is None + request = SimpleNamespace( + request_id="r1", + num_computed_tokens=32, + block_hashes=[_hash(0), _hash(1)], + ) + + delay_free, params = scheduler.request_finished( + request, + ([11, 12, 13], [0, 0, 22, 30, 31, 32]), + ) + + assert delay_free is True + assert params is None + + scheduler.update_connector_output(SimpleNamespace(finished_sending={"r1"})) + assert "r1" in scheduler._block_hashes + + metadata = scheduler.build_connector_meta( + SimpleNamespace( + scheduled_new_reqs=[], + scheduled_cached_reqs=SimpleNamespace( + req_ids=[], + resumed_req_ids=set(), + new_block_ids=[], + num_computed_tokens=[], + ), + num_scheduled_tokens={}, + preempted_req_ids=set(), + ) + ) + assert metadata.ready_save_intents["r1"] == SaveIntent( + block_ids_by_group=((11, 12), (0, 22)), + block_hashes=(_hash(0), _hash(1)), + ) + assert metadata.save_intents == {} + + scheduler.update_connector_output(SimpleNamespace(finished_sending={"r1"})) + assert "r1" not in scheduler._block_hashes + assert "r1" not in scheduler._held_requests + + def test_recurrent_save_does_not_predict_current_step_state(): scheduler = _make_recurrent_scheduler() diff --git a/python/tests/test_connector_fault_tolerance.py b/python/tests/test_connector_fault_tolerance.py index 7d3e105f..dbf5df12 100644 --- a/python/tests/test_connector_fault_tolerance.py +++ b/python/tests/test_connector_fault_tolerance.py @@ -47,6 +47,7 @@ def __init__(self) -> None: self.register_response: tuple[bool, str] = (True, "ok") self.register_exception: Exception | None = None self.register_calls: list[tuple] = [] + self.register_kwargs: list[dict] = [] self.unregister_calls: list[str] = [] self.release_calls: list[bytes] = [] @@ -76,8 +77,9 @@ def load( return (False, "simulated load failure") return (True, "ok") - def register_context_batch(self, *args) -> tuple[bool, str]: + def register_context_batch(self, *args, **kwargs) -> tuple[bool, str]: self.register_calls.append(args) + self.register_kwargs.append(kwargs) if self.register_exception is not None: raise self.register_exception return self.register_response diff --git a/python/tests/test_connector_save_lifecycle.py b/python/tests/test_connector_save_lifecycle.py index eb3503f6..fdd33e37 100644 --- a/python/tests/test_connector_save_lifecycle.py +++ b/python/tests/test_connector_save_lifecycle.py @@ -109,6 +109,50 @@ def test_later_save_reopens_completed_request(): assert finished_sending == {"request"} +def test_committed_save_runs_on_a_no_forward_step(): + worker = make_worker() + worker._cache_groups = SimpleNamespace(has_recurrent_state=True) + worker._registered_layers = ["layer"] + worker._ctx.engine_client.save.return_value = (True, "") + metadata = PegaConnectorMetadata( + ready_save_intents={ + "request": SaveIntent( + block_ids_by_group=((1,),), + block_hashes=(b"hash",), + ) + } + ) + + with patch("torch.cuda.synchronize"): + worker.start_load_kv(metadata, None) + + worker._ctx.engine_client.save.assert_called_once() + finished_sending, _ = worker.get_finished({"request"}) + assert finished_sending == {"request"} + + +def test_current_step_save_waits_for_forward_completion(): + worker = make_worker() + worker._cache_groups = SimpleNamespace(has_recurrent_state=True) + worker._registered_layers = ["layer"] + worker._ctx.engine_client.save.return_value = (True, "") + metadata = PegaConnectorMetadata( + save_intents={ + "request": SaveIntent( + block_ids_by_group=((1,),), + block_hashes=(b"hash",), + ) + } + ) + + worker.start_load_kv(metadata, None) + worker._ctx.engine_client.save.assert_not_called() + + with patch("torch.cuda.synchronize"): + worker.wait_for_save() + worker._ctx.engine_client.save.assert_called_once() + + def test_preemption_waits_for_every_save_task(): worker = make_worker() completion = enqueue_save(worker) diff --git a/python/tests/test_hybrid_reconcile.py b/python/tests/test_hybrid_reconcile.py new file mode 100644 index 00000000..f776aba4 --- /dev/null +++ b/python/tests/test_hybrid_reconcile.py @@ -0,0 +1,128 @@ +"""Hybrid (HMA) prefix reconcile: attention prefix + recurrent membership. + +Pure-function contract covering the scheduler-side hit derivation: + +- a recurrent checkpoint at query position ``k`` resumes ``k + 1`` blocks + (the state block ends with block ``k``'s tokens, vLLM convention); +- checkpoints outside the attention prefix are unusable; +- every recurrent group AND every TP shard must hold the boundary; +- ``usable`` positions drive re-derivation when the token budget shrinks + the reconciled hit. +""" + +from __future__ import annotations + +from .unit_stubs import install_connector_unit_stubs + +install_connector_unit_stubs() + +from pegaflow.connector.common import ( # noqa: E402 + CacheGroupLayout, + reconcile_hybrid_hit, +) + +from .test_cache_group_layout import ( # noqa: E402 + _config, + _full_attention, + _group, + _mamba, +) + + +def _hits(*per_shard_positions: tuple[int, ...]) -> tuple[tuple[int, ...], ...]: + """One recurrent group's per-shard hit positions.""" + return tuple(tuple(p) for p in per_shard_positions) + + +class TestReconcileHybridHit: + def test_rightmost_checkpoint_within_prefix(self): + # attn=111, recur=001 -> hit = 2 + 1 = 3 + hit, checkpoint, usable = reconcile_hybrid_hit(3, (_hits((2,)),)) + assert hit == 3 + assert checkpoint == 2 + assert usable == frozenset({2}) + + def test_checkpoint_beyond_attention_prefix_is_unusable(self): + # The checkpoint at position 2 outruns the 2-block attention prefix. + hit, checkpoint, usable = reconcile_hybrid_hit(2, (_hits((2,)),)) + assert hit == 0 + assert checkpoint is None + assert usable == frozenset() + + def test_no_checkpoint_means_recompute(self): + # Attention alone cannot resume a recurrent model. + assert reconcile_hybrid_hit(3, (_hits(()),)) == (0, None, frozenset()) + + def test_shards_must_agree_on_the_boundary(self): + # shard0 has {1, 2}, shard1 has {2, 3}; prefix is 4 blocks. + hit, checkpoint, usable = reconcile_hybrid_hit(4, (_hits((1, 2), (2, 3)),)) + assert hit == 3 + assert checkpoint == 2 + assert usable == frozenset({2}) + + def test_all_recurrent_groups_must_hold_the_boundary(self): + # Two recurrent groups (e.g. mamba + gated-deltanet stacks). + hit, checkpoint, usable = reconcile_hybrid_hit(3, (_hits((0, 2)), _hits((1, 2)))) + assert hit == 3 + assert checkpoint == 2 + assert usable == frozenset({2}) + + # Group 2 lost its position-2 checkpoint: fall back to nothing. + hit, checkpoint, usable = reconcile_hybrid_hit(3, (_hits((0, 2)), _hits((1,)))) + assert hit == 0 + assert checkpoint is None + + def test_mid_prefix_gap_is_fine_below_rightmost(self): + # recur=101 with attn prefix 3: the missing middle checkpoint does + # not matter; rightmost hit at 2 still resumes 3 blocks. + hit, checkpoint, usable = reconcile_hybrid_hit(3, (_hits((0, 2)),)) + assert hit == 3 + assert checkpoint == 2 + assert usable == frozenset({0, 2}) + + def test_zero_attention_prefix_means_no_hit(self): + assert reconcile_hybrid_hit(0, (_hits((0,)),)) == (0, None, frozenset()) + + def test_earlier_boundary_survives_for_budget_rederivation(self): + # usable carries every legal boundary so a token-budget clamp can + # re-derive hit=2 from checkpoint 1 after 3 was trimmed away. + _, checkpoint, usable = reconcile_hybrid_hit(4, (_hits((1, 3)),)) + assert checkpoint == 3 + assert usable == frozenset({1, 3}) + + +class TestStorageGroupIds: + def test_attention_first_layout(self): + config = _config( + _group("attn", _full_attention()), + _group("mamba", _mamba()), + ) + layout = CacheGroupLayout.from_config(config) + assert layout.storage_group_ids == (0, 1) + assert layout.storage_group_of(1) == 1 + + def test_recurrent_group_can_come_first(self): + # vLLM group order is not guaranteed; attention must always map to + # storage group 0 regardless of connector position. + config = _config( + _group("mamba", _mamba()), + _group("attn", _full_attention()), + ) + layout = CacheGroupLayout.from_config(config) + assert layout.hash_group_index == 1 + assert layout.storage_group_ids == (1, 0) + + def test_single_group_defaults_to_zero(self): + config = _config(_group("attn", _full_attention())) + layout = CacheGroupLayout.from_config(config) + assert layout.storage_group_ids == (0,) + + def test_multiple_recurrent_groups_get_dense_ids(self): + config = _config( + _group("attn", _full_attention()), + _group("mamba_a", _mamba()), + _group("mamba_b", _mamba()), + ) + layout = CacheGroupLayout.from_config(config) + assert layout.storage_group_ids == (0, 1, 2) + assert layout.recurrent_group_indices == frozenset({1, 2}) From 157279a2fbe0ed679a0ab2f57c686c78dbafec85 Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Mon, 17 Aug 2026 12:07:22 +0800 Subject: [PATCH 04/16] fix(core): keep leased blocks indexed under pressure (#434) ## Summary - reclaim only cache entries exclusively owned by the resident cache - keep leased, queried, loading, and weakly referenced blocks indexed during memory pressure - preserve reclaimable-before-retained ordering with a bounded LRU scan - strengthen the GPU eviction integration test to prove an unleased batch is reclaimed while a leased batch remains queryable and loadable ## Testing - `cargo test --no-default-features --features cuda-13,rdma -p pegaflow-core storage::read_cache::tests -- --nocapture` - `cargo test --no-default-features --features cuda-13,rdma -p pegaflow-core --test eviction -- --nocapture` - `cargo clippy --workspace --all-targets --no-default-features --features cuda-13,rdma -- -D warnings` - `prek run` (includes release-mode Rust test suite) Signed-off-by: xiaguan <751080330@qq.com> --- pegaflow-core/src/cache.rs | 10 ++++ pegaflow-core/src/storage/mod.rs | 10 ---- pegaflow-core/src/storage/read_cache.rs | 64 ++++++++++++++++++++++--- pegaflow-core/tests/eviction.rs | 19 ++++++-- 4 files changed, 81 insertions(+), 22 deletions(-) diff --git a/pegaflow-core/src/cache.rs b/pegaflow-core/src/cache.rs index 90e209e7..ba8e623b 100644 --- a/pegaflow-core/src/cache.rs +++ b/pegaflow-core/src/cache.rs @@ -72,6 +72,16 @@ impl TinyLfuCache { self.lru.contains_key(key) } + /// Returns true only when the cache owns the block exclusively. + /// + /// `Arc::get_mut` rejects both other strong references and weak references, + /// so no holder outside the cache can keep or reacquire the allocation. + pub(crate) fn is_exclusively_owned(&mut self, key: &BlockKey) -> bool { + self.lru + .peek_mut(key) + .is_some_and(|block| Arc::get_mut(block).is_some()) + } + /// Insert with TinyLFU admission. If the candidate is colder than the /// current LRU victim it is dropped. pub(crate) fn insert(&mut self, key: BlockKey, value: ArcSealedBlock) -> CacheInsertOutcome { diff --git a/pegaflow-core/src/storage/mod.rs b/pegaflow-core/src/storage/mod.rs index 632a5f57..e17aee96 100644 --- a/pegaflow-core/src/storage/mod.rs +++ b/pegaflow-core/src/storage/mod.rs @@ -515,22 +515,12 @@ impl StorageEngine { } let mut batch_bytes = 0u64; - let mut still_referenced = 0u64; for (_key, block) in &evicted { let b = block.memory_footprint(); batch_bytes = batch_bytes.saturating_add(b); - if Arc::strong_count(block) > 1 { - still_referenced += 1; - } core_metrics().cache_resident_bytes.add(-(b as i64), &[]); } - if still_referenced > 0 { - core_metrics() - .cache_block_evictions_still_referenced - .add(still_referenced, &[]); - } - freed_bytes = freed_bytes.saturating_add(batch_bytes); freed_blocks += evicted.len(); diff --git a/pegaflow-core/src/storage/read_cache.rs b/pegaflow-core/src/storage/read_cache.rs index 325f5e05..e125cfd6 100644 --- a/pegaflow-core/src/storage/read_cache.rs +++ b/pegaflow-core/src/storage/read_cache.rs @@ -160,13 +160,19 @@ impl ReadCache { let removed = { let mut inner = self.inner.lock(); let mut removed = Vec::with_capacity(batch_size); - while removed.len() < batch_size { - let next = remove_lru(&mut inner, ResidentClass::Reclaimable) - .or_else(|| remove_lru(&mut inner, ResidentClass::Retained)); - let Some(block) = next else { - break; - }; - removed.push(block); + remove_lru_batch_from_class( + &mut inner, + ResidentClass::Reclaimable, + batch_size, + &mut removed, + ); + if removed.len() < batch_size { + remove_lru_batch_from_class( + &mut inner, + ResidentClass::Retained, + batch_size, + &mut removed, + ); } removed }; @@ -353,6 +359,35 @@ fn remove_lru(inner: &mut ReadCacheInner, class: ResidentClass) -> Option, +) { + let candidates = class_lru(inner, class).len(); + for _ in 0..candidates { + if removed.len() == batch_size { + break; + } + + let Some(key) = class_lru(inner, class) + .iter() + .next() + .map(|(key, _)| key.clone()) + else { + break; + }; + if inner.cache.is_exclusively_owned(&key) { + let block = remove_lru(inner, class) + .expect("exclusive LRU candidate must remain resident while locked"); + removed.push(block); + } else { + class_lru(inner, class).get(&key); + } + } +} + fn record_residence_durations( removed: Vec, attributes: &[opentelemetry::KeyValue], @@ -461,6 +496,21 @@ mod tests { ); } + #[test] + fn pressure_reclaim_waits_for_weak_references() { + let cache = make_cache(); + let key = BlockKey::new("ns".into(), vec![1]); + let block = make_block(); + let weak = Arc::downgrade(&block); + cache.batch_insert(vec![(key.clone(), block)]); + + assert!(cache.remove_lru_batch(1).is_empty()); + assert_class(&cache, &key, ResidentClass::Retained); + + drop(weak); + assert_eq!(cache.remove_lru_batch(1)[0].0, key); + } + #[test] fn local_hit_refreshes_recency_without_changing_class() { let cache = make_cache(); diff --git a/pegaflow-core/tests/eviction.rs b/pegaflow-core/tests/eviction.rs index bf5069e9..180b2a37 100644 --- a/pegaflow-core/tests/eviction.rs +++ b/pegaflow-core/tests/eviction.rs @@ -11,7 +11,7 @@ use pegaflow_core::StorageConfig; const BLOCK_SIZE: usize = 4096; const NUM_BLOCKS: usize = 4; -/// Pool fits one batch with headroom for metadata, but not two. +/// Pool fits two batches, so a third batch must reclaim one of them. const POOL_SIZE: usize = NUM_BLOCKS * BLOCK_SIZE * 2; fn eviction_storage_config() -> StorageConfig { @@ -61,11 +61,20 @@ async fn leased_blocks_survive_eviction_pressure() { env.save_and_wait(&hashes).await; let lease = env.assert_all_hit_lease(&hashes).await; - // Second batch creates eviction pressure while the first is leased. - let pressure = make_block_hashes(NUM_BLOCKS, 20); - env.save_layer(0, &pressure).await; + // Fill the remaining pool, then force a third batch to reclaim it while + // the first batch remains leased. + let reclaimed = make_block_hashes(NUM_BLOCKS, 20); + env.save_layer_and_flush(0, &reclaimed).await; + let pressure = make_block_hashes(NUM_BLOCKS, 21); + env.save_layer_and_flush(0, &pressure).await; + assert_eq!(env.count_hits_then_release(&reclaimed).await, 0); - // First batch is still leased — load must succeed. + // The lease pins both the allocation and its cache index, so an exact + // follow-up query must still see the same prefix after pressure. + let second_lease = env.assert_all_hit_lease(&hashes).await; + env.release(&second_lease); + + // The original lease must still load the pinned data. env.data().zero_gpu(); env.load_to_gpu(lease, hashes.len()).await; env.data().assert_gpu_matches_expected(); From 96e1ff6d75c9b63a85542b9471c8c7a64cdc3b48 Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Mon, 17 Aug 2026 14:02:54 +0800 Subject: [PATCH 05/16] chore(release): bump version to 0.23.12 (#437) ## Summary - bump the Rust workspace packages to `0.23.12` - bump the Python package and Commitizen version to `0.23.12` - refresh the workspace package versions in `Cargo.lock` ## Validation - `cargo metadata --locked --no-deps --format-version 1` - `prek run` (includes Clippy and release-mode Rust tests) Signed-off-by: xiaguan <751080330@qq.com> --- Cargo.lock | 16 ++++++++-------- Cargo.toml | 2 +- python/pyproject.toml | 4 ++-- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c36ecde6..63d68d69 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1896,7 +1896,7 @@ checksum = "19b17cddbe7ec3f8bc800887bab5e717348c95ea2ca0b1bf0837fb964dc67099" [[package]] name = "pegaflow-common" -version = "0.23.11" +version = "0.23.12" dependencies = [ "colored", "libc", @@ -1906,7 +1906,7 @@ dependencies = [ [[package]] name = "pegaflow-core" -version = "0.23.11" +version = "0.23.12" dependencies = [ "ahash", "bytesize", @@ -1939,7 +1939,7 @@ dependencies = [ [[package]] name = "pegaflow-metaserver" -version = "0.23.11" +version = "0.23.12" dependencies = [ "axum", "clap", @@ -1959,7 +1959,7 @@ dependencies = [ [[package]] name = "pegaflow-pd-wire" -version = "0.23.11" +version = "0.23.12" dependencies = [ "serde", "serde_json", @@ -1967,7 +1967,7 @@ dependencies = [ [[package]] name = "pegaflow-proto" -version = "0.23.11" +version = "0.23.12" dependencies = [ "prost", "tonic", @@ -1977,7 +1977,7 @@ dependencies = [ [[package]] name = "pegaflow-py" -version = "0.23.11" +version = "0.23.12" dependencies = [ "log", "mea", @@ -1995,7 +1995,7 @@ dependencies = [ [[package]] name = "pegaflow-server" -version = "0.23.11" +version = "0.23.12" dependencies = [ "axum", "clap", @@ -2027,7 +2027,7 @@ dependencies = [ [[package]] name = "pegaflow-transfer" -version = "0.23.11" +version = "0.23.12" dependencies = [ "anyhow", "bincode", diff --git a/Cargo.toml b/Cargo.toml index f3c502f1..9ec28224 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,7 @@ members = [ resolver = "2" [workspace.package] -version = "0.23.11" +version = "0.23.12" edition = "2024" license = "Apache-2.0" diff --git a/python/pyproject.toml b/python/pyproject.toml index 68406f8e..be561c3f 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "pegaflow-llm" -version = "0.23.11" +version = "0.23.12" description = "High-performance key-value storage engine with Python bindings" readme = "README.md" requires-python = ">=3.10" @@ -114,6 +114,6 @@ profile = "black" [tool.commitizen] name = "cz_conventional_commits" -version = "0.23.11" +version = "0.23.12" version_files = ["pyproject.toml:^version", "../Cargo.toml:^version"] tag_format = "v$version" From d53f886c971e27dd24d3c0ea2cbc5c1f48c24ad7 Mon Sep 17 00:00:00 2001 From: GentleCold Date: Mon, 17 Aug 2026 14:14:27 +0800 Subject: [PATCH 06/16] fix(core): allow SSD weak refs during pressure eviction (#436) ## Summary - Treat only additional strong `Arc` owners as pressure-eviction pins. - Allow SSD write-queue `Weak` references to yield to memory pressure, preserving the existing fire-and-forget SSD contract. - Add regression coverage for weak-only eviction and strong-reference retention. ## Validation Ran on RTX 4090 / Linux with CUDA 13 and RDMA enabled: - `cargo fmt --all -- --check` - `cargo test --no-default-features --features cuda-13,rdma -p pegaflow-core storage::read_cache::tests -- --nocapture` (22 passed) - `cargo clippy --workspace --all-targets --no-default-features --features cuda-13,rdma -- -D warnings` - `cargo test --no-default-features --features cuda-13,rdma -p pegaflow-core --test eviction -- --nocapture --test-threads=1` (3 passed) The existing `ssd_cache` integration suite was also attempted on the 4090 host, but its fixed 32 KiB test pool is split across the host's NUMA-local pools and exhausts before the first save; all 11 failures are `pinned pool exhausted`, before SSD assertions run. --- pegaflow-core/src/cache.rs | 12 ++++++------ pegaflow-core/src/storage/read_cache.rs | 22 ++++++++++++++++++---- 2 files changed, 24 insertions(+), 10 deletions(-) diff --git a/pegaflow-core/src/cache.rs b/pegaflow-core/src/cache.rs index ba8e623b..55620006 100644 --- a/pegaflow-core/src/cache.rs +++ b/pegaflow-core/src/cache.rs @@ -72,14 +72,14 @@ impl TinyLfuCache { self.lru.contains_key(key) } - /// Returns true only when the cache owns the block exclusively. + /// Returns true when the cache is the only strong owner of the block. /// - /// `Arc::get_mut` rejects both other strong references and weak references, - /// so no holder outside the cache can keep or reacquire the allocation. - pub(crate) fn is_exclusively_owned(&mut self, key: &BlockKey) -> bool { + /// Weak references from fire-and-forget backing-store work do not pin the + /// resident block: they may fail to upgrade after pressure eviction. + pub(crate) fn is_cache_owned_only(&self, key: &BlockKey) -> bool { self.lru - .peek_mut(key) - .is_some_and(|block| Arc::get_mut(block).is_some()) + .peek(key) + .is_some_and(|block| Arc::strong_count(block) == 1) } /// Insert with TinyLFU admission. If the candidate is colder than the diff --git a/pegaflow-core/src/storage/read_cache.rs b/pegaflow-core/src/storage/read_cache.rs index e125cfd6..5457bb26 100644 --- a/pegaflow-core/src/storage/read_cache.rs +++ b/pegaflow-core/src/storage/read_cache.rs @@ -378,9 +378,9 @@ fn remove_lru_batch_from_class( else { break; }; - if inner.cache.is_exclusively_owned(&key) { + if inner.cache.is_cache_owned_only(&key) { let block = remove_lru(inner, class) - .expect("exclusive LRU candidate must remain resident while locked"); + .expect("cache-owned LRU candidate must remain resident while locked"); removed.push(block); } else { class_lru(inner, class).get(&key); @@ -497,18 +497,32 @@ mod tests { } #[test] - fn pressure_reclaim_waits_for_weak_references() { + fn pressure_reclaim_ignores_weak_references() { let cache = make_cache(); let key = BlockKey::new("ns".into(), vec![1]); let block = make_block(); let weak = Arc::downgrade(&block); cache.batch_insert(vec![(key.clone(), block)]); + assert_eq!(cache.remove_lru_batch(1)[0].0, key); + assert!(weak.upgrade().is_none()); + } + + #[test] + fn pressure_reclaim_waits_for_external_strong_reference() { + let cache = make_cache(); + let key = BlockKey::new("ns".into(), vec![1]); + let block = make_block(); + let external = Arc::clone(&block); + let weak = Arc::downgrade(&block); + cache.batch_insert(vec![(key.clone(), block)]); + assert!(cache.remove_lru_batch(1).is_empty()); assert_class(&cache, &key, ResidentClass::Retained); - drop(weak); + drop(external); assert_eq!(cache.remove_lru_batch(1)[0].0, key); + assert!(weak.upgrade().is_none()); } #[test] From d3ed31a7a4ee9bc7cdc5f112f0a1eb26efe63675 Mon Sep 17 00:00:00 2001 From: GentleCold Date: Mon, 17 Aug 2026 14:37:41 +0800 Subject: [PATCH 07/16] feat(metrics): improve local HLL accuracy (#435) ## Summary - Increase the default local HLL size from `bucket_bits=14` (16,384 registers) to `16` (65,536 registers), reducing the theoretical standard error from about 0.8% to about 0.4%. - Record only the final miss suffix in HLL cardinality while keeping all queried blocks in `pegaflow_hll_total_requests`. - Add `pegaflow_hll_estimated_hit_rate`, calculated and clamped in the tracker, while retaining the existing cardinality and total metrics and labels. - Use stable `(namespace, block hash)` identities so local measurements do not collide across namespaces. - Update the metrics documentation and default windows to `15m,1h,1d`. ## Compatibility This is not an `/metrics` protocol breaking change: the HTTP endpoint, existing metric names, types, and labels remain available. The new gauge is additive. The existing PromQL remains valid, but the HLL cardinality semantics are now explicitly miss-only and the default precision is higher. MetaServer HLL aggregation, heartbeat, protobuf, and MetaServer metrics are intentionally excluded and will be handled by a follow-up PR. ## Validation - `cargo fmt --all -- --check` - `cargo test -p pegaflow-common` (34 passed) - `git diff --check` Server compilation could not complete on the macOS development host because the workspace's Linux-only `io-uring`/CUDA dependencies require Linux headers/toolchains; Linux CI remains the authoritative server build check. --- docs/metrics.md | 42 +++++++--- pegaflow-common/src/hll.rs | 140 ++++++++++++++++++++++++++++----- pegaflow-core/src/lib.rs | 5 ++ pegaflow-server/src/lib.rs | 9 ++- pegaflow-server/src/metric.rs | 24 +++++- pegaflow-server/src/service.rs | 13 ++- 6 files changed, 195 insertions(+), 38 deletions(-) diff --git a/docs/metrics.md b/docs/metrics.md index d2b1ec58..3b36e410 100644 --- a/docs/metrics.md +++ b/docs/metrics.md @@ -141,18 +141,25 @@ PegaFlow exposes the following metrics for monitoring KV cache operations: ### HLL Reuse Metrics - **pegaflow_hll_cardinality** (Gauge) - - Estimated distinct block hashes observed in a configured sliding window + - Estimated distinct `(namespace, block hash)` objects classified as misses in a configured sliding window - Labels: `window` (`15m`, `1h`, `1d` by default) - Use case: Derive approximate prefix reuse over longer windows without storing every block hash - **pegaflow_hll_total_requests** (Gauge) - - Total block hash observations in the same configured sliding window + - Total queried blocks in the same configured sliding window, including ready blocks and duplicates - Labels: `window` (`15m`, `1h`, `1d` by default) - - Use case: Denominator for HLL-based estimated hit-rate PromQL + - Use case: Denominator for HLL-based reference reuse rate -PegaFlow does not export a separate HLL hit-rate gauge. Use PromQL so the -ratio is computed from values in the same scrape: +- **pegaflow_hll_estimated_hit_rate** (Gauge) + - Server-computed miss-based infinite-cache reuse reference from the same + HLL snapshot as the two gauges above + - Labels: `window` + - Value is clamped to `[0, 1]` + +The existing cardinality and total metrics are retained. Existing PromQL +continues to work, but new dashboards should prefer the direct gauge because +it applies the same cardinality clamp as the tracker: ```promql 1 - ( @@ -162,6 +169,18 @@ ratio is computed from values in the same scrape: ) ``` +```promql +pegaflow_hll_estimated_hit_rate{window="1h"} +``` + +This is a metrics semantic update, not an `/metrics` protocol breaking change: +metric names, types, existing labels, and the HTTP endpoint are unchanged; +the new gauge is additive. The default HLL size changes from 16,384 registers +(`bucket_bits=14`, about 0.8% standard error) to 65,536 registers +(`bucket_bits=16`, about 0.4%). Three default windows use about 192 KiB of +register storage; sliding slots make the live tracker a few MiB per server. +The setting remains configurable with `--metric-hll-bucket-bits`. + ### Save Metrics (GPU → CPU) - **pegaflow_save_bytes_total** (Counter) - Total bytes saved from GPU to CPU storage @@ -245,15 +264,17 @@ tier counters. - Only used when `--metrics-otel-endpoint` is set - `--metric-hll-windows`: Comma-separated HLL sliding windows for estimated - prefix reuse (default: `15m,1h,24h`) + prefix reuse (default: `15m,1h,1d`) - Supported units: `s`, `m`, `h`, `d` - Each configured duration becomes a canonical `window` label. For example, the default config exports `window="15m"`, `window="1h"`, and `window="1d"`. - Empty entries such as `15m,,1h` and duplicate durations such as `1h,60m` are rejected at startup. -- `--metric-hll-bucket-bits`: HLL bucket index bits (default: `14`) - - Higher values use more memory and lower estimation error. +- `--metric-hll-bucket-bits`: HLL bucket index bits (default: `16`) + - `2^16 = 65,536` registers per window and about 0.4% standard error. + - Higher values use more memory and lower estimation error; `18` remains + the supported maximum. **Example: Prometheus Metrics** ```bash @@ -484,7 +505,10 @@ sum by (le) ( rate(pegaflow_cache_residence_duration_seconds_bucket{reason="pressure"}[5m]) ) -# HLL estimated hit rate for the 1h window +# HLL estimated hit rate for the 1h window (preferred) +pegaflow_hll_estimated_hit_rate{window="1h"} + +# Backward-compatible derivation from the retained gauges 1 - ( pegaflow_hll_cardinality{window="1h"} / diff --git a/pegaflow-common/src/hll.rs b/pegaflow-common/src/hll.rs index 924d4182..cc817b19 100644 --- a/pegaflow-common/src/hll.rs +++ b/pegaflow-common/src/hll.rs @@ -22,6 +22,7 @@ pub const MAX_BUCKET_BITS: u8 = 18; /// Input hashes are expected to have good uniformity (e.g. SHA-256). /// The full hash is used: top `bucket_bits` bits select the register, /// remaining bits are scanned for leading zeros. +#[derive(Debug)] pub struct HyperLogLog { registers: Vec, bucket_bits: u8, @@ -34,8 +35,8 @@ impl HyperLogLog { /// Create a new HyperLogLog with the given bucket bits. /// /// `bucket_bits` determines the number of buckets (2^bucket_bits) and estimation - /// accuracy (~1.04 / sqrt(2^bucket_bits)). 14 gives 16384 buckets - /// and ~0.8% standard error. + /// accuracy (~1.04 / sqrt(2^bucket_bits)). 16 gives 65536 buckets + /// and ~0.4% standard error. pub fn new(bucket_bits: u8) -> Self { assert!( (MIN_BUCKET_BITS..=MAX_BUCKET_BITS).contains(&bucket_bits), @@ -205,7 +206,7 @@ fn alpha_m(m: usize) -> f64 { /// Metric snapshot returned by [`HllTracker::metric`]. #[derive(Debug, Clone)] pub struct HllMetric { - /// Estimated number of distinct block hashes in the window. + /// Estimated number of distinct miss block identities in the window. pub cardinality: f64, /// Total block requests (including duplicates) in the window. pub total_requests: u64, @@ -225,7 +226,7 @@ struct WindowSlot { /// /// Divides time into fixed-duration slots and maintains a ring of HLLs. /// The merged cardinality across all active slots approximates the number -/// of distinct blocks requested in the window. From this we derive: +/// of distinct miss blocks recorded in the window. From this we derive: /// /// ```text /// hit_rate = (total_requests - cardinality) / total_requests @@ -249,7 +250,7 @@ impl HllTracker { /// /// - `slot_duration`: how long each time slot lasts (e.g. 1 hour) /// - `window_duration`: total sliding window (e.g. 24 hours) - /// - `bucket_bits`: HLL bucket index bits (4..=18, default 14) + /// - `bucket_bits`: HLL bucket index bits (4..=18, default 16) pub fn new(slot_duration: Duration, window_duration: Duration, bucket_bits: u8) -> Self { Self { slots: VecDeque::new(), @@ -268,6 +269,20 @@ impl HllTracker { /// time drift. For example with 1h slots: if the first slot starts at 0:00 /// and the next request arrives at 1:30, the new slot starts at 1:00 (not 1:30). pub fn record(&mut self, hash: &[u8]) { + let hashes = [hash]; + self.record_hashes_with_total(&hashes, 1); + } + + /// Record a batch of distinct identities while accounting for a possibly + /// larger observation count. The identities are inserted into HLL, while + /// `total_requests` is used as the denominator. This is used by the + /// miss-only reference: only cache misses enter HLL, but every queried + /// block still contributes to the total observation count. + pub fn record_hashes_with_total>(&mut self, hashes: &[T], total_requests: u64) { + if total_requests == 0 { + return; + } + let now = Instant::now(); let need_new_slot = match self.slots.back() { @@ -295,15 +310,15 @@ impl HllTracker { } let slot = self.slots.back_mut().unwrap(); - slot.hll.insert(hash); - slot.request_count += 1; + for hash in hashes { + slot.hll.insert(hash.as_ref()); + } + slot.request_count += total_requests; } /// Record a batch of block hashes from a gRPC request. pub fn record_hashes(&mut self, hashes: &[Vec]) { - for hash in hashes { - self.record(hash); - } + self.record_hashes_with_total(hashes, hashes.len() as u64); } /// Compute and return the current metric snapshot. @@ -417,6 +432,30 @@ impl MultiWindowHllTracker { } } + /// Record all queried blocks in the denominator but insert only the + /// identities that were misses into HLL. Repeated misses remain deduped by + /// HLL, preserving the infinite-cache reuse reference without discarding + /// historical observations. + pub fn record_namespaced_misses( + &mut self, + namespace: &str, + total_requests: u64, + miss_hashes: &[Vec], + ) { + if total_requests == 0 { + return; + } + debug_assert!(miss_hashes.len() as u64 <= total_requests); + + let namespaced_hashes: Vec<[u8; 8]> = miss_hashes + .iter() + .map(|hash| namespaced_hash(namespace, hash)) + .collect(); + for (_, tracker) in &mut self.windows { + tracker.record_hashes_with_total(&namespaced_hashes, total_requests); + } + } + /// Snapshot every window. Returned in insertion order. pub fn metrics(&mut self) -> Vec<(String, HllMetric)> { self.windows @@ -426,6 +465,38 @@ impl MultiWindowHllTracker { } } +/// Stable identity hash for `(namespace, block_hash)`. +/// +/// Cluster aggregation unions raw HLL registers across nodes, so every node +/// must map the same object to the exact same bits regardless of platform, +/// architecture, or build. FNV-1a with a splitmix64 finalizer is fully +/// specified byte-by-byte; do not replace it with a hasher that does not +/// guarantee cross-platform stability (e.g. ahash, SipHash with random keys). +fn namespaced_hash(namespace: &str, block_hash: &[u8]) -> [u8; 8] { + const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325; + const FNV_PRIME: u64 = 0x0000_0100_0000_01b3; + let mut h = FNV_OFFSET; + // Length prefix keeps (namespace, hash) boundaries unambiguous. + for &byte in (namespace.len() as u64) + .to_le_bytes() + .iter() + .chain(namespace.as_bytes()) + .chain(block_hash) + { + h = (h ^ u64::from(byte)).wrapping_mul(FNV_PRIME); + } + splitmix64(h).to_be_bytes() +} + +/// splitmix64 finalizer: strengthens FNV-1a's avalanche so the top bits used +/// for HLL bucket selection are uniformly distributed. +fn splitmix64(mut x: u64) -> u64 { + x = x.wrapping_add(0x9e37_79b9_7f4a_7c15); + x = (x ^ (x >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + x = (x ^ (x >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + x ^ (x >> 31) +} + fn derive_slot_duration(window: Duration) -> Duration { const MIN_SLOT: Duration = Duration::from_secs(60); const MAX_SLOT: Duration = Duration::from_secs(3600); @@ -766,7 +837,7 @@ mod tests { vec![ ("15m".into(), Duration::from_secs(15 * 60)), ("1h".into(), Duration::from_secs(3600)), - ("24h".into(), Duration::from_secs(86400)), + ("1d".into(), Duration::from_secs(86400)), ], 14, ); @@ -781,7 +852,7 @@ mod tests { assert_eq!(metrics.len(), 3); assert_eq!(metrics[0].0, "15m"); assert_eq!(metrics[1].0, "1h"); - assert_eq!(metrics[2].0, "24h"); + assert_eq!(metrics[2].0, "1d"); for (label, m) in metrics { assert_eq!(m.total_requests, 2000, "{label}: total"); assert!( @@ -792,6 +863,44 @@ mod tests { } } + #[test] + fn namespaced_misses_keep_all_observations_in_denominator() { + let mut tracker = + MultiWindowHllTracker::new(vec![("15m".into(), Duration::from_secs(15 * 60))], 14); + let miss = vec![sha256_like(1).to_vec(), sha256_like(2).to_vec()]; + tracker.record_namespaced_misses("model", 5, &miss); + + let metric = tracker.metrics().remove(0).1; + assert_eq!(metric.total_requests, 5); + assert!(metric.cardinality > 1.0 && metric.cardinality < 3.5); + assert!(metric.estimated_hit_rate > 0.3); + } + + #[test] + fn namespaced_misses_allow_all_hit_observation_without_hll_insert() { + let mut tracker = + MultiWindowHllTracker::new(vec![("15m".into(), Duration::from_secs(15 * 60))], 14); + tracker.record_namespaced_misses("model", 4, &[]); + + let metric = tracker.metrics().remove(0).1; + assert_eq!(metric.total_requests, 4); + assert_eq!(metric.cardinality, 0.0); + assert_eq!(metric.estimated_hit_rate, 1.0); + } + + #[test] + fn namespaced_misses_separate_equal_raw_hashes() { + let mut tracker = + MultiWindowHllTracker::new(vec![("15m".into(), Duration::from_secs(15 * 60))], 16); + let raw_hash = vec![42; 32]; + tracker.record_namespaced_misses("model-a", 1, std::slice::from_ref(&raw_hash)); + tracker.record_namespaced_misses("model-b", 1, std::slice::from_ref(&raw_hash)); + + let metric = tracker.metrics().remove(0).1; + assert_eq!(metric.total_requests, 2); + assert!((1.5..2.5).contains(&metric.cardinality)); + } + #[test] fn derive_slot_clamps() { assert_eq!( @@ -849,11 +958,4 @@ mod tests { hash[24..32].copy_from_slice(&m3.to_le_bytes()); hash } - - fn splitmix64(mut x: u64) -> u64 { - x = x.wrapping_add(0x9e3779b97f4a7c15); - x = (x ^ (x >> 30)).wrapping_mul(0xbf58476d1ce4e5b9); - x = (x ^ (x >> 27)).wrapping_mul(0x94d049bb133111eb); - x ^ (x >> 31) - } } diff --git a/pegaflow-core/src/lib.rs b/pegaflow-core/src/lib.rs index be7020c4..7c15b49b 100644 --- a/pegaflow-core/src/lib.rs +++ b/pegaflow-core/src/lib.rs @@ -483,6 +483,11 @@ impl PegaEngine { .collect() } + /// Return the namespace associated with a registered instance. + pub fn instance_namespace(&self, instance_id: &str) -> Result { + Ok(self.get_instance(instance_id)?.namespace().to_string()) + } + /// Count prefix hit blocks with SSD prefetch support. /// /// Argument contract: diff --git a/pegaflow-server/src/lib.rs b/pegaflow-server/src/lib.rs index e5117d20..97321513 100644 --- a/pegaflow-server/src/lib.rs +++ b/pegaflow-server/src/lib.rs @@ -174,11 +174,11 @@ pub struct Cli { /// HLL sliding-window list for hit-rate estimation. Comma-separated humantime /// durations; each becomes a canonical `window` label in metrics (e.g. `15m,1h,1d`). /// Slot duration is derived as `clamp(window/24, 1min, 1h)`. - #[arg(long, default_value = "15m,1h,24h", value_parser = parse_hll_windows_arg)] + #[arg(long, default_value = "15m,1h,1d", value_parser = parse_hll_windows_arg)] pub metric_hll_windows: String, - /// HLL bucket index bits 4–18 (default: 14 → 16384 buckets, ~0.8% error) - #[arg(long, default_value_t = 14, value_parser = parse_hll_bucket_bits)] + /// HLL bucket index bits 4–18 (default: 16 → 65536 buckets, ~0.4% error) + #[arg(long, default_value_t = 16, value_parser = parse_hll_bucket_bits)] pub metric_hll_bucket_bits: u8, /// Transfer lock timeout in seconds. Blocks held for cross-node RDMA transfer are @@ -211,7 +211,7 @@ fn parse_hll_windows_arg(s: &str) -> Result { Ok(s.to_string()) } -/// Parse a comma-separated list of humantime windows (e.g. `15m,1h,24h`). +/// Parse a comma-separated list of humantime windows (e.g. `15m,1h,1d`). /// Each entry becomes `(label, duration)` where label is canonicalized from /// the parsed duration. fn parse_hll_windows(s: &str) -> Result, String> { @@ -732,6 +732,7 @@ mod tests { parse_hll_windows(&cli.metric_hll_windows).unwrap(), expected_hll_windows() ); + assert_eq!(cli.metric_hll_bucket_bits, 16); } #[test] diff --git a/pegaflow-server/src/metric.rs b/pegaflow-server/src/metric.rs index fc4b8f8b..920edf39 100644 --- a/pegaflow-server/src/metric.rs +++ b/pegaflow-server/src/metric.rs @@ -8,6 +8,7 @@ use tonic::Status; struct HllGaugeHandles { _cardinality: ObservableGauge, _total_requests: ObservableGauge, + _estimated_hit_rate: ObservableGauge, } static HLL_GAUGES: OnceLock = OnceLock::new(); @@ -15,18 +16,20 @@ static HLL_GAUGES: OnceLock = OnceLock::new(); /// Register HLL observable gauges backed by the multi-window tracker. /// /// Emits one time series per configured window, labeled with `window=