From 2ae34345ee5919130a6a08fe91f969d174b8b69e Mon Sep 17 00:00:00 2001 From: 0z5a <0z5a@users.noreply.github.com> Date: Fri, 11 Sep 2026 09:13:51 +0800 Subject: [PATCH] [Perf][Engram] Overlap DeepSeek-V4.1 lookup with decoder compute Rebase the existing PR onto dsv41-feat at c9d909e802a39292f54101bff8a36096761ea605. Signed-off-by: 0z5a <0z5a@users.noreply.github.com> Assisted-by: OpenAI Codex --- tests/kernels/test_engram.py | 33 +++++++--- vllm/config/engram.py | 3 + vllm/models/deepseek_v4_1/common/engram.py | 71 ++++++++++++++++++++-- vllm/models/deepseek_v4_1/nvidia/model.py | 14 ++++- 4 files changed, 103 insertions(+), 18 deletions(-) diff --git a/tests/kernels/test_engram.py b/tests/kernels/test_engram.py index f7aa6dd03e79..09e237f227d2 100644 --- a/tests/kernels/test_engram.py +++ b/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 + ) diff --git a/vllm/config/engram.py b/vllm/config/engram.py index 2aede1868d55..6191eb2f84d6 100644 --- a/vllm/config/engram.py +++ b/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 diff --git a/vllm/models/deepseek_v4_1/common/engram.py b/vllm/models/deepseek_v4_1/common/engram.py index f30c3a5606d1..3b9dd9159c28 100644 --- a/vllm/models/deepseek_v4_1/common/engram.py +++ b/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: diff --git a/vllm/models/deepseek_v4_1/nvidia/model.py b/vllm/models/deepseek_v4_1/nvidia/model.py index 6d803736ec5c..7a1890d1dad4 100644 --- a/vllm/models/deepseek_v4_1/nvidia/model.py +++ b/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