launcher/dcp2/diffs/07-models_deepseek_v4_1_common_engram.py.diff

--- baked-ours/models/deepseek_v4_1/common/engram.py	2026-09-10 20:13:20.000000000 +0200
+++ baked/models/deepseek_v4_1/common/engram.py	2026-09-11 10:13:47.000000000 +0200
@@ -634,15 +634,105 @@
 _KAI_THREADS = int(_kai_os.environ.get("DSV41_ENGRAM_DISK_THREADS", "32"))
 _KAI_CHUNK = int(_kai_os.environ.get("DSV41_ENGRAM_DISK_CHUNK", "16"))
 _KAI_POOL: _KaiPool | None = None
+# Stage 6: gather rows through a read-only memory map (one NumPy fancy-index copy
+# per chunk) instead of one preadv syscall per row.
+_KAI_MMAP = _kai_os.environ.get("DSV41_ENGRAM_MMAP", "0") == "1"
+# The gather releases the GIL, so more threads = more page faults in flight
+# (3 of 4 ranks fault rows in over NFS; each cold row is one round trip).
+_KAI_MMAP_THREADS = int(_kai_os.environ.get("DSV41_ENGRAM_MMAP_THREADS", "96"))
+# Stage 8: hash-ahead prefetch of the next prefill chunk's rows into the page
+# cache while the current chunk computes (the prompt is known up front).
+_KAI_PREFETCH = _kai_os.environ.get("DSV41_ENGRAM_PREFETCH", "0") == "1"
+_KAI_PREFETCH_MODE = _kai_os.environ.get("DSV41_ENGRAM_PREFETCH_MODE", "fadvise")  # fadvise | pread
+_KAI_PREFETCH_THREADS = int(_kai_os.environ.get("DSV41_ENGRAM_PREFETCH_THREADS", "32"))
+_KAI_PREFETCH_TASK_ROWS = 512
+# Gathers smaller than this many unique rows (per table) keep the preadv path
+# (0 = memmap for every size). Lets decode steps (72 rows) and prefill chunks
+# (tens of thousands) use different paths.
+_KAI_MMAP_MIN_ROWS = int(_kai_os.environ.get("DSV41_ENGRAM_MMAP_MIN_ROWS", "0"))
+# v3 (window 8: 72 decode rows cost 5-6 ms = one NVMe latency per row, serialized by
+# the per-task faults): gathers up to this many rows skip the pool -- first an async
+# readahead of every row page (posix_fadvise WILLNEED, ~2 us each, all reads in
+# flight together), then ONE fancy-index gather per table on the calling thread.
+_KAI_MMAP_INLINE_ROWS = int(_kai_os.environ.get("DSV41_ENGRAM_MMAP_INLINE_ROWS", "2048"))
 
 
 def _kai_pool() -> _KaiPool:
     global _KAI_POOL
     if _KAI_POOL is None:
-        _KAI_POOL = _KaiPool(max_workers=_KAI_THREADS, thread_name_prefix="engram-disk")
+        _KAI_POOL = _KaiPool(
+            max_workers=_KAI_MMAP_THREADS if _KAI_MMAP else _KAI_THREADS,
+            thread_name_prefix="engram-disk",
+        )
     return _KAI_POOL
 
 
+def _kai_prefetch_rows(table, rows: list) -> None:
+    """Warm rank-local rows `rows` of `table` in the page cache: weight rows by
+    exact byte range, scale rows by 4 KiB page (32 rows share one)."""
+    if not rows:
+        return
+    dim, sb = table.dim, table.sb
+    if _KAI_PREFETCH_MODE == "pread":
+        buf = bytearray(dim)
+        for r in rows:
+            try:
+                _kai_os.preadv(table.w_fd, [buf], table.w_off + r * dim)
+            except OSError:
+                pass
+        sbuf = bytearray(4096)
+        for p in {(table.s_off + r * sb) >> 12 for r in rows}:
+            try:
+                _kai_os.preadv(table.s_fd, [sbuf], p << 12)
+            except OSError:
+                pass
+        return
+    adv = _kai_os.POSIX_FADV_WILLNEED
+    for r in rows:
+        try:
+            _kai_os.posix_fadvise(table.w_fd, table.w_off + r * dim, dim, adv)
+        except OSError:
+            pass
+    for p in {(table.s_off + r * sb) >> 12 for r in rows}:
+        try:
+            _kai_os.posix_fadvise(table.s_fd, p << 12, 4096, adv)
+        except OSError:
+            pass
+
+
+def _kai_map_file(path: str):
+    """Read-only shared mapping of a whole shard with readahead DISABLED
+    (MADV_RANDOM, the mmap twin of the preadv path's POSIX_FADV_RANDOM): a
+    fault brings exactly the 4 KiB page of the row. Window 7 (2026-09-11)
+    showed why: without it every cold-row fault pulled an NFS readahead window
+    (rsize 1 MiB) — 9K prefill 14.7 s vs 8.1 s, code decode 53 vs 61 tok/s."""
+    import mmap as _mmap
+    import numpy as _np
+
+    fd = _kai_os.open(path, _kai_os.O_RDONLY)
+    try:
+        mm = _mmap.mmap(fd, 0, access=_mmap.ACCESS_READ)
+    finally:
+        _kai_os.close(fd)
+    try:
+        mm.madvise(_mmap.MADV_RANDOM)
+    except (AttributeError, OSError):  # pragma: no cover - old kernels
+        pass
+    return mm, _np.frombuffer(mm, dtype=_np.uint8)
+
+
+def _kai_mmap_rows(view, rel, lo: int, hi: int, row_bytes: int, buf) -> None:
+    """Copy rows rel[lo:hi] of the [num_rows, row_bytes] memmap view into buf
+    (a flat uint8 memoryview). One fancy-index gather: C memcpy for resident
+    pages, page faults for cold ones; the GIL is released inside NumPy."""
+    import numpy as _np
+
+    dst = _np.frombuffer(buf, dtype=_np.uint8)[lo * row_bytes : hi * row_bytes].reshape(
+        hi - lo, row_bytes
+    )
+    _np.take(view, _np.asarray(rel[lo:hi], dtype=_np.int64), axis=0, out=dst)
+
+
 def _kai_pread_rows(
     fd: int, base: int, rel: list, lo: int, hi: int, row_bytes: int, buf
 ) -> None:
@@ -665,18 +755,41 @@
     total = sum(len(job[2]) for job in jobs)
     if total == 0:
         return
+
+    def _run(job, lo, hi):
+        fd, base, rel, row_bytes, buf = job[:5]
+        view = job[5] if len(job) > 5 else None
+        if view is not None:
+            _kai_mmap_rows(view, rel, lo, hi, row_bytes, buf)
+        else:
+            _kai_pread_rows(fd, base, rel, lo, hi, row_bytes, buf)
+
     if total == 1:
-        for fd, base, rel, row_bytes, buf in jobs:
-            _kai_pread_rows(fd, base, rel, 0, len(rel), row_bytes, buf)
+        for job in jobs:
+            _run(job, 0, len(job[2]))
         return
-    chunk = max(1, min(_KAI_CHUNK, -(-total // _KAI_THREADS)))
+    # Memmap gathers: spread the rows over every pool thread (cold rows fault
+    # in sequentially within a task, so more tasks = more reads in flight).
+    if any(len(job) > 5 and job[5] is not None for job in jobs):
+        if total <= _KAI_MMAP_INLINE_ROWS and all(len(job) > 5 and job[5] is not None for job in jobs):
+            adv = _kai_os.POSIX_FADV_WILLNEED
+            for fd, base, rel, row_bytes, buf, view in jobs:
+                for r in rel:
+                    try:
+                        _kai_os.posix_fadvise(fd, base + r * row_bytes, row_bytes, adv)
+                    except OSError:
+                        pass
+            for fd, base, rel, row_bytes, buf, view in jobs:
+                _kai_mmap_rows(view, rel, 0, len(rel), row_bytes, buf)
+            return
+        chunk = max(1, -(-total // _KAI_MMAP_THREADS))
+    else:
+        chunk = max(1, min(_KAI_CHUNK, -(-total // _KAI_THREADS)))
     pool = _kai_pool()
     futs = [
-        pool.submit(
-            _kai_pread_rows, fd, base, rel, lo, min(lo + chunk, len(rel)), row_bytes, buf
-        )
-        for fd, base, rel, row_bytes, buf in jobs
-        for lo in range(0, len(rel), chunk)
+        pool.submit(_run, job, lo, min(lo + chunk, len(job[2])))
+        for job in jobs
+        for lo in range(0, len(job[2]), chunk)
     ]
     for fut in futs:
         fut.result()
@@ -734,6 +847,8 @@
         sname = f"layers.{layer_id}.engram.embed.scale"
         self.w_fd, self.w_off, self.w_shape = self._open(model_dir, weight_map[wname], wname)
         self.s_fd, self.s_off, self.s_shape = self._open(model_dir, weight_map[sname], sname)
+        self._w_path = _kai_os.path.join(model_dir, weight_map[wname])
+        self._s_path = _kai_os.path.join(model_dir, weight_map[sname])
         self.dim = dim
         self.sb = dim // block_size
         assert self.w_shape[1] == dim, (self.w_shape, dim)
@@ -750,6 +865,21 @@
         self.s_off += row_start * self.sb
         self.row_start = row_start
         self.num_rows = num_rows
+        # Stage 6: [num_rows, row_bytes] uint8 views of this rank's rows through a
+        # read-only memory map of each shard (opened once; pages come and go
+        # with the OS page cache, nothing is pinned).
+        self.w_view = self.s_view = None
+        if _KAI_MMAP:
+            import numpy as _np
+
+            self._w_mm, w_buf = _kai_map_file(self._w_path)
+            self._s_mm, s_buf = _kai_map_file(self._s_path)
+            self.w_view = _np.ndarray(
+                (num_rows, dim), dtype=_np.uint8, buffer=w_buf, offset=self.w_off
+            )
+            self.s_view = _np.ndarray(
+                (num_rows, self.sb), dtype=_np.uint8, buffer=s_buf, offset=self.s_off
+            )
         self.threads = _KAI_THREADS
         self.chunk = _KAI_CHUNK
         self.pool = _kai_pool()  # shared by all tables (was one pool per table)
@@ -780,6 +910,11 @@
 
     def read_jobs(self, rel_l: list, w: torch.Tensor, s: torch.Tensor) -> list:
         """Read jobs for rank-local rows `rel_l` into uint8 w [R, dim], s [R, sb]."""
+        if self.w_view is not None and len(rel_l) >= _KAI_MMAP_MIN_ROWS:
+            return [
+                (self.w_fd, self.w_off, rel_l, self.dim, memoryview(w.numpy()).cast("B"), self.w_view),
+                (self.s_fd, self.s_off, rel_l, self.sb, memoryview(s.numpy()).cast("B"), self.s_view),
+            ]
         return [
             (self.w_fd, self.w_off, rel_l, self.dim, memoryview(w.numpy()).cast("B")),
             (self.s_fd, self.s_off, rel_l, self.sb, memoryview(s.numpy()).cast("B")),
@@ -1251,6 +1386,18 @@
         hash_ids: torch.Tensor,
         token_mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
+        # Stage 7: phase "engram_fwd" of the prefill budget (inert unless enabled).
+        from vllm.model_executor.kernels.attention.dsa.candidate_blocks import _PrefillTimer
+
+        with _PrefillTimer("engram_fwd", int(hash_ids.shape[0])):
+            return self._forward_impl(hidden_states, hash_ids, token_mask)
+
+    def _forward_impl(
+        self,
+        hidden_states: torch.Tensor,
+        hash_ids: torch.Tensor,
+        token_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
         """hidden_states: [T, hc_mult, dim]; hash_ids: [T, n_hash_cols] (all
         tokens, pre sequence-parallel shard); token_mask: [T], False shuts
         the gate so those positions pass through untouched."""
@@ -1352,6 +1499,27 @@
         ]
         self.hashes_ready = torch.cuda.Event()
         self.num_staged = 0
+        # DSV41_ENGRAM_TIMING=1: host wall-clock of stage() phases (diagnostic).
+        self._timing = _kai_os.environ.get("DSV41_ENGRAM_TIMING", "0") == "1"
+        # two buckets (decode: n <= 16 tokens, prefill: n > 16): sync, gather+dequant, tail, total, calls, rows, unique rows
+        self._t_acc = [[0.0, 0.0, 0.0, 0.0, 0, 0, 0] for _ in range(2)]
+        self._t_last_total_ms = 0.0  # stage 7 reads it after each timed prefill step
+        # Stage 8: hash-ahead prefetch state (see _hash_ahead).
+        self._prefetch = _KAI_PREFETCH
+        if self._prefetch:
+            self.hash_host_next = torch.empty_like(self.hash_host)
+            self.hashes_next_ready = torch.cuda.Event()
+            self.last_pos_host = torch.zeros(256, dtype=torch.int64, device="cpu", pin_memory=True)
+            self._prefetch_pool = _KaiPool(
+                max_workers=_KAI_PREFETCH_THREADS, thread_name_prefix="engram-prefetch"
+            )
+            self._prefetch_futs: list = []
+            self._prefetch_stats = [0, 0, 0.0]  # lookaheads, rows, host seconds
+            logger.info(
+                "Engram DISK prefetch: hash-ahead of the next prefill chunk (%s, %d threads)",
+                _KAI_PREFETCH_MODE, _KAI_PREFETCH_THREADS,
+            )
+
         logger.info(
             "Engram DISK rows staged before the forward (graph-safe): %d layers, "
             "heads [%d, %d) of %d, up to %d tokens/step, %d read threads",
@@ -1363,6 +1531,72 @@
             _KAI_THREADS,
         )
 
+    def _hash_ahead(self, input_batch, req_states, positions, query_start_loc, lookback_token_ids) -> None:
+        """Stage 8: for the batch's longest still-prefilling request, hash the
+        NEXT chunk of its (known) prompt with the same kernel and warm this
+        rank's rows in the page cache while the current step computes."""
+        if torch.cuda.is_current_stream_capturing():
+            return
+        import time as _kai_time
+
+        from .mm_preprocess import image_sentinel_mask
+
+        t0 = _kai_time.perf_counter()
+        n_reqs = int(input_batch.num_reqs)
+        if n_reqs <= 0:
+            return
+        idx_np = input_batch.idx_mapping_np[:n_reqs]
+        prompt_len = req_states.prompt_len.np[idx_np]
+        next_start = self.last_pos_host[:n_reqs].numpy() + 1  # valid after hashes_ready.synchronize()
+        remaining = prompt_len - next_start
+        i = int(remaining.argmax())
+        if remaining[i] <= 0:
+            return
+        req_idx = int(idx_np[i])
+        s = int(next_start[i])
+        L = int(min(remaining[i], self.max_tokens))
+        tok = req_states.all_token_ids.gpu
+        ids = tok[req_idx, s : s + L]
+        pos = torch.arange(s, s + L, dtype=positions.dtype, device=positions.device)
+        qsl = torch.tensor([0, L], dtype=query_start_loc.dtype, device=query_start_loc.device)
+        depth = int(lookback_token_ids.shape[1])
+        lb = torch.full((1, depth), -1, dtype=lookback_token_ids.dtype, device=lookback_token_ids.device)
+        k = min(depth, s)
+        if k > 0:
+            lb[0, :k] = tok[req_idx, s - k : s].flip(0).to(lb.dtype)
+        hashes = self.hash_state(
+            ids, pos, qsl, image_sentinel_mask(ids), lb, image_sentinel_mask(lb), None, None
+        )
+        host = self.hash_host_next[:L]
+        host.copy_(hashes[:, :, self.head_start : self.head_end], non_blocking=True)
+        self.hashes_next_ready.record()
+        self.hashes_next_ready.synchronize()
+        self._prefetch_futs = [f for f in self._prefetch_futs if not f.done()]
+        total_rows = 0
+        for engram in self.engrams:
+            local = host[:, engram.layer_hash_index, :].to(torch.int64)
+            rel, owned = engram.embed_tokens.disk_rel_owned(local)
+            rows = torch.unique(rel[owned]).tolist()
+            total_rows += len(rows)
+            table = engram.embed_tokens.disk
+            for lo in range(0, len(rows), _KAI_PREFETCH_TASK_ROWS):
+                self._prefetch_futs.append(
+                    self._prefetch_pool.submit(
+                        _kai_prefetch_rows, table, rows[lo : lo + _KAI_PREFETCH_TASK_ROWS]
+                    )
+                )
+        st = self._prefetch_stats
+        st[0] += 1
+        st[1] += total_rows
+        st[2] += _kai_time.perf_counter() - t0
+        if st[0] % 16 == 0 or L >= 4096:
+            logger.info(
+                "Engram prefetch: next chunk %d tokens from %d (prompt %d), %d unique rows queued, "
+                "host %.1f ms (%d lookaheads so far, avg %.0f rows, %.1f ms)",
+                L, s, int(prompt_len[i]), total_rows, 1e3 * (_kai_time.perf_counter() - t0),
+                st[0], st[1] / st[0], 1e3 * st[2] / st[0],
+            )
+
     @torch.inference_mode()
     def stage(
         self,
@@ -1371,6 +1605,8 @@
         query_start_loc: torch.Tensor,
         lookback_token_ids: torch.Tensor,
         num_tokens: int,
+        input_batch=None,
+        req_states=None,
     ) -> int:
         """Stage rows for the first `num_tokens` (unpadded) tokens of this step.
         Returns the number staged; 0 while the KV cache is unbound (profiling),
@@ -1380,6 +1616,12 @@
             return 0
         from .mm_preprocess import image_sentinel_mask
 
+        timing = self._timing
+        if timing:
+            import time as _kai_time
+
+            _t0 = _kai_time.perf_counter()
+
         ids = input_ids[:n]
         hashes = self.hash_state(
             ids,
@@ -1393,19 +1635,90 @@
         )
         host = self.hash_host[:n]
         host.copy_(hashes[:, :, self.head_start : self.head_end], non_blocking=True)
+        prefetch = self._prefetch and input_batch is not None and req_states is not None
+        if prefetch:
+            # last position of every request in this step -> next chunk start
+            n_reqs = int(input_batch.num_reqs)
+            if n_reqs > self.last_pos_host.shape[0]:
+                self.last_pos_host = torch.zeros(
+                    n_reqs * 2, dtype=torch.int64, device="cpu", pin_memory=True
+                )
+            ends = torch.as_tensor(
+                input_batch.query_start_loc_np[1 : n_reqs + 1].astype("int64") - 1,
+                device=positions.device,
+            )
+            self.last_pos_host[:n_reqs].copy_(positions[ends], non_blocking=True)
         self.hashes_ready.record()
         # The one host sync per step. It also orders this step's rewrite of the
         # pinned buffers after the previous step's async H2D copies from them.
         self.hashes_ready.synchronize()
+        if timing:
+            _t1 = _kai_time.perf_counter()
         requests = []
         for engram in self.engrams:
             local = host[:, engram.layer_hash_index, :].to(torch.int64)
             rel, owned = engram.embed_tokens.disk_rel_owned(local)
             requests.append((engram.embed_tokens.disk, rel, owned))
         rows_per_layer = gather_dequant_many(requests)
+        if timing:
+            _t2 = _kai_time.perf_counter()
+            _acc = self._t_acc[0 if n <= 16 else 1]
+            _acc[0] += _t1 - _t0
+            _acc[1] += _t2 - _t1
+            _acc[4] += 1
+            _rows = sum(int(r[1].numel()) for r in requests)
+            _uniq = sum(int(torch.unique(r[1]).numel()) for r in requests)
+            _acc[5] += _rows
+            _acc[6] += _uniq
+            # the H2D tail is timed by the caller-visible end of stage(): record
+            # the total when the call returns (see the return sites below)
+            self._t_mid = _t2
+            self._t_last = (_rows, _uniq, _t2 - _t1)
         for engram, rows, buf in zip(self.engrams, rows_per_layer, self.rows_host):
             staged = buf[:n]
             staged.copy_(rows.view(n, self.local_heads, self.dim))
             engram.staged_rows[:n].copy_(staged, non_blocking=True)
         self.num_staged = n
+        if timing:
+            _t3 = _kai_time.perf_counter()
+            _acc = self._t_acc[0 if n <= 16 else 1]
+            _acc[2] += _t3 - self._t_mid
+            _acc[3] += _t3 - _t0
+            self._t_last_total_ms = 1e3 * (_t3 - _t0)
+            # stage 6 defines _KAI_MMAP (memmap gather); resolved at call time
+            _mode = (
+                "mmap"
+                if globals().get("_KAI_MMAP", False)
+                and self._t_last[0] >= globals().get("_KAI_MMAP_MIN_ROWS", 0) * len(self.engrams)
+                else "preadv"
+            )
+            if n >= 1024:
+                logger.info(
+                    "Engram prefill stage [%s]: %d tokens, %d rows (%d unique), "
+                    "gather+dequant %.1f ms, total %.1f ms",
+                    _mode, n, self._t_last[0], self._t_last[1],
+                    1e3 * self._t_last[2], 1e3 * (_t3 - _t0),
+                )
+            if _acc[4] % (200 if n <= 16 else 20) == 0:
+                logger.info(
+                    "Engram stage timing [%s, %s]: total %.2f ms/step (sync wait %.2f, "
+                    "disk+dequant %.2f, h2d+rest %.2f), %.1f rows/step (%.1f unique), "
+                    "avg of %d steps",
+                    _mode, "decode" if n <= 16 else "prefill",
+                    1e3 * _acc[3] / _acc[4],
+                    1e3 * _acc[0] / _acc[4],
+                    1e3 * _acc[1] / _acc[4],
+                    1e3 * _acc[2] / _acc[4],
+                    _acc[5] / _acc[4],
+                    _acc[6] / _acc[4],
+                    _acc[4],
+                )
+                for i in range(7):
+                    _acc[i] = 0
+        if prefetch:
+            try:
+                self._hash_ahead(input_batch, req_states, positions, query_start_loc, lookback_token_ids)
+            except Exception as e:  # noqa: BLE001 - prefetch is best effort
+                logger.warning_once("Engram prefetch disabled after error: %r", e)
+                self._prefetch = False
         return n