From f3bc5bb540a71e74859a68ba0edc2c02a7332e41 Mon Sep 17 00:00:00 2001 From: doxu Date: Mon, 19 Oct 2026 21:50:32 -0700 Subject: [PATCH 1/2] [KV Offload] Emit self-describing events for recurrent state groups OffloadingEventsTracker treated every group with a load window as a sliding-window group and emitted a block_size=0 placeholder BlockStored for it. MambaSpec groups have a one-chunk window because they depend on a single state, not because their hashes are meaningless: the offloaded chunk is the recurrent state checkpoint at the chunk's last hash. Consumers that index offload residency (e.g. a KV-aware router) therefore could not see where a hybrid prefix is reusable and overestimated it from the attention group. Describe recurrent groups like full-attention chunks (hashes, tokens, parent, block_size, group_idx) with kv_cache_spec_kind "mamba"; windowed attention groups keep the placeholder. Signed-off-by: doxu --- docs/features/kv_offloading_usage.md | 2 +- .../unit/offloading_connector/test_events.py | 90 +++++++++++++++++-- .../kv_connector/v1/offloading/events.py | 48 +++++++--- 3 files changed, 118 insertions(+), 22 deletions(-) diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index 7450e0dacc13..351da4ab828d 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -111,7 +111,7 @@ vllm serve \ | `tier_poll_interval_s` | no | `0.001` | multi-tier | Pause between tier-poll rounds for tiers that must be serviced between engine steps (currently `p2p`). See [Threading model](#threading-model). Set to `0` to disable the thread and service those tiers once per engine step instead. | | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. | | `sync_load_budget_bytes` | no | `0` | both | Per-step budget of KV bytes (per GPU) loaded within the same step, so those requests run immediately instead of waiting a step for an async load. Hits that don't fit the remaining budget, or would leave nothing to compute, stay async. `0` disables. | -| `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Partial recurrent tails emit the hash-aligned portion from the physical block start through the tail boundary. Other sliding-window/SSM chunks keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. | +| `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Recurrent (Mamba / linear-attention) state groups are self-describing too, with `kv_cache_spec_kind` `mamba`: the payload lists the chunk's hashes and tokens like an attention chunk, but the residency it announces is the state checkpoint at the chunk's last hash, so a hybrid prefix is reusable only up to a boundary where every group has a residency. Partial recurrent tails emit the hash-aligned portion from the physical block start through the tail boundary. Sliding-window and chunked-local attention chunks keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. | | `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). | | `canonical_layout` | no | `false` | both | Store each offloaded block in a common canonical representation instead of each worker's local GPU page (enables sharing across TP sizes and single-copy storage of replicated KV). See [Canonical Layout](#canonical-layout). | diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_events.py b/tests/v1/kv_connector/unit/offloading_connector/test_events.py index 45c5711306d1..b9a7db7037c3 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_events.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_events.py @@ -29,7 +29,9 @@ FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, + KVCacheSpec, KVCacheSpecKind, + MambaSpec, ) from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm.v1.kv_offload.base import ( @@ -50,6 +52,10 @@ kv_cache_spec_kind=KVCacheSpecKind.FULL_ATTENTION.value, kv_cache_spec_sliding_window=None, ) +_MAMBA_EVENT_SPEC = OffloadingEventGroupSpec( + kv_cache_spec_kind=KVCacheSpecKind.MAMBA.value, + kv_cache_spec_sliding_window=None, +) def _tracker( @@ -89,17 +95,31 @@ def _group_config( blocks_per_chunk: int = 1, tokens_per_hash: int | None = None, sliding_window_size_in_chunks: int | None = None, + mamba: bool = False, ) -> GroupOffloadConfig: if tokens_per_hash is None: tokens_per_hash = block_size tokens_per_chunk = block_size * blocks_per_chunk assert tokens_per_chunk % tokens_per_hash == 0 - kv_spec = FullAttentionSpec( - block_size=block_size, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - ) + kv_spec: KVCacheSpec + if mamba: + # A recurrent group: one state per chunk, offloaded in align mode. + kv_spec = MambaSpec( + block_size=block_size, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + sliding_window_size_in_chunks = 1 + event_spec = _MAMBA_EVENT_SPEC + else: + kv_spec = FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ) + event_spec = _FULL_ATTENTION_EVENT_SPEC manager_cls = KVCacheSpecRegistry.get_manager_class(kv_spec) assert manager_cls is not None return GroupOffloadConfig( @@ -110,7 +130,8 @@ def _group_config( sliding_window_size_in_chunks=sliding_window_size_in_chunks, kv_cache_spec=kv_spec, manager_cls=manager_cls, - kv_event_group_spec=_FULL_ATTENTION_EVENT_SPEC, + kv_event_group_spec=event_spec, + requires_cow_source=mamba, ) @@ -515,12 +536,63 @@ def test_take_events_opt_out_keeps_placeholders(): assert len(events[3].block_hashes) == 3 +def test_mamba_chunk_store_emits_self_describing_event(): + """A recurrent group's chunk is announced like an attention chunk, with + kv_cache_spec_kind "mamba", so a consumer can align it by hash.""" + tracker = _tracker() + group_config = _group_config( + group_idx=1, block_size=8, tokens_per_hash=4, mamba=True + ) + req = _request(block_hashes=[_hash(i) for i in range(4)], token_count=16) + keys = _record_chunks(tracker, req, group_config, num_chunks=2) + + events = list(tracker.take_events([_stored_event(keys[1:])])) + + assert len(events) == 1 + event = events[0] + assert isinstance(event, BlockStored) + assert event.block_hashes == [_wire_hash(_hash(3))] + assert event.parent_block_hash == _wire_hash(_hash(1)) + assert event.token_ids == list(range(9, 17)) + assert event.block_size == 8 + assert event.group_idx == 1 + assert event.kv_cache_spec_kind == KVCacheSpecKind.MAMBA.value + assert event.kv_cache_spec_sliding_window is None + + [removed] = tracker.take_events([_removed_event(keys[1:])]) + assert isinstance(removed, BlockRemoved) + assert removed.block_hashes == [_wire_hash(_hash(3))] + assert removed.group_idx == 1 + + +@pytest.mark.parametrize( + "record_method", ["record_partial_store", "record_partial_lookup"] +) +def test_mamba_partial_tail_event_is_self_describing(record_method): + tracker = _tracker() + group_config = _group_config( + block_size=16, blocks_per_chunk=1, mamba=True + )._replace(hashes_per_chunk=4) + req = _request(block_hashes=[_hash(i) for i in range(8)], token_count=32) + key = make_offload_key(req.block_hashes[6], group_config.group_idx) + + getattr(tracker, record_method)(req, group_config, 28, key) + [event] = tracker.take_events([_stored_event([key])]) + + assert isinstance(event, BlockStored) + assert event.block_hashes == [_wire_hash(_hash(i)) for i in range(4, 7)] + assert event.parent_block_hash == _wire_hash(_hash(3)) + assert event.token_ids == list(range(17, 29)) + assert event.block_size == 4 + assert event.kv_cache_spec_kind == KVCacheSpecKind.MAMBA.value + + @pytest.mark.parametrize( "sliding_window_size_in_chunks", [1, 2], - ids=["ssm", "sliding-window"], + ids=["one-chunk-window", "sliding-window"], ) -def test_event_metadata_skips_non_full_attention_group( +def test_event_metadata_skips_windowed_attention_group( sliding_window_size_in_chunks: int, ): tracker = _tracker() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py index 23fe932451d3..fa48595706d7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py @@ -13,6 +13,14 @@ of the same hash. Opt-in via ``kv_connector_extra_config["self_describing_kv_events"]``; inert unless KV cache events are enabled. See the PR description for the full design. + +Recurrent (Mamba / linear-attention) state groups are described the same way +as full-attention chunks: the payload lists the chunk's hashes and tokens so a +consumer can place it at the same radix position as the attention blocks, and +``kv_cache_spec_kind`` is ``"mamba"``. The residency it announces is the state +checkpoint at the chunk's last hash, not a reusable token range: a hybrid +prefix is reusable only up to a boundary where every group has a residency. +Windowed attention groups keep the legacy placeholder payload. """ from collections.abc import Iterable @@ -35,6 +43,7 @@ ) from vllm.v1.kv_cache_interface import ( KVCacheGroupSpec, + MambaSpec, get_kv_cache_spec_kind, get_kv_cache_spec_sliding_window, ) @@ -99,6 +108,18 @@ class _OffloadEventMetadata: active_residencies: set[tuple[Medium, str | None]] +def describes_group(group_config: "GroupOffloadConfig") -> bool: + """Whether the tracker builds self-describing payloads for this group. + + Full-attention groups and recurrent state groups (``MambaSpec``) do; + windowed attention groups (sliding window, chunked local) keep the legacy + placeholder because a window's hashes do not identify a reusable prefix. + """ + if group_config.sliding_window_size_in_chunks is None: + return True + return isinstance(group_config.kv_cache_spec, MambaSpec) + + class OffloadingEventsTracker: """Tracks offloaded chunks' KV event payloads from store to eviction. @@ -132,12 +153,12 @@ def record_store( ) -> None: """Snapshot the KV cache event payload for one offloaded chunk. - No-op when self-describing event capture is disabled or for - sliding-window / SSM groups, which keep the legacy placeholder payload. + No-op when self-describing event capture is disabled or for windowed + attention groups, which keep the legacy placeholder payload. """ if not self.self_describing_enabled: return - if group_config.sliding_window_size_in_chunks is not None: + if not describes_group(group_config): return meta = self._build_event_metadata(req, group_config, chunk_idx) if existing := self._pending_event_metadata.get(offload_key): @@ -154,7 +175,7 @@ def record_lookup( """Snapshot metadata for a ready or write-pending primary-tier hit.""" if not self.self_describing_enabled: return - if group_config.sliding_window_size_in_chunks is not None: + if not describes_group(group_config): return if offload_key not in self._pending_event_metadata: self._pending_event_metadata[offload_key] = self._build_event_metadata( @@ -176,7 +197,7 @@ def record_partial_store( offload_key: OffloadKey, ) -> None: """Snapshot metadata for a newly stored partial recurrent tail.""" - if group_config.sliding_window_size_in_chunks is not None: + if not describes_group(group_config): return self._record_partial_tail(req, group_config, boundary_tokens, offload_key) @@ -188,7 +209,7 @@ def record_partial_lookup( offload_key: OffloadKey, ) -> None: """Backfill metadata for a partial recurrent tail lookup hit.""" - if group_config.sliding_window_size_in_chunks is not None: + if not describes_group(group_config): return if offload_key not in self._pending_event_metadata: self._record_partial_tail(req, group_config, boundary_tokens, offload_key) @@ -254,9 +275,9 @@ def _record_partial_tail( def take_events(self, events: Iterable[OffloadingEvent]) -> Iterable[KVCacheEvent]: """Translate raw OffloadingEvents into self-describing KV events. - Complete metadata is available only for full-attention groups when - the tracker is enabled. Other shapes retain the legacy placeholder - payload so consumers can ignore them. + Complete metadata is available for full-attention and recurrent + state groups when the tracker is enabled. Windowed attention groups + retain the legacy placeholder payload so consumers can ignore them. Yields: ``BlockStored`` or ``BlockRemoved`` events corresponding to @@ -314,9 +335,12 @@ def _build_event_metadata( group_config.tokens_per_chunk // group_config.tokens_per_block ) - if group_config.sliding_window_size_in_chunks is not None: + if not describes_group(group_config): # The recording methods filter these out before calling this helper. - raise AssertionError("self-describing events only support full attention") + raise AssertionError( + "self-describing events only support full-attention and " + "recurrent state groups" + ) parent_block_hash: BlockHash | None if first_hash_idx == 0: @@ -394,7 +418,7 @@ def _take_stored_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]: logger.warning_once( "OffloadingEventsTracker: no event metadata for " "offload key during BlockStored emission; emitting a " - "placeholder payload. Expected for non-full-attention " + "placeholder payload. Expected for windowed attention " "groups and promotions not observed as a primary-tier " "hit before translation." ) From 4172118ded34c0ba9914dc64cad2bd8623c10e9e Mon Sep 17 00:00:00 2001 From: doxu Date: Tue, 27 Oct 2026 14:13:36 -0700 Subject: [PATCH 2/2] [KV Offload] Announce recurrent groups as single state checkpoints A recurrent group's offloaded chunk holds one state, valid at the chunk's last hash. Emitting the whole hash range and token span let a consumer treat every hash in a multi-hash chunk as a resumable position, and kept a per-chunk token copy alive for each resident state. The BlockStored now carries that last hash, its parent, and no tokens; BlockRemoved fans out to the same single hash. Attention groups are unchanged. Signed-off-by: doxu --- docs/features/kv_offloading_usage.md | 2 +- .../unit/offloading_connector/test_events.py | 37 ++++++++++---- .../kv_connector/v1/offloading/events.py | 48 +++++++++++++++---- 3 files changed, 70 insertions(+), 17 deletions(-) diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index 351da4ab828d..fe84c9087368 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -111,7 +111,7 @@ vllm serve \ | `tier_poll_interval_s` | no | `0.001` | multi-tier | Pause between tier-poll rounds for tiers that must be serviced between engine steps (currently `p2p`). See [Threading model](#threading-model). Set to `0` to disable the thread and service those tiers once per engine step instead. | | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. | | `sync_load_budget_bytes` | no | `0` | both | Per-step budget of KV bytes (per GPU) loaded within the same step, so those requests run immediately instead of waiting a step for an async load. Hits that don't fit the remaining budget, or would leave nothing to compute, stay async. `0` disables. | -| `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Recurrent (Mamba / linear-attention) state groups are self-describing too, with `kv_cache_spec_kind` `mamba`: the payload lists the chunk's hashes and tokens like an attention chunk, but the residency it announces is the state checkpoint at the chunk's last hash, so a hybrid prefix is reusable only up to a boundary where every group has a residency. Partial recurrent tails emit the hash-aligned portion from the physical block start through the tail boundary. Sliding-window and chunked-local attention chunks keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. | +| `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Recurrent (Mamba / linear-attention) state groups announce state checkpoints, with `kv_cache_spec_kind` `mamba`: an offloaded chunk or partial tail holds one recurrent state, valid at its last hash, so its `BlockStored` carries that single hash, its parent hash and an empty `token_ids`. Consumers align it by hash with the attention group's blocks; a hybrid prefix is reusable only up to a boundary where every group has a residency. Sliding-window and chunked-local attention chunks keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. | | `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). | | `canonical_layout` | no | `false` | both | Store each offloaded block in a common canonical representation instead of each worker's local GPU page (enables sharing across TP sizes and single-copy storage of replicated KV). See [Canonical Layout](#canonical-layout). | diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_events.py b/tests/v1/kv_connector/unit/offloading_connector/test_events.py index b9a7db7037c3..d8c767f0c360 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_events.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_events.py @@ -536,9 +536,10 @@ def test_take_events_opt_out_keeps_placeholders(): assert len(events[3].block_hashes) == 3 -def test_mamba_chunk_store_emits_self_describing_event(): - """A recurrent group's chunk is announced like an attention chunk, with - kv_cache_spec_kind "mamba", so a consumer can align it by hash.""" +def test_mamba_chunk_store_emits_state_checkpoint_event(): + """A recurrent group's chunk is announced as one state checkpoint: the + chunk's last hash, chained to its predecessor, with no tokens and + kv_cache_spec_kind "mamba", so a consumer aligns it by hash.""" tracker = _tracker() group_config = _group_config( group_idx=1, block_size=8, tokens_per_hash=4, mamba=True @@ -553,7 +554,7 @@ def test_mamba_chunk_store_emits_self_describing_event(): assert isinstance(event, BlockStored) assert event.block_hashes == [_wire_hash(_hash(3))] assert event.parent_block_hash == _wire_hash(_hash(1)) - assert event.token_ids == list(range(9, 17)) + assert event.token_ids == [] assert event.block_size == 8 assert event.group_idx == 1 assert event.kv_cache_spec_kind == KVCacheSpecKind.MAMBA.value @@ -565,10 +566,30 @@ def test_mamba_chunk_store_emits_self_describing_event(): assert removed.group_idx == 1 +def test_mamba_multi_hash_chunk_announces_only_last_hash(): + """In chunk mode a recurrent chunk spans several hashes, but only the + last one is a resumable checkpoint; the others are not announced.""" + tracker = _tracker() + group_config = _group_config(block_size=4, blocks_per_chunk=4, mamba=True) + req = _request(block_hashes=[_hash(i) for i in range(8)], token_count=32) + keys = _record_chunks(tracker, req, group_config, num_chunks=2) + + [event] = tracker.take_events([_stored_event(keys[1:])]) + assert isinstance(event, BlockStored) + assert event.block_hashes == [_wire_hash(_hash(7))] + assert event.parent_block_hash == _wire_hash(_hash(6)) + assert event.token_ids == [] + assert event.block_size == 4 + + [removed] = tracker.take_events([_removed_event(keys[1:])]) + assert isinstance(removed, BlockRemoved) + assert removed.block_hashes == [_wire_hash(_hash(7))] + + @pytest.mark.parametrize( "record_method", ["record_partial_store", "record_partial_lookup"] ) -def test_mamba_partial_tail_event_is_self_describing(record_method): +def test_mamba_partial_tail_event_is_state_checkpoint(record_method): tracker = _tracker() group_config = _group_config( block_size=16, blocks_per_chunk=1, mamba=True @@ -580,9 +601,9 @@ def test_mamba_partial_tail_event_is_self_describing(record_method): [event] = tracker.take_events([_stored_event([key])]) assert isinstance(event, BlockStored) - assert event.block_hashes == [_wire_hash(_hash(i)) for i in range(4, 7)] - assert event.parent_block_hash == _wire_hash(_hash(3)) - assert event.token_ids == list(range(17, 29)) + assert event.block_hashes == [_wire_hash(_hash(6))] + assert event.parent_block_hash == _wire_hash(_hash(5)) + assert event.token_ids == [] assert event.block_size == 4 assert event.kv_cache_spec_kind == KVCacheSpecKind.MAMBA.value diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py index fa48595706d7..99d9fc8b8e9e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py @@ -14,11 +14,11 @@ ``kv_connector_extra_config["self_describing_kv_events"]``; inert unless KV cache events are enabled. See the PR description for the full design. -Recurrent (Mamba / linear-attention) state groups are described the same way -as full-attention chunks: the payload lists the chunk's hashes and tokens so a -consumer can place it at the same radix position as the attention blocks, and -``kv_cache_spec_kind`` is ``"mamba"``. The residency it announces is the state -checkpoint at the chunk's last hash, not a reusable token range: a hybrid +Recurrent (Mamba / linear-attention) state groups announce state checkpoints, +not token ranges: an offloaded chunk holds the recurrent state at the chunk's +last hash only. Their ``BlockStored`` therefore carries that single hash, its +parent, an empty ``token_ids`` and ``kv_cache_spec_kind`` ``"mamba"``. A +consumer aligns it by hash with the attention group's blocks, and a hybrid prefix is reusable only up to a boundary where every group has a residency. Windowed attention groups keep the legacy placeholder payload. """ @@ -117,9 +117,30 @@ def describes_group(group_config: "GroupOffloadConfig") -> bool: """ if group_config.sliding_window_size_in_chunks is None: return True + return _is_recurrent_group(group_config) + + +def _is_recurrent_group(group_config: "GroupOffloadConfig") -> bool: return isinstance(group_config.kv_cache_spec, MambaSpec) +def _as_state_checkpoint(meta: _OffloadEventMetadata) -> _OffloadEventMetadata: + """Reduce a chunk payload to the recurrent state checkpoint it holds. + + A recurrent group's offloaded chunk is one state, valid at the chunk's + last hash. Announcing the whole hash range would let a consumer treat + every hash in it as a resumable position, so the payload keeps only the + last hash (the OffloadKey), chained to the hash before it, and no tokens: + the consumer aligns it with the attention group's blocks by hash. + """ + assert meta.block_hashes + if len(meta.block_hashes) > 1: + meta.parent_block_hash = meta.block_hashes[-2] + meta.block_hashes = meta.block_hashes[-1:] + meta.token_ids = () + return meta + + class OffloadingEventsTracker: """Tracks offloaded chunks' KV event payloads from store to eviction. @@ -268,6 +289,8 @@ def _record_partial_tail( kv_cache_spec=group_config.kv_event_group_spec, active_residencies={_PRIMARY_RESIDENCY}, ) + if _is_recurrent_group(group_config): + meta = _as_state_checkpoint(meta) if existing := self._pending_event_metadata.get(offload_key): meta.active_residencies.update(existing.active_residencies) self._pending_event_metadata[offload_key] = meta @@ -313,7 +336,8 @@ def _build_event_metadata( ) -> _OffloadEventMetadata: """Build the payload snapshot for one offloaded chunk: its constituent per-block hashes, the whole chunk's tokens, and the - per-block ``block_size``.""" + per-block ``block_size``. A recurrent group's chunk is reduced to + its state checkpoint (see :func:`_as_state_checkpoint`).""" hashes_per_chunk = group_config.hashes_per_chunk assert hashes_per_chunk > 0 assert chunk_idx >= 0 @@ -352,7 +376,12 @@ def _build_event_metadata( tok_start = chunk_idx * group_config.tokens_per_chunk tok_end = tok_start + group_config.tokens_per_chunk assert tok_end <= len(req.all_token_ids) - token_ids = tuple(req.all_token_ids[tok_start:tok_end]) + # A recurrent chunk's payload carries no tokens; skip the copy. + token_ids: tuple[int, ...] = ( + () + if _is_recurrent_group(group_config) + else tuple(req.all_token_ids[tok_start:tok_end]) + ) lora_id: int | None = None lora_name: str | None = None @@ -360,7 +389,7 @@ def _build_event_metadata( lora_id = req.lora_request.adapter_id lora_name = req.lora_request.name - return _OffloadEventMetadata( + meta = _OffloadEventMetadata( block_hashes=tuple(chunk_hashes), parent_block_hash=parent_block_hash, token_ids=token_ids, @@ -372,6 +401,9 @@ def _build_event_metadata( kv_cache_spec=group_config.kv_event_group_spec, active_residencies={_PRIMARY_RESIDENCY}, ) + if _is_recurrent_group(group_config): + meta = _as_state_checkpoint(meta) + return meta def _placeholder_stored( self,