Optional resident mode for the quantized tables (VLLM_PLE_QUANT_RESIDENT) (#2)

- Optional resident mode for the quantized tables (VLLM_PLE_QUANT_RESIDENT) (5e57b1daaed1c6161fbbdbbe226cd6b680bfb777)
- Preflight: measure available memory and the actual shard bytes (4bb76906884e27b8e5030d7d607db255684ce239)


Co-authored-by: mgabor <mgabor3141@users.noreply.huggingface.co>

README.md

@@ -54,6 +54,35 @@ Drop `-e VLLM_GDN_DECODE_KERNEL=triton` when serving the plain NVFP4 build. No c
 memory cap needed: unlike the BF16 disk path, the quantized tables fit the page cache next to
 checkpoint streaming. Two flags were added to the block on 2026-09-06 after a field report, and both failures reproduce on our box: the image's default `--max-num-seqs` (1024) exceeds the Mamba state blocks left at 0.92 utilization (598) and CUDA graph capture aborts, and at the native 262,144 context with MTP the KV cache needs 7.57 GiB where 5.95 GiB is left after the draft head. 32,768 / 36 is what the table below was measured with; the report ran 135,168 / 64 with MTP.
 
+### Optional: keep the table in anonymous RAM
+
+`-e VLLM_PLE_QUANT_RESIDENT=1` copies the shards into anonymous memory at load instead of
+serving them mmapped. Unset, nothing changes. Contributed by
+[mgabor3141](https://huggingface.co/primitive-ai/Qwen3.8-Flash-Next-PLE-quant/discussions/2).
+
+This buys cold-start, not throughput. Measured on one RTX PRO 6000, INT4 table, mixed
+checkpoint, page cache dropped before each arm:
+
+| | first request | warm asymptote | requests to within 5% |
+|---|---|---|---|
+| mmapped | 19.0 tok/s | 93.0 | ~150 |
+| resident | 74.3 | 96.8 | ~3 |
+| mmapped, MTP nspec 3 | 12.6 | 125.4 | ~125 |
+| resident, MTP nspec 3 | 86.5 | 128.2 | ~25 |
+
+Warm, the two are 2–4% apart and the MTP pair is inside the run-to-run spread. Cold they are
+not close: loading the 172 GB checkpoint evicts the sidecar from page cache, so the mmapped
+table serves its first request at 19 tok/s and takes roughly 150 requests to converge.
+Speculation makes that worse, not better. `FileHugePages` is 0 on this kernel while the
+resident copy shows 32 GiB of `anon_thp`, which is where the warm 2–4% comes from.
+
+Worth the RAM if the server restarts often or takes a burst right after boot; not otherwise.
+KV pool is unchanged either way (426,639 tokens without MTP, 134,192 with, identical in both
+arms), and the copy costs 4–11 s of boot. The table needs 48.9 / 29.8 / 26.8 GiB of anonymous
+memory for the fp8 / int4 / nvfp4 sidecars, on top of the roughly 8 GiB the worker already
+holds. A preflight refuses with a clear message rather than OOM-ing part-way in; it reads
+`MemAvailable` and, in a container, the cgroup limit less its reclaimable page cache.
+
 Hosts with less than about 100 GB of RAM plus swap: the `ple_layer_quant.py` in this repo since 2026-09-06 builds the 95 GB table parameter on the meta device, so the offload worker no longer requests that allocation and `vm.overcommit_memory` can stay at its default. Before that revision the worker constructed the BF16 parameter first and the kernel's heuristic refused it (`DefaultCPUAllocator: can't allocate memory: you tried to allocate 102400491520 bytes`), which a 64 GB host reported and we reproduced under an emulated 67 GiB commit limit; the revised file boots under the same limit and scores 77.0 (three runs: 76.0, 76.0, 79.0, against 78.2 for the same table with the previous overlay) on the tool-calling suite. The reporter then confirmed it on the real machine (64 GB, CommitLimit 40 GiB, kernel-default overcommit): 341 s to healthy, KV pool unchanged at 216,778 tokens, host RAM about 10 GB, with the previous overlay still failing in the same session ([log](https://github.com/MarcoPizeta/flash-next-rtxpro6000-bench/blob/main/logs/test-revised-ple-overlay-no-overcommit-20260906.log)). Re-download the overlay if yours predates 2026-09-06. The BF16 disk overlay on the model cards still needs `sysctl vm.overcommit_memory=1` on such hosts.
 
 ## Measured

worker_image_quant.py

@@ -340,8 +340,112 @@ class _PleQuantTable:
         import json
         import os
 
+        import ctypes
+
         from safetensors import safe_open
 
+        resident = os.environ.get("VLLM_PLE_QUANT_RESIDENT", "0") == "1"
+
+        def _cgroup_headroom() -> int | None:
+            """Bytes this cgroup may still allocate, or None when unlimited.
+
+            `memory.current` counts the page cache too, and on this path most of
+            it is the checkpoint we just streamed, which the kernel will reclaim
+            rather than OOM over. Subtract the file pages, the same way
+            MemAvailable does at the host level, or a container that has read a
+            large checkpoint refuses a copy it could serve.
+            """
+            for limit_p, usage_p, stat_p, keys in (
+                ("/sys/fs/cgroup/memory.max",
+                 "/sys/fs/cgroup/memory.current",
+                 "/sys/fs/cgroup/memory.stat",
+                 ("inactive_file", "active_file")),
+                ("/sys/fs/cgroup/memory/memory.limit_in_bytes",
+                 "/sys/fs/cgroup/memory/memory.usage_in_bytes",
+                 "/sys/fs/cgroup/memory/memory.stat",
+                 ("total_inactive_file", "total_active_file")),
+            ):
+                try:
+                    raw = open(limit_p).read().strip()
+                    if raw == "max":
+                        return None
+                    limit = int(raw)
+                    # cgroup v1 spells "unlimited" as a sentinel near 2**63
+                    if limit >= (1 << 62):
+                        return None
+                    usage = int(open(usage_p).read().strip())
+                except (OSError, ValueError):
+                    continue
+                reclaimable = 0
+                try:
+                    for line in open(stat_p):
+                        field, _, value = line.partition(" ")
+                        if field in keys:
+                            reclaimable += int(value)
+                except (OSError, ValueError):
+                    pass
+                return max(limit - max(usage - reclaimable, 0), 0)
+            return None
+
+        def _avail_bytes() -> int | None:
+            """Memory the copy can actually claim, or None if it cannot be read.
+
+            SC_AVPHYS_PAGES is MemFree: it excludes the page cache, which on this
+            path is largely the sidecar and the checkpoint, and it reports the
+            host's pages even inside a container with `--memory`. Both errors
+            bite: a host with 173 GiB of MemAvailable and 1.9 GiB free refuses a
+            copy that would succeed, and a capped container passes the check and
+            is killed part-way through it.
+
+            None means "could not tell", and the caller then proceeds. This
+            preflight exists to fail fast and legibly, not to gate correctness;
+            a platform whose memory it cannot read should still be able to boot.
+            """
+            avail = None
+            try:
+                for line in open("/proc/meminfo"):
+                    if line.startswith("MemAvailable:"):
+                        avail = int(line.split()[1]) * 1024
+                        break
+            except (OSError, ValueError, IndexError):
+                pass
+            if avail is None:
+                try:
+                    avail = os.sysconf("SC_AVPHYS_PAGES") * os.sysconf("SC_PAGE_SIZE")
+                except (ValueError, OSError):
+                    return None
+            headroom = _cgroup_headroom()
+            return avail if headroom is None else min(avail, headroom)
+
+        def _drop_cache(path: str) -> None:
+            """Release the page cache for a shard we have already copied.
+
+            Without this the cgroup briefly accounts both the anonymous copy and
+            the file pages it came from. The kernel would reclaim the latter under
+            pressure anyway; doing it eagerly keeps the steady-state footprint
+            honest and predictable.
+            """
+            POSIX_FADV_DONTNEED = 4
+            try:
+                fd = os.open(path, os.O_RDONLY)
+            except OSError:
+                return
+            try:
+                libc = ctypes.CDLL("libc.so.6", use_errno=True)
+                # posix_fadvise returns the errno instead of setting it, so a
+                # bare try/except here would never fire.
+                rc = libc.posix_fadvise(
+                    ctypes.c_int(fd), ctypes.c_long(0), ctypes.c_long(0),
+                    ctypes.c_int(POSIX_FADV_DONTNEED),
+                )
+                if rc != 0:
+                    logger.debug("posix_fadvise(DONTNEED) on %s returned %d",
+                                 path, rc)
+            except OSError:
+                pass
+            finally:
+                os.close(fd)
+
         meta = json.load(open(os.path.join(quant_dir, "META.json")))
         self.layout = meta["layout"]
         assert meta["rows"] == total_rows and meta["width"] == width, (
@@ -358,19 +462,60 @@ class _PleQuantTable:
         else:
             key = "weight_i4"
         self._q, self._s, self._s2 = [], [], []
+        if resident:
+            # Preflight: the table size is known before any of it is read, so an
+            # undersized host should fail now rather than part-way through.
+            # The shard files are the copy, so measure them. rows*width is exact
+            # for no layout: it omits the scales (fp8 is 1.19 GiB short) and
+            # ignores 4-bit packing (int4 1.60x over, nvfp4 1.78x).
+            need = sum(
+                os.path.getsize(os.path.join(quant_dir, f"shard_{n}.safetensors"))
+                for n in range(n_shards)
+            )
+            avail = _avail_bytes()
+            if avail is not None and need > avail:
+                raise MemoryError(
+                    f"PLE resident mode needs about {need / 2**30:.1f} GiB for the "
+                    f"table but only {avail / 2**30:.1f} GiB is available to this "
+                    "process; unset VLLM_PLE_QUANT_RESIDENT to serve it mmapped"
+                )
+            logger.info(
+                "PLE quant table: resident mode, copying about %.1f GiB into "
+                "anonymous RAM (%s available)",
+                need / 2**30,
+                "headroom unknown" if avail is None else f"{avail / 2**30:.1f} GiB",
+            )
+
         for n in range(n_shards):
-            f = safe_open(os.path.join(quant_dir, f"shard_{n}.safetensors"),
-                          framework="pt")
-            self._q.append(f.get_tensor(key))
-            self._s.append(f.get_tensor("weight_scale"))
+            shard_path = os.path.join(quant_dir, f"shard_{n}.safetensors")
+            f = safe_open(shard_path, framework="pt")
+            q = f.get_tensor(key)
+            s = f.get_tensor("weight_scale")
+            if resident:
+                # clone() copies into anonymous, THP-eligible memory. Do it per
+                # shard so the transient cost is one shard, not one table.
+                q = q.clone()
+                s = s.clone()
+            self._q.append(q)
+            self._s.append(s)
             self._s2.append(
                 f.get_tensor("weight_scale_2").item()
                 if "weight_scale_2" in f.keys() else 1.0
             )
         self.width = width
         self._lut = None
-        logger.info("PLE quant table: %s, %d shards mmapped from %s",
-                    self.layout, n_shards, quant_dir)
+        if resident:
+            del f
+            for n in range(n_shards):
+                _drop_cache(os.path.join(quant_dir, f"shard_{n}.safetensors"))
+            logger.info(
+                "PLE quant table: %s, %d shards resident in anonymous RAM from %s "
+                "(no disk reads on the gather path)",
+                self.layout, n_shards, quant_dir,
+            )
+        else:
+            logger.info("PLE quant table: %s, %d shards mmapped from %s",
+                        self.layout, n_shards, quant_dir)
 
     def gather_into(self, ids: torch.Tensor, out: torch.Tensor) -> None:
         ids = ids.long()