[KV Offload] Emit self-describing events for recurrent state groups

docs/features/kv_offloading_usage.md

@@ -113,7 +113,7 @@ Secondary tiers are not supported yet on engines whose workers span several node
 | `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 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). |
 

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,84 @@ def test_take_events_opt_out_keeps_placeholders():
     assert len(events[3].block_hashes) == 3
 
 
+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
+    )
+    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 == []
+    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
+
+
+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_state_checkpoint(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(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
+
+
 @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()

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 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.
 """
 
 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,39 @@ 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 _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.
 
@@ -132,12 +174,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 +196,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 +218,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 +230,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)
@@ -247,16 +289,18 @@ 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
 
     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
@@ -292,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
@@ -314,9 +359,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:
@@ -328,15 +376,20 @@ 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
         if req.lora_request is not None:
             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,
@@ -348,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,
@@ -394,7 +450,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."
                     )