fix(mp): preserve async store events and opaque page tails

csrc/cuda/mp_mem_kernels.cu

@@ -447,8 +447,6 @@ void multi_layer_block_kv_transfer(
     PageBufferShapeDesc shape_desc, int lmcache_chunk_size,
     EngineKVFormat engine_kv_format, int skip_prefix_n_blocks) {
   int head_bytes = shape_desc.hs * shape_desc.element_size;
-  TORCH_CHECK(head_bytes % sizeof(uint16_t) == 0, "head_size * element_size (",
-              head_bytes, ") must be divisible by 2 for vectorized access");
 
   if (engine_kv_format == EngineKVFormat::NL_X_NB_BSV_BSS) {
     // Blocked-scale indexer cache: the per-token fp32 scale must be a whole
@@ -464,8 +462,14 @@ void multi_layer_block_kv_transfer(
     LAUNCH_TEMPLATED(uint4);  // 16 bytes per copy
   } else if (head_bytes % sizeof(uint32_t) == 0) {
     LAUNCH_TEMPLATED(uint32_t);  // 4 bytes per copy
+  } else if (head_bytes % sizeof(uint16_t) == 0) {
+    LAUNCH_TEMPLATED(uint16_t);  // 2 bytes per copy
   } else {
-    LAUNCH_TEMPLATED(uint16_t);  // 2 bytes per copy (minimum granularity)
+    // Opaque model-owned page tails can make a logical row byte-odd (for
+    // example GLM-5.3 C4's 561-byte FP8 row). The page and object are still
+    // byte-addressable, so retain correctness with a scalar-byte fallback
+    // instead of rejecting the whole store or truncating the tail.
+    LAUNCH_TEMPLATED(uint8_t);  // 1 byte per copy (correctness fallback)
   }
 }
 

lmcache/integration/vllm/kv_cache_group_edits.py

@@ -443,9 +443,13 @@ def apply(
         groups in the shared pool.
 
         head_size = ceil(row / block_size), rounded up to the kernels'
-        vector alignment, and block_size * head_size may exceed the row
-        by at most this layer's own page padding
-        (spec.page_size_bytes), never reaching sibling bytes.
+        vector alignment.  If the page has enough padding, block_size *
+        head_size may exceed the semantic row while remaining inside this
+        layer's declared page.  Some exact-fit pages cannot be split into
+        vector-aligned token rows (GLM-5.3-Flash DCP1 is 4096 * 265 bytes).
+        Those pages use one opaque physical slot instead; the adapter carries
+        block_size independently as tokens_per_block, so logical scheduling
+        and cache keys remain unchanged.
         """
         assert isinstance(kv_cache, torch.Tensor), (
             "single-layer KV cache must be a torch.Tensor"
@@ -475,27 +479,57 @@ def apply(
             if candidate * block_size * elem <= page_bytes:
                 head_size = candidate
                 break
-        if head_size == 0:
+        if head_size != 0:
+            return kv_cache.as_strided(
+                (num_blocks, block_size, head_size),
+                (kv_cache.stride(0), head_size, 1),
+            )
+
+        # The semantic row is still a valid opaque page even when it cannot
+        # be factored into vector-aligned per-token rows.  Preserve the whole
+        # declared page as one physical slot.  EngineGroupInfo retains the
+        # logical block_size, and KVLayerGroupsManager therefore maps the one
+        # slot back to exactly one engine block ID.
+        if page_bytes % elem:
+            raise ValueError(
+                f"declared Mamba page size {page_bytes} bytes is not aligned "
+                f"to tensor element size {elem}"
+            )
+        page_elems = page_bytes // elem
+        if page_elems < row:
+            raise ValueError(
+                f"declared Mamba page has {page_elems} elements but the state "
+                f"row requires {row}"
+            )
+        if page_elems > block_step:
             raise ValueError(
-                f"cannot tile a {row}-element state row into {block_size} "
-                f"aligned tokens within the {page_bytes}-byte page"
+                f"declared Mamba page has {page_elems} elements but the "
+                f"physical block stride is only {block_step}"
             )
         return kv_cache.as_strided(
-            (num_blocks, block_size, head_size),
-            (kv_cache.stride(0), head_size, 1),
+            (num_blocks, 1, page_elems),
+            (kv_cache.stride(0), page_elems, 1),
         )
 
 
 class _PaddedAttentionPageViewEdit(KVCacheGroupEdit):
     """Canonicalize a padded attention layer as an opaque rank-3 page.
 
     Current vLLM HMA layouts expose MLA as ``[B, H=1, N, C]`` and a replicated
-    DFlash page as e.g. ``[B, H=2, N=16, C=256]``. In both cases the inner page
-    is tightly packed, while sibling pages create a gap between dim-0 rows.
-    LMCache's opaque ``[B, N, C]`` format supports that authoritative padded
-    block stride. Re-factoring all inner page elements over the engine's
-    logical block size is a zero-copy view and preserves every payload byte;
+    DFlash page as e.g. ``[B, H=2, N=16, C=256]``. The semantic inner page is
+    tightly packed, while the declared physical page can also append opaque
+    model-owned state before sibling pages create the remaining gap between
+    dim-0 rows. LMCache's opaque ``[B, N, C]`` format supports that
+    authoritative padded block stride. The complete declared page is exposed
+    as a zero-copy view, preserving both semantic KV and any opaque page tail;
     the resulting dimensions are addressing metadata, not semantic K/V axes.
+
+    Ordinary pages retain one physical slot per token. Packed pages whose byte
+    count cannot be factored over the logical block size (for example GLM-5.3
+    NVFP4 KV, whose per-page scale/tail records are not token-uniform) use one
+    physical slot for the whole page. ``EngineGroupInfo.tokens_per_block``
+    independently retains the logical token count, so LMCache's compressed
+    geometry maps one block ID to one complete opaque page.
     """
 
     name = "padded-attention-page-view"
@@ -520,22 +554,42 @@ def apply(
         kv_cache: RegisteredKVCache,
         _layout_hints: LayoutHints,
     ) -> torch.Tensor:
-        """Return a padded-stride-preserving opaque ``[B, BS, HS]`` view.
+        """Return a padded-stride-preserving opaque ``[B, slots, width]`` view.
 
         Raises:
-            ValueError: If one physical page cannot be factored evenly over
-                the engine's logical block size.
+            ValueError: If the declared page is not element-aligned, is smaller
+                than the semantic tensor page, or exceeds the physical dim-0
+                stride.
         """
         assert isinstance(kv_cache, torch.Tensor)
-        page_elems = kv_cache.shape[1:].numel()
-        if page_elems % spec.block_size:
+        element_size = kv_cache.element_size()
+        page_bytes = spec.page_size_bytes
+        if page_bytes % element_size:
+            raise ValueError(
+                f"declared attention page size {page_bytes} bytes is not aligned "
+                f"to tensor element size {element_size}"
+            )
+        page_elems = page_bytes // element_size
+        semantic_page_elems = kv_cache.shape[1:].numel()
+        if page_elems < semantic_page_elems:
+            raise ValueError(
+                f"declared attention page has {page_elems} elements but the "
+                f"semantic tensor page requires {semantic_page_elems}"
+            )
+        if page_elems > kv_cache.stride(0):
             raise ValueError(
-                f"a {page_elems}-element attention page cannot be factored "
-                f"over block_size={spec.block_size}"
+                f"declared attention page has {page_elems} elements but the "
+                f"physical block stride is only {kv_cache.stride(0)}"
             )
-        hidden_size = page_elems // spec.block_size
+        # A packed page need not have a uniform byte width per logical token.
+        # Treat such a page as one opaque physical slot.  The vLLM adapter
+        # carries spec.block_size separately as tokens_per_block, so the group
+        # manager derives the correct compression ratio and still consumes one
+        # engine block ID per logical page.
+        physical_slots = spec.block_size if page_elems % spec.block_size == 0 else 1
+        hidden_size = page_elems // physical_slots
         return kv_cache.as_strided(
-            (kv_cache.shape[0], spec.block_size, hidden_size),
+            (kv_cache.shape[0], physical_slots, hidden_size),
             (kv_cache.stride(0), hidden_size, 1),
         )
 

lmcache/integration/vllm/lmcache_mp_metadata.py

@@ -566,11 +566,11 @@ def aggregate(
     ) -> "KVConnectorWorkerMetadata":
         assert isinstance(other, LMCacheMPWorkerMetadata)
         merged_requests = dict(self.completed_store_requests)
-        for k, v in other.completed_store_requests.items():
-            merged_requests[k] = merged_requests.get(k, 0) + v
+        for request_id, count in other.completed_store_requests.items():
+            merged_requests[request_id] = merged_requests.get(request_id, 0) + count
         merged_jobs = dict(self.completed_store_jobs)
-        for k, v in other.completed_store_jobs.items():
-            merged_jobs[k] = merged_jobs.get(k, 0) + v
+        for job_id, count in other.completed_store_jobs.items():
+            merged_jobs[job_id] = merged_jobs.get(job_id, 0) + count
         return LMCacheMPWorkerMetadata(
             completed_store_requests=merged_requests,
             completed_store_jobs=merged_jobs,

lmcache/v1/multiprocess/transfer_context/worker_transfer.py

@@ -670,6 +670,11 @@ def submit_store(
             RequestType.STORE,
             [key, instance_id, block_ids, event_ipc_handle],
         ).to_device_future(device=self._device)
+        # Multiple incremental stores for one request overwrite the adapter's
+        # request-keyed event slot. Tie every producer event to its own remote
+        # future so the exported IPC event remains valid until the sidecar has
+        # finished waiting on it and reading the corresponding GPU pages.
+        future.retain_reference(event)
         self._inflight_store_futures.add(future)
         return future
 

lmcache/v1/platform/cuda/cumem_ipc.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 """Share vLLM CUDA VMM allocations through same-UID POSIX file descriptors."""
 
+# Future
 from __future__ import annotations
 
 # Standard

tests/v1/multiprocess/test_engine_driven_transfer.py

@@ -572,8 +572,8 @@ def test_create_transfer_context_force_lmcache_driven_mode() -> None:
     assert isinstance(context, LMCacheDrivenTransferContext)
 
 
-def test_lmcache_driven_preemption_waits_for_remote_store_futures() -> None:
-    """Handle-path flush waits on remote completion, not the worker device."""
+def test_lmcache_driven_preemption_retains_each_store_event_and_waits() -> None:
+    """Two stores for one request retain both events until remote completion."""
     # First Party
     from lmcache.v1.multiprocess.transfer_context import (
         LMCacheDrivenTransferContext,
@@ -582,14 +582,24 @@ def test_lmcache_driven_preemption_waits_for_remote_store_futures() -> None:
 
     context = LMCacheDrivenTransferContext()
     registration_future = MagicMock(name="registration_future")
-    raw_store_future = MagicMock(name="raw_store_future")
-    pending = MagicMock(name="pending_store_future")
-    raw_store_future.to_device_future.return_value = pending
+    raw_store_futures = [
+        MagicMock(name="raw_store_future_1"),
+        MagicMock(name="raw_store_future_2"),
+    ]
+    pending = [
+        MagicMock(name="pending_store_future_1"),
+        MagicMock(name="pending_store_future_2"),
+    ]
+    for pending_future in pending:
+        pending_future.query.return_value = False
+    for raw_future, pending_future in zip(raw_store_futures, pending, strict=True):
+        raw_future.to_device_future.return_value = pending_future
     send_request = MagicMock(
-        name="send_request", side_effect=[registration_future, raw_store_future]
+        name="send_request", side_effect=[registration_future, *raw_store_futures]
     )
     event_backend = MagicMock(name="event_backend")
-    event_backend.export_event.return_value = b"event"
+    event_backend.export_event.side_effect = [b"event-1", b"event-2"]
+    events = [MagicMock(name="event_1"), MagicMock(name="event_2")]
 
     with (
         patch.object(
@@ -609,21 +619,27 @@ def test_lmcache_driven_preemption_waits_for_remote_store_futures() -> None:
             mq_timeout=2.5,
             send_request=send_request,
         )
-        context.submit_store(
-            _request_id="request",
-            key="key",
-            instance_id=1,
-            kv_caches={},
-            block_ids=[[0]],
-            event=MagicMock(name="event"),
-            _blocks_in_chunk=1,
-        )
+        for index, event in enumerate(events):
+            context.submit_store(
+                _request_id="request",
+                key=f"key-{index}",
+                instance_id=1,
+                kv_caches={},
+                block_ids=[[index]],
+                event=event,
+                _blocks_in_chunk=1,
+            )
+
+    pending[0].retain_reference.assert_called_once_with(events[0])
+    pending[1].retain_reference.assert_called_once_with(events[1])
 
     context.flush_inflight_stores()
 
-    pending.result.assert_called_once_with(timeout=2.5)
+    for pending_future in pending:
+        pending_future.result.assert_called_once_with(timeout=2.5)
     context.flush_inflight_stores()
-    pending.result.assert_called_once()
+    for pending_future in pending:
+        pending_future.result.assert_called_once()
 
 
 def test_lmcache_driven_preemption_without_stores_is_noop() -> None:

tests/v1/multiprocess/test_ipc_memory_reclaim.py

@@ -17,6 +17,7 @@
 # Standard
 # Standard Library
 from types import SimpleNamespace
+from typing import cast
 from unittest.mock import MagicMock
 import threading
 import time
@@ -153,7 +154,7 @@ def block_resolve(*_args: object, **_kwargs: object) -> list[list[object]]:
             raise TimeoutError("active STORE was not released")
         raise RuntimeError("stop after lifetime check")
 
-    module.context.resolve_obj_keys.side_effect = block_resolve
+    cast(MagicMock, module.context.resolve_obj_keys).side_effect = block_resolve
 
     def run_store() -> None:
         try: