Disk-backed primary tier, and correct the concurrency claim

UPSTREAM_ISSUE.md

@@ -267,98 +267,8 @@ larger sessions, re-send it):
 
 ---
 
-# Third defect: a secondary tier silently corrupts KV when the TP group spans nodes
 
-This one produces **wrong output with no error**, so it matters more than the
-other two.
-
-`FileSystemTierManager` is constructed with a memoryview of the *primary* CPU
-tier (`tiering/spec.py`: `primary_kv_view = primary_tier.get_kv_memoryview()`),
-and both directions operate on it alone:
-
-```python
-def submit_store(self, job_metadata):        # CPU -> disk
-    task = functools.partial(batch_store_block, paths, self._primary_kv_view, ...)
-def submit_load(self, job_metadata):         # disk -> CPU
-    batch_load_block(paths, self._primary_kv_view, ...)
-```
-
-That CPU tier is a per-node `/dev/shm` mmap
-(`shared_offload_region.py`: `Created mmap file /dev/shm/vllm_offload_<engine_id>.mmap`),
-and the secondary tier is instantiated **scheduler-side only** — in EngineCore
-and the API server, never in the workers.
-
-So when a tensor-parallel group spans hosts (`--nnodes 2`, one `vllm serve`
-process per node):
-
-* rank 0's shard is cascaded CPU -> disk and promoted disk -> CPU normally;
-* rank 1's CPU region is on the other host and is never touched by any FS tier;
-* the scheduler nonetheless tracks one logical block set and issues the load to
-  every worker (`pending_count = num_workers`);
-* each worker copies CPU -> GPU from **its own** region, so rank 1 copies
-  whatever stale bytes occupy those slots.
-
-The request is then marked as having a valid cached prefix while half the heads
-hold garbage. Measured on our pair: the restore reported `cached_tokens=13312`
-and completed in 5.63 s instead of 16.26 s, and the model emitted 512
-consecutive tokens of neither `content` nor `reasoning_content` — the same
-prompt answers correctly on a cold prefill and on a GPU prefix-cache hit.
-Confirmation that only one rank participates:
-
-```
-head   : 15 GB under kvoffload/, directories ..._r0/...
-worker : kvoffload/ empty, 0 files
-```
-
-`replicated_layout` in `offloading/config.py` already refuses to engage when
-`parallel_config.nnodes_within_dp != 1`, with the comment "Shared /dev/shm mmap
-layout is single-node mp only" — so the single-node assumption is understood in
-one place but not enforced where it silently changes results.
-
-**Suggested fix:** refuse to configure a secondary tier when the KV-parallel
-group spans nodes (fail closed at startup), or instantiate the tier per rank so
-each cascades its own CPU region. A CPU-only tier is unaffected, because the
-scheduler drives every rank's region symmetrically.
-
-### We implemented the per-rank option, out-of-tree
-
-`peer_kv_agent.py` + `PeerMirroredFileSystemTierManager` (in
-`vllm_bounded_fs_tier_peer.py`) make the disk tier correct on a multi-host TP
-group without patching vLLM, using the supported out-of-tree
-`secondary_tiers[].module_path` extension point.
-
-The head-side tier mirrors every cascade and promotion to an agent running in
-each remote rank's container. The agent performs byte-identical I/O against
-*that* rank's own `/dev/shm` region and its own `_r<rank>` files —
-`FileMapper.base_path` excludes rank ("rank lives outside the hash"), so the
-shards are siblings in one namespace. Crucially the head does not release a
-job's `JobResult` until every peer has acknowledged, which keeps disk -> CPU
-ordered before CPU -> GPU on all ranks. A peer failure fails the job, so
-`mark_miss` turns it into a recompute — a wasted prefill instead of a silently
-half-restored prefix.
-
-Measured on the same probe, restores served from disk (8 FS-tier block hits per
-session, both ranks holding 220 `.bin` files / 11 GB):
-
-| session | cold TTFT | restored TTFT | speedup | cached | needle |
-|---|---|---|---|---|---|
-| 1 | 15.45 s | 2.59 s | **6.0x** | 13,312 | OK |
-| 2 | 15.57 s | 2.67 s | **5.8x** | 13,312 | OK |
-
-No cross-session contamination. Before the per-rank cascade the identical
-configuration returned 512 consecutive empty tokens.
-
-Known limitations of our workaround, all easy to lift and none affecting
-correctness: the peer's disk is unbounded (only the head tier enforces
-`max_bytes`, and a peer-side miss degrades to a recompute); each rank stores
-whole block rows, so the other rank's inert half doubles disk usage; and the CPU
-primary tier must be large enough to accept a promotion — a 4 GiB tier against a
-1 GiB GPU pool produced `kv_offload_tiering_promotion_allocation_failures_total`
-and every lookup missed, while 8 GiB worked.
-
----
-
-# Fourth defect: the CPU offload region is never unlinked, and it breaks the next boot
+# Third defect: the CPU offload region is never unlinked, and it breaks the next boot
 
 `SharedOffloadRegion` creates `/dev/shm/vllm_offload_<engine_id>.mmap` and nothing
 removes it when the engine stops or crashes. Every boot mints a fresh
@@ -383,7 +293,7 @@ it by clearing `/dev/shm/vllm_offload_*.mmap` on both hosts before launch.
 
 ---
 
-# Fifth issue: one uniform region row size makes fine-grained groups pathological
+# Fourth issue: one uniform region row size makes fine-grained groups pathological
 
 Not a bug so much as a scaling cliff worth documenting. The offload region uses a
 single row size for every KV group (derived from an average over all blocks), but

patch_offload_mmap_region.py

@@ -0,0 +1,146 @@
+#!/usr/bin/env python3
+"""Back the CPU staging region with a sparse file on local NVMe.
+
+Ported from MiaAI PR#58 (overlay/patch_kv_offload_groups.py, its `sor_edits`
+and the cudaHostRegister skip), re-expressed in this kit's overlay style. Their
+group-exclusion half is NOT taken: we already have
+overlay/patch_offload_group_filter.py, which does the same job and additionally
+carries the region-copies fix (rows sized for world_size while every worker uses
+slot 0) that halved the row from 51.81 to 25.91 MiB. PR#58 does not have that.
+
+WHY THIS REPLACES THE STREAMED RESTORE. vLLM admits a restore only when every
+group's chunks are resident in the primary tier TOGETHER -- PR#58 states the
+sizing rule directly: cpu_bytes_to_use >= row_bytes * (restore_tokens/3584 + 8),
+and "undersized, stores still succeed and metrics still climb while no restore
+ever serves", which is exactly the signature measured here (store count rising,
+load count frozen). Our step-2b streamed restore tried to dodge that invariant
+by reporting HIT and making chunks resident batch by batch; it stalls at N>=6.
+Backing the region with NVMe instead makes the tier large enough to satisfy the
+invariant, so the restore window stops being bounded by RAM without fighting the
+scheduler.
+
+SIZING. Do not take a row size on trust -- read it off your own boot. vLLM sets
+num_blocks = cpu_bytes_to_use // aligned_kv_bytes_per_chunk and then creates a
+region of exactly num_blocks * aligned_kv_bytes_per_chunk, so the "Created mmap
+file ... (N GB)" line divided by num_blocks gives you the real row. On this
+stack it is 54,329,344 B (51.81 MiB), confirmed twice: 32 GiB -> 632 rows and
+260 GB -> 4785 rows, both exact.
+
+The rule that actually matters is not a byte count, it is a comparison: THE TIER
+MUST BE LARGER THAN THE GPU KV POOL. Both are LRU over the same stream of
+blocks, so a tier smaller than the pool holds a strict subset of what the pool
+already has -- the GPU answers first and the tier can never be the thing that
+serves. Sizing it below the pool is not a smaller cache, it is no cache. The
+file is sparse, so over-sizing costs nothing until written.
+
+Gated on GLM53_OFFLOAD_MMAP_DIR being set; unset leaves vLLM's /dev/shm
+behaviour untouched. Idempotent, sentinel-guarded, fails closed on drift.
+"""
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+VLLM = Path("/usr/local/lib/python3.12/dist-packages/vllm")
+SOR = VLLM / "v1/kv_offload/cpu/shared_offload_region.py"
+GWK = VLLM / "v1/kv_offload/cpu/gpu_worker.py"
+MARK = "[glm53-offload-mmap]"
+
+# 1. put the region file on NVMe instead of /dev/shm
+PATH_OLD = '''        self.mmap_path = f"/dev/shm/vllm_offload_{engine_id}.mmap"'''
+PATH_NEW = '''        # LOCAL [glm53-offload-mmap]: a sparse file on this node's own NVMe, so
+        # the staging tier's capacity is disk rather than RAM. Each rank keeps
+        # its own file, which is also why no cross-node mirroring is needed.
+        _mmap_dir = os.environ.get("GLM53_OFFLOAD_MMAP_DIR")
+        self._glm53_disk_backed = bool(_mmap_dir)
+        if _mmap_dir:
+            os.makedirs(_mmap_dir, exist_ok=True)
+        else:
+            _mmap_dir = "/dev/shm"
+        self.mmap_path = f"{_mmap_dir}/vllm_offload_{engine_id}.mmap"'''
+PATH_SENTINEL = "[glm53-offload-mmap]: a sparse file"
+
+# 2. the /dev/shm free-space check is meaningless for a sparse file on NVMe
+SHM_OLD = '''                check_shm_free_space(self.total_size_bytes)
+                os.ftruncate(self.fd, self.total_size_bytes)'''
+SHM_NEW = '''                if not self._glm53_disk_backed:  # LOCAL [glm53-offload-mmap]
+                    check_shm_free_space(self.total_size_bytes)
+                os.ftruncate(self.fd, self.total_size_bytes)'''
+SHM_SENTINEL = "if not self._glm53_disk_backed:"
+
+# 3. never pre-fault: writing every page would materialise the whole file
+POP_OLD = '''        populate_write_fn = _get_populate_write_fn(self.mmap_obj)
+
+        if rank is not None:'''
+POP_NEW = '''        # LOCAL [glm53-offload-mmap]: pre-faulting a sparse NVMe file would
+        # materialise the entire region on disk at boot.
+        if self._glm53_disk_backed:
+            populate_write_fn = lambda *_a, **_k: None
+            _glm53_skip_populate = True
+        else:
+            populate_write_fn = _get_populate_write_fn(self.mmap_obj)
+            _glm53_skip_populate = False
+
+        if _glm53_skip_populate:
+            pass
+        elif rank is not None:'''
+POP_SENTINEL = "_glm53_skip_populate"
+
+SOR_EDITS = [
+    ("sor:mmap-path", PATH_OLD, PATH_NEW, PATH_SENTINEL),
+    ("sor:skip-shm-check", SHM_OLD, SHM_NEW, SHM_SENTINEL),
+    ("sor:skip-populate", POP_OLD, POP_NEW, POP_SENTINEL),
+]
+
+# 4. cudaHostRegister on a disk-backed mapping is wrong and slow
+GWK_OLD = '''    if not current_platform.is_cuda_alike():'''
+GWK_NEW = '''    import os as _os  # LOCAL [glm53-offload-mmap]
+    if _os.environ.get("GLM53_OFFLOAD_MMAP_DIR"):
+        logger.info("%s disk-backed staging region: skipping cudaHostRegister",
+                    "[glm53-offload-mmap]")
+        return
+    if not current_platform.is_cuda_alike():'''
+GWK_SENTINEL = "skipping cudaHostRegister"
+
+GWK_EDITS = [("gwk:skip-hostregister", GWK_OLD, GWK_NEW, GWK_SENTINEL)]
+
+
+def apply_edits(path: Path, edits) -> str:
+    if not path.is_file():
+        return f"MISSING {path} - not patched"
+    src = path.read_text()
+    applied, skipped = [], []
+    for label, old, new, sentinel in edits:
+        if sentinel in src:
+            skipped.append(label)
+            continue
+        n = src.count(old)
+        if n != 1:
+            raise SystemExit(
+                f"{MARK} FAIL {label}: anchor matched {n} times (expected 1). "
+                f"Upstream drifted; refusing to patch.")
+        src = src.replace(old, new)
+        applied.append(label)
+    if applied:
+        compile(src, str(path), "exec")
+        path.write_text(src)
+    parts = []
+    if applied:
+        parts.append("applied " + ",".join(applied))
+    if skipped:
+        parts.append("already " + ",".join(skipped))
+    return "; ".join(parts) or "nothing to do"
+
+
+def main() -> int:
+    if not os.environ.get("GLM53_OFFLOAD_MMAP_DIR"):
+        print(f"{MARK} disabled (GLM53_OFFLOAD_MMAP_DIR unset); vLLM pristine")
+        return 0
+    print(f"{MARK} {apply_edits(SOR, SOR_EDITS)}")
+    print(f"{MARK} {apply_edits(GWK, GWK_EDITS)}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())