recipes/docker-images/mia-tp4/adapter/engram_prefetch.py

"""Engram rows fetched on a side stream as soon as the hash ids exist.

Measured 2026-09-18 with the live profiler at c1: 2.2-2.4 ms of every 48 ms step was GPU idle
before the two `_engram_gather_kernel` launches on real text (1.0 ms on sparkDash prose). The
stall is NVMe latency: the host callback that serves the lookup reads every cache miss
(O_DIRECT 8 KB pread, ~200 us) while the graph waits, once per Engram layer, and the all-reduce
then waits for the slowest rank.

The hash ids for both Engram layers are computed at the start of the forward, ~1 ms of GPU work
before layer 1 needs them and ~15 ms before layer 14 does. This adapter forks a side stream
right after the hasher returns, runs the ids copy, the host lookup and the row copies for every
Engram layer there, and joins the side stream only when the layer's gather runs. Inside a CUDA
graph that is an ordinary fork/join branch with a host node on it; the model, the hasher and
the row store are unchanged, and the gather reads the same staging buffers as before.

Gate ``DSV41_ENGRAM_PREFETCH=1``. ``DSV41_ENGRAM_PREFETCH_CHECK=1`` additionally runs the old
synchronous lookup after each prefetched gather and accumulates the number of gathers whose
output differed (logged every DSV41_STATS_SECONDS; expected 0).
"""
import ctypes as C
import logging
import os
import threading
import time

import torch

import engram_backend as eb

logger = logging.getLogger(__name__)
_EMBEDS = []
_state = {"disabled": False}


def enabled():
    return os.environ.get("DSV41_ENGRAM_PREFETCH", "0").strip() in ("1", "on", "true")


def _checking():
    return os.environ.get("DSV41_ENGRAM_PREFETCH_CHECK", "0").strip() == "1"


def _staging(self, indices, count):
    """Same staging layout as engram_backend.owned, created here when the prefetch is first."""
    capacity = 1 << (count - 1).bit_length()
    key = (indices.device.index, capacity)
    if key not in self._staging:
        if torch.cuda.is_current_stream_capturing():
            raise RuntimeError(f"Engram staging {key} must be warmed before graph capture")
        ids = torch.empty(capacity, dtype=torch.int64, pin_memory=True)
        w = torch.empty((capacity, 256), dtype=torch.uint8, pin_memory=True)
        s = torch.empty((capacity, 8), dtype=torch.uint8, pin_memory=True)
        dw, ds = w.to(indices.device), s.to(indices.device)
        sequential = torch.arange(capacity, dtype=torch.int64, device=indices.device)
        self._staging[key] = (ids, w, s, dw, ds, sequential)
    ids, w, s, dw, ds, sequential = self._staging[key]
    work_key = (indices.device.index, count)
    if work_key not in self._works:
        self._works[work_key] = eb.Work(self._store, ids.data_ptr(), w.data_ptr(), s.data_ptr(), count)
    return ids, w, s, dw, ds, sequential, self._works[work_key]


def _prefetch(self, indices):
    count = indices.numel()
    self._pending = None
    if not count or self.rows == 0:
        return
    if self._side is None:
        if torch.cuda.is_current_stream_capturing():
            raise RuntimeError("Engram prefetch stream must exist before graph capture")
        self._side = torch.cuda.Stream(device=indices.device)
    ids, w, s, dw, ds, sequential, work = _staging(self, indices, count)
    main = torch.cuda.current_stream(indices.device)
    self._side.wait_stream(main)  # fork: the hash ids are ready on main
    with torch.cuda.stream(self._side):
        ids[:count].copy_(indices.reshape(-1), non_blocking=True)
        error = eb._cuda.cudaLaunchHostFunc(self._side.cuda_stream,
                                            C.cast(eb._lib.row_store_lookup, eb.P), C.addressof(work))
        if error:
            raise RuntimeError(f"CUDA Engram host callback failed: {error}")
        dw[:count].copy_(w[:count], non_blocking=True)
        ds[:count].copy_(s[:count], non_blocking=True)
    self._pending = (indices.data_ptr(), tuple(indices.shape), count, dw, ds, sequential)


def _prefetch_all(hash_ids):
    if _state["disabled"] or hash_ids.dim() != 3 or hash_ids.shape[0] == 0:
        return
    for emb in _EMBEDS:
        if emb._hash_index is None:
            continue
        _prefetch(emb, hash_ids[:, emb._hash_index])


def install(module):
    if not enabled():
        return
    Emb, Eng, Hasher = module.EngramEmbedding, module.Engram, module.EngramHasher
    for name in ("_staging", "_works", "_store", "rows"):
        pass  # set by engram_backend.install's __init__
    orig_init = Emb.__init__

    def init(self, *args, **kwargs):
        orig_init(self, *args, **kwargs)
        self._hash_index = None
        self._pending = None
        self._side = torch.cuda.Stream() if torch.cuda.is_available() else None
        self._check_acc = None
        _EMBEDS.append(self)

    Emb.__init__ = init
    orig_eng_init = Eng.__init__

    def eng_init(self, *args, **kwargs):
        orig_eng_init(self, *args, **kwargs)
        self.embed._hash_index = self.layer_hash_index

    Eng.__init__ = eng_init
    orig_owned = Emb._owned_rows
    check = _checking()

    def owned(self, indices):
        p = self._pending
        if not p or p[0] != indices.data_ptr() or p[1] != tuple(indices.shape) or p[2] != indices.numel():
            return orig_owned(self, indices)
        self._pending = None
        from sglang.kernels.ops.embeddings.engram_gather import engram_gather

        _, _, count, dw, ds, sequential = p
        torch.cuda.current_stream(indices.device).wait_stream(self._side)  # join
        out = self._empty(indices)
        engram_gather(dw.data_ptr(), ds.data_ptr(), sequential[:count], out.view(-1, 256), 256, 32)
        if check:
            reference = orig_owned(self, indices)
            if self._check_acc is None:
                self._check_acc = torch.zeros((), dtype=torch.int64, device=indices.device)
            self._check_acc += (out != reference).any().to(torch.int64)
        return out

    Emb._owned_rows = owned
    orig_forward = Hasher.forward

    def forward(self, input_ids, forward_batch):
        hash_ids = orig_forward(self, input_ids, forward_batch)
        try:
            _prefetch_all(hash_ids)
        except Exception as exc:
            if not _state["disabled"]:
                _state["disabled"] = True
                logger.warning("DSV41 Engram prefetch disabled after error: %r", exc)
            for emb in _EMBEDS:
                emb._pending = None
        return hash_ids

    Hasher.forward = forward
    logger.warning("DSV41 Engram prefetch ARMED: rows fetched on a side stream after the hasher%s",
                   " (CHECK mode: old path re-run and compared)" if check else "")
    if check:
        period = float(os.getenv("DSV41_STATS_SECONDS", "60"))

        def loop():
            while True:
                time.sleep(period)
                counts = [int(e._check_acc.item()) if e._check_acc is not None else -1 for e in _EMBEDS]
                logger.warning("DSV41 Engram prefetch CHECK: gathers differing from the synchronous path per layer = %s", counts)

        threading.Thread(target=loop, daemon=True, name="engram-prefetch-check").start()