[Perf][Engram] Overlap DeepSeek-V4.1 lookup with decoder compute

tests/kernels/test_engram.py

@@ -662,13 +662,14 @@ def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens):
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
 @pytest.mark.parametrize("cpu_offload", [False, True])
 @pytest.mark.parametrize("capture", ["eager", "full", "breakable"])
-def test_engram_prepared_rows_survive_graph_breaks(cpu_offload, capture):
+@pytest.mark.parametrize("overlap", [False, True])
+def test_engram_prepared_rows_survive_graph_breaks(cpu_offload, capture, overlap):
     """Consume early lookup results after a break, with fresh IDs each replay."""
-    _run_engram_prepared_rows(cpu_offload, capture)
+    _run_engram_prepared_rows(cpu_offload, capture, overlap=overlap)
 
 
 def _run_engram_prepared_rows(
-    cpu_offload, capture, tp_size=1, rank=0, use_sequence_parallel=False
+    cpu_offload, capture, tp_size=1, rank=0, use_sequence_parallel=False, overlap=False
 ):
     from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
 
@@ -689,6 +690,7 @@ def _run_engram_prepared_rows(
         dtype=torch.bfloat16,
         device="cuda",
     )
+    engram._init_lookup_staging(1, overlap)
     # Match the non-contiguous per-layer slice of the model hash tensor.
     hashes = torch.randint(
         0,
@@ -762,17 +764,28 @@ def _engram_tp_worker(rank, tp_size, port):
         for cpu_offload in (False, True):
             for sp in (False, True):
                 for capture in ("eager", "compiled", "full", "breakable"):
-                    _run_engram_prepared_rows(cpu_offload, capture, tp_size, rank, sp)
+                    for overlap in (False, True):
+                        _run_engram_prepared_rows(
+                            cpu_offload, capture, tp_size, rank, sp, overlap=overlap
+                        )
     finally:
         cleanup_dist_env_and_memory()
 
 
-@pytest.mark.distributed(num_gpus=2)
+@pytest.mark.parametrize(
+    "tp_size",
+    [
+        pytest.param(2, marks=pytest.mark.distributed(num_gpus=2)),
+        pytest.param(4, marks=pytest.mark.distributed(num_gpus=4)),
+    ],
+)
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
-def test_engram_head_collectives_survive_graph_breaks():
-    """All-gather preserves head order and local SP tokens across graph replay."""
+def test_engram_head_collectives_survive_graph_breaks(tp_size):
+    """All-gather preserves staged SP tokens across graph and compile replay."""
     from vllm.utils.network_utils import get_open_port
 
-    if torch.accelerator.device_count() < 2:
-        pytest.skip("Requires two GPUs")
-    torch.multiprocessing.spawn(_engram_tp_worker, args=(2, get_open_port()), nprocs=2)
+    if torch.accelerator.device_count() < tp_size:
+        pytest.skip(f"Requires {tp_size} GPUs")
+    torch.multiprocessing.spawn(
+        _engram_tp_worker, args=(tp_size, get_open_port()), nprocs=tp_size
+    )

vllm/config/engram.py

@@ -47,6 +47,9 @@ class EngramConfig:
     """Shard embeddings across TP and all DP ranks when enabled.
     Otherwise, each DP rank has a separate TP-sharded embedding replica."""
 
+    lookup_overlap: bool = False
+    """Prefetch Engram rows on a separate CUDA stream while decoder layers run."""
+
     def verify_model_config(self, model_config: "ModelConfig | None") -> None:
         """Reject Engram configuration for models without n-gram embeddings."""
         from vllm.platforms import current_platform

vllm/models/deepseek_v4_1/common/engram.py

@@ -40,6 +40,7 @@
 import torch
 from torch import nn
 
+from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
 from vllm.config import VllmConfig, get_current_vllm_config
 from vllm.distributed import (
     get_tensor_model_parallel_rank,
@@ -53,6 +54,7 @@
 from vllm.triton_utils import tl, triton
 from vllm.utils.platform_utils import is_uva_available
 from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
+from vllm.v1.worker.ubatching import dbo_current_ubatch_id
 
 logger = init_logger(__name__)
 
@@ -939,14 +941,70 @@ def __init__(
             layout.head_dim,
             dtype=torch.bfloat16,
         )
+        parallel_config = get_current_vllm_config().parallel_config
+        self._init_lookup_staging(
+            parallel_config.num_ubatches if parallel_config.use_ubatching else 1,
+            engram_config.lookup_overlap if engram_config else False,
+        )
 
-    def prepare_embeddings(self, hash_ids: torch.Tensor) -> None:
-        """Gather this layer's rows on the main stream before decoder layers."""
-        self.embed_tokens.lookup(hash_ids, self.staged_rows[: hash_ids.shape[0]])
+    def _init_lookup_staging(self, num_slots: int, overlap: bool) -> None:
+        self._lookup_rows = [self.staged_rows] + [
+            torch.empty_like(self.staged_rows) for _ in range(num_slots - 1)
+        ]
+        self._lookup_streams = (
+            [
+                torch.cuda.Stream(device=self.staged_rows.device)
+                for _ in range(num_slots)
+            ]
+            if overlap
+            else []
+        )
+        self._lookup_ready = [torch.cuda.Event() for _ in self._lookup_streams]
+
+    def _lookup_slot(self) -> int:
+        rows = getattr(self, "_lookup_rows", None)
+        return dbo_current_ubatch_id() if rows is not None and len(rows) > 1 else 0
+
+    def wait_for_embeddings(self) -> None:
+        if (
+            getattr(self, "_lookup_streams", None)
+            and not BreakableCUDAGraphCapture.is_active()
+        ):
+            torch.cuda.current_stream().wait_event(
+                self._lookup_ready[self._lookup_slot()]
+            )
 
-    def embed(self, hash_ids: torch.Tensor) -> torch.Tensor:
+    def prepare_embeddings(self, hash_ids: torch.Tensor) -> torch.Tensor:
+        """Stage local heads after hashing, optionally overlapping decoder compute."""
+        slot = self._lookup_slot()
+        buffers = getattr(self, "_lookup_rows", None)
+        rows = (buffers[slot] if buffers is not None else self.staged_rows)[
+            : hash_ids.shape[0]
+        ]
+        streams = getattr(self, "_lookup_streams", None)
+        if streams and not BreakableCUDAGraphCapture.is_active():
+            stream = streams[slot]
+            stream.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(stream):
+                self.embed_tokens.lookup(hash_ids, rows, background=True)
+                self._lookup_ready[slot].record(stream)
+        else:
+            self.embed_tokens.lookup(hash_ids, rows)
+        return rows
+
+    def embed(
+        self, hash_ids: torch.Tensor, prepared_rows: torch.Tensor | None = None
+    ) -> torch.Tensor:
         """Gather heads, returning only local tokens when SP is enabled."""
-        rows = self.staged_rows[: hash_ids.shape[0]]
+        if prepared_rows is None:
+            self.wait_for_embeddings()
+            buffers = getattr(self, "_lookup_rows", None)
+            prepared_rows = (
+                buffers[self._lookup_slot()]
+                if buffers is not None
+                else self.staged_rows
+            )
+        rows = prepared_rows[: hash_ids.shape[0]]
         if self.embed_tokens.tp_size == 1:
             return rows
         if self.use_sequence_parallel:
@@ -974,11 +1032,12 @@ def forward(
         hidden_states: torch.Tensor,
         hash_ids: torch.Tensor,
         token_mask: torch.Tensor | None = None,
+        prepared_rows: 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."""
-        kv = self.wkv(self.embed(hash_ids).flatten(-2))
+        kv = self.wkv(self.embed(hash_ids, prepared_rows).flatten(-2))
         num_kv_tokens = hash_ids.shape[0]
         assert token_mask is None or token_mask.shape == (num_kv_tokens,)
         if self.use_sequence_parallel:

vllm/models/deepseek_v4_1/nvidia/model.py

@@ -293,6 +293,7 @@ def forward(
         residual: torch.Tensor | None = None,
         engram_hashes: torch.Tensor | None = None,
         engram_mask: torch.Tensor | None = None,
+        engram_rows: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         # The reference collapses each sublayer's input with the *previous*
         # sublayer's pre-mix: attention uses the pre-mix carried in (identity
@@ -343,6 +344,7 @@ def forward(
                     residual,
                     engram_hashes[:, self.engram.layer_hash_index],
                     engram_mask,
+                    prepared_rows=engram_rows,
                 )
             post_mix, res_mix, x, attn_pre = mhc_pre_delayed_tilelang(
                 residual,
@@ -556,6 +558,7 @@ def forward(
         # profile runs (KV cache unbound).
         engram_hashes: torch.Tensor | None = None
         engram_mask: torch.Tensor | None = None
+        engram_rows: dict[int, torch.Tensor] = {}
         if (
             self.engram_hash is not None
             and input_ids is not None
@@ -598,8 +601,10 @@ def forward(
                 for layer in islice(self.layers, self.start_layer, self.end_layer):
                     engram = getattr(layer, "engram", None)
                     if engram is not None:
-                        engram.prepare_embeddings(
-                            engram_hashes[:, engram.layer_hash_index]
+                        engram_rows[engram.layer_hash_index] = (
+                            engram.prepare_embeddings(
+                                engram_hashes[:, engram.layer_hash_index]
+                            )
                         )
 
         full_num_tokens = positions.shape[0]
@@ -623,6 +628,10 @@ def forward(
             islice(self.layers, self.start_layer, self.end_layer),
             start=self.start_layer,
         ):
+            prepared_rows = None
+            if layer.engram is not None and engram_hashes is not None:
+                layer.engram.wait_for_embeddings()
+                prepared_rows = engram_rows[layer.engram.layer_hash_index]
             hidden_states, residual, post_mix, res_mix, pre_mix = layer(
                 hidden_states,
                 positions,
@@ -633,6 +642,7 @@ def forward(
                 residual,
                 engram_hashes,
                 engram_mask,
+                prepared_rows,
             )
             if idx + 1 in self.aux_hidden_state_layers:
                 # Reconstruct the aux hidden state for draft models