[Bugfix][DeepSeek V4.1] Inject Engram at pipeline stage entry

.buildkite/test_areas/models_basic.yaml

@@ -70,11 +70,12 @@ steps:
   - tests/models/test_deepseek_v41_decoder_replay_layers.py
   - tests/models/test_deepseek_v41_replay_batch.py
   - tests/models/test_deepseek_v41_pcp.py
+  - tests/models/test_deepseek_v41_pipeline.py
   - tests/models/transformers/test_layer_registry_e2e.py
   - tests/models/test_deepseek_v41_replay_start.py
   - tests/models/test_deepseek_v4_vl_stream_weights.py
   commands:
-    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_pcp.py models/test_deepseek_v4_vl_stream_weights.py
+    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_pcp.py models/test_deepseek_v41_pipeline.py models/test_deepseek_v4_vl_stream_weights.py
   mirror:
     amd:
       label: ":amd: (MI355) Basic Models (Other)"

tests/models/test_deepseek_v41_pipeline.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from importlib import import_module
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.skip_global_cleanup
+def test_stage_entry_injects_engram_before_mhc_pre(monkeypatch):
+    model_module = import_module("vllm.models.deepseek_v41.nvidia.model")
+
+    class StopAfterMhcPre(Exception):
+        pass
+
+    hidden = torch.arange(24, dtype=torch.float32).reshape(2, 4, 3)
+    hashes = torch.arange(6).reshape(2, 3)
+    mask = torch.tensor([True, False])
+    pre_mix = torch.randn(2, 4)
+    engram_calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]] = []
+    mhc_pre_calls: list[tuple[torch.Tensor, torch.Tensor | None]] = []
+
+    class Engram:
+        layer_hash_index = 1
+
+        def __call__(self, residual, layer_hashes, engram_mask):
+            engram_calls.append((residual, layer_hashes, engram_mask))
+            return residual + 10
+
+    def mhc_pre(residual, *args, pre_mix=None, **kwargs):
+        mhc_pre_calls.append((residual, pre_mix))
+        raise StopAfterMhcPre
+
+    layer = SimpleNamespace(
+        mhc_stream=None,
+        engram=Engram(),
+        hc_attn_fn=None,
+        hc_attn_scale=None,
+        hc_attn_base=None,
+        rms_norm_eps=1e-6,
+        hc_eps=1e-6,
+        hc_post_alpha=1.0,
+        hc_sinkhorn_iters=1,
+        attn_norm=SimpleNamespace(weight=None, variance_epsilon=1e-6),
+    )
+    monkeypatch.setattr(model_module, "mhc_pre_delayed_tilelang", mhc_pre)
+
+    with pytest.raises(StopAfterMhcPre):
+        model_module.DeepseekV4DecoderLayer.forward(
+            layer,
+            hidden,
+            torch.arange(2),
+            None,
+            pre_mix=pre_mix,
+            engram_hashes=hashes,
+            engram_mask=mask,
+        )
+
+    assert len(engram_calls) == len(mhc_pre_calls) == 1
+    engram_residual, layer_hashes, engram_mask = engram_calls[0]
+    assert engram_residual is hidden
+    assert engram_mask is mask
+    torch.testing.assert_close(layer_hashes, hashes[:, 1])
+    mhc_residual, mhc_pre_mix = mhc_pre_calls[0]
+    torch.testing.assert_close(mhc_residual, hidden + 10)
+    assert mhc_pre_mix is pre_mix

vllm/models/deepseek_v41/nvidia/model.py

@@ -542,6 +542,12 @@ def forward(
                 )
             else:
                 residual = x
+                if self.engram is not None and engram_hashes is not None:
+                    residual = self.engram(
+                        residual,
+                        engram_hashes[:, self.engram.layer_hash_index],
+                        engram_mask,
+                    )
                 post_mix, res_mix, x, attn_pre = mhc_pre(
                     residual,
                     self.hc_attn_fn,