[wangyicong52]

Overview

DeepSeek-V4.1 applies trained Engram modules at configured transformer layers. If a pipeline partition starts a non-initial stage on one of those layers, the stage-entry control-flow branch calls mHC pre directly and bypasses the normal Engram branch. Certain valid PP partitions can therefore silently skip an Engram operation for every request. The real DeepSeek-V4.1-Flash configuration triggers this at layer 14 with VLLM_PP_LAYER_PARTITION=14,26.

This PR restores the missing NVIDIA operation by applying the local stage-entry Engram after receiving the previous stage's completed residual stream and before mHC pre.

AMD is affected by the same issue. We currently only have access to NVIDIA hardware, so this PR fixes only the NVIDIA implementation for now.

Claims

• Restore Engram when its layer is the first layer of a non-initial pipeline stage.
• Preserve the incoming pre_mix, selected Engram hash column, mask, and existing no-hash behavior.
• Add a non-identity regression test to the NVIDIA model CI job.

Validation

Environment

• Hardware: two 8 x NVIDIA H20 96GB nodes.
• Image: vllm/vllm-openai:nightly-49cb2bbf71cfd36ee5aa33d1e184b1cddff55294.

vLLM launch

Both arms used identical arguments:

VLLM_PP_LAYER_PARTITION=14,26 \
VLLM_MOE_SKIP_PADDING=1 \
vllm serve /models/DeepSeek-V4.1-Flash \
  --host 0.0.0.0 \
  --port 18001 \
  --served-model-name deepseek-v4.1-flash \
  --language-model-only \
  --tokenizer-mode deepseek_v41 \
  --attention-backend FLASHMLA_SPARSE_DSV41 \
  --kv-cache-dtype fp8_ds_mla \
  --block-size 64 \
  --tensor-parallel-size 4 \
  --pipeline-parallel-size 2 \
  --decode-context-parallel-size 1 \
  --prefill-context-parallel-size 1 \
  --distributed-executor-backend mp \
  --enable-expert-parallel \
  --all2all-backend allgather_reducescatter \
  --engram-config '{"cpu_offload":true}' \
  --moe-backend humming \
  --max-model-len auto \
  --max-num-batched-tokens 8192 \
  --max-num-seqs 64 \
  --gpu-memory-utilization 0.9 \
  --no-enable-prefix-caching \
  --no-swa-bounded-replay \
  --enforce-eager \
  --no-async-scheduling

Verify Engram execution at the pipeline boundary

This check verifies that layer 14 executes its trained Engram module before mHC pre when it is the first layer of a pipeline stage.

We compared three configurations using the real DeepSeek-V4.1-Flash checkpoint:

• Without this PR, PP=1: the reference configuration. Layer 14 runs within the same pipeline stage as the preceding layers.
• Without this PR, PP=2 (14,26): the configuration that exposes the bug. The first stage contains 14 layers, so the second stage starts at layer 14.
• With this PR, PP=2 (14,26): the same pipeline partition with the fix applied, tested independently on both H20 nodes.

Each configuration received the same fixed 256-token prompt and generated two tokens, each constrained to a single allowed token ID. This kept the token sequence fixed across configurations. Each configuration was run twice in eager mode.

Was layer 14's Engram actually executed?

Configuration	Layer 14 Engram executed?	Input passed to layer 14's mHC pre
Without this PR, PP=1 — reference	Yes	Engram output
Without this PR, PP=2 (14,26)	No	Tensor received from the previous stage, without applying Engram
With this PR, PP=2 (14,26) — both nodes	Yes	Engram output

The unpatched PP=2 configuration had valid Engram hashes and masks available at layer 14, but never called its Engram module. The missing operation was therefore caused by the stage-entry control flow.

Did mHC receive the correct tensor?

We recorded SHA-256 checksums of the full tensors at the pipeline handoff, Engram input/output, and mHC pre input. These checksums track tensor contents; they are separate from the token hashes used by Engram.

For every TP rank, measured step, and repeat, the checks confirmed these paths:

• Without this PR, PP=2: the tensor sent by the previous stage was passed directly to mHC pre. Engram was skipped.
• With this PR, PP=2: the received tensor was passed to Engram, and Engram's output was passed to mHC pre.
• PP=1 reference: Engram's output was also passed to mHC pre.

These comparisons verify the data flow within each run; they do not require bitwise-identical outputs between PP=1 and PP=2.

Did Engram make a measurable change?

For runs that executed Engram, we measured the difference between its input and output at layer 14:

Step	RMS of the Engram update	Tensor elements whose stored values changed
Prefill	0.03445–0.03557	67.22%–67.82%
Decode	0.05899–0.06505	81.83%–83.14%

The update RMS is RMS(Engram output − Engram input). These nonzero differences confirm that skipping this trained Engram module changes the tensor passed to mHC pre.

This check confirms that the affected PP partition skips a trained Engram operation and that this PR restores it before mHC pre on both tested nodes. It does not establish an accuracy improvement: captured logits varied across repeated runs and nodes by an amount comparable to the patch effect. The GSM8K results below report the separate end-to-end accuracy evaluation.

Full GSM8K comparison

Both configurations used PP=2 (14,26) with identical evaluation settings: 4-shot, temperature 0, seed 42, and a maximum of 8,192 output tokens.

evalscope eval \
  --eval-type openai_api \
  --api-url http://<node-ip>:18001/v1 \
  --api-key EMPTY \
  --model deepseek-v4.1-flash \
  --model-id deepseek-v4.1-flash \
  --datasets gsm8k \
  --dataset-hub modelscope \
  --dataset-args '{"gsm8k":{"few_shot_num":4}}' \
  --seed 42 \
  --eval-batch-size 64 \
  --generation-config '{"temperature":0.0,"top_p":1.0,"max_tokens":8192,"reasoning_effort":"high","seed":42,"timeout":10000}' \
  --timeout 10000 \
  --no-timestamp

Configuration	Correct / Total	Accuracy	API errors	Pod restarts
Without this PR	1283 / 1319	97.27%	0	0
With this PR	1283 / 1319	97.27%	0	0

Seven questions were correct only without the patch, and seven only with it; aggregate accuracy was unchanged.

We have not observed an end-to-end accuracy loss in this GSM8K evaluation. However, the confirmed omission of a trained Engram operation remains a correctness bug with potential model-quality impact that this benchmark has not demonstrated.

The focused regression test and full pre-commit suite also passed.

Details

At the end of a pipeline stage, the previous rank completes its final mHC post and sends the full 3-D residual stream together with pre_mix. The next stage starts with residual=None, assigns that received stream to residual, and calls mHC pre directly. Because normal Engram injection lives in the following elif, an Engram on the stage's first layer is silently skipped.

This change injects the local Engram after accepting the completed stream and before mHC pre. It does not repeat the previous stage's post or alter first-layer embeddings, in-stage Engram execution, or batches without Engram hashes.

The real DeepSeek-V4.1-Flash configuration has Engram layers 1 and 14. A pipeline partition beginning its second stage at layer 14 can therefore reach this path. Open PR searches for the same Engram pipeline-stage fix found no direct duplicate; related PCP work such as #62183 does not implement this injection.

AI assistance was used for implementation support and a read-only Codex CLI review. Human line-by-line review is still required before requesting maintainer review.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[ZJY0516]

Thanks for fixing it

[github-actions]

✅ @wangyicong52, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[ZJY0516]

/ci run

[github-actions]

✅ Triggered Buildkite CI #95966 for commit 0b33a5375ad6.

ℹ️ This PR is 94 commits behind upstream main. Merged upstream main into this branch as 0b33a5375ad6 before starting CI. Pull the branch before pushing more changes.