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

[0z5a]

Purpose

Add the default-off --engram-config '{"lookup_overlap":true}' option. Stage each Engram layer's local embedding rows on its own CUDA stream, then wait for the completion event immediately before the decoder consumes those rows. Keep staging storage separate between microbatches and pass the prepared tensor through the decoder call. Breakable CUDA graph capture uses the existing main-stream lookup path.

This change preserves the upstream SP all-gather exchange. It does not include the all-to-all work in #56219 or the lookback reconstruction in #56224. The duplicate-work check also found #56230, which changes DP table ownership/communication; that is different from the lookup scheduling here.

Based on #56214 (dsv41-feat, c9d909e802a39292f54101bff8a36096761ea605). The PR targets that feature branch so the model implementation is not repeated in this diff.

Rebased on 2026-09-11 after the feature branch was rewritten. Changed-file Python parsing and all applicable pre-commit hooks passed. The older H100 integration measurements below belong to the prior revision based on e47aa780bccf59f59dfa2cbb18e17a10b4fe69ba. A selected same-head H200 screening result is reported separately below, with an unresolved concurrent-output difference.

AI assistance: OpenAI Codex assisted with implementation, review, test execution and preparation of this PR. This draft does not claim that a human has completed a line-by-line review.

Test Plan

Fresh split: all applicable pre-commit hooks passed for the four changed files, and all changed Python files parsed successfully.

The existing GPU tests were extended for lookup overlap on/off, eager/full/breakable capture, CPU/GPU table placement and TP2/TP4. The integration validation ran the prepared-row and collective graph tests in tests/kernels/test_engram.py; CUDA execution was not repeated on the split branch.

Test Result

2026-09-11: selected positive H200 workload on the current head

Both arms used 2ae34345ee5919130a6a08fe91f969d174b8b69e with lookup_overlap=false/true, the pinned V4.1 Flash checkpoint, 4 x H200 NVL over PCIe (no NVLink), TP4/PP1/DP1 + EP, eager V1, Engram CPU offload and zero generic CPU offload. Prefix caching was enabled; the cache was reset before four warmup requests and sixteen measured requests per arm.

Input → output tokens	Concurrency	Overlap off, output tokens/s	Overlap on, output tokens/s	Observed change
512 → 128	4	71.39797	72.97452	+2.21%

All sixteen measured requests per arm succeeded and generated exactly 128 output tokens. This is one A/B pair and a selected positive workload, not the complete workload matrix or evidence of a general/statistically established speedup. Startup time is excluded.

Correctness remains unresolved: the separate serial and fixed-prefix probes matched, but one of four concurrent responses diverged at the fifth generated token; no A/A repeat was run to establish the cause. This result is not an output-equivalence or merge qualification.

Benchmark timing records, exact flags and provenance. OpenAI Codex assisted with execution and reporting.

Earlier combined H100 integration evidence

Recorded full-model runs used deepseek-ai/DeepSeek-V4.1-Flash, revision df42c109f1defefcbfcedbe7d905718a12266e40, on 4×SXM H100 80 GiB. Each positive matrix case checked six short questions, cold/cached long input with chunked prefill, and four unequal-length concurrent requests. “Text equal” means short-answer text equality to the TP4 reference, not logits/token-probability parity or a standard accuracy benchmark.

Evidence scope: GPU results below come from the earlier combined integration based on #56214, including companion changes. These older rows do not describe the separate same-head H200 screening reported above. Fresh split checks are listed separately; passing the integrated run does not establish isolated-branch equivalence.

Integrated configuration	QA answers correct	QA text equal to TP4	Long / chunked + prefix	Concurrent requests	Strict result
TP4 eager, synchronous reference	6/6	6/6	2/2 correct; cache reuse checked	4/4 correct	PASS
TP4 eager, lookup overlap	6/6	6/6	2/2 correct; cache reuse checked	4/4 correct	PASS
TP4 breakable decode graph, synchronous	6/6	6/6	2/2 correct; cache reuse checked	4/4 correct	PASS
TP4 breakable decode graph, overlap option	6/6	6/6	2/2 correct; cache reuse checked	4/4 correct	PASS
DP2×TP2 SP eager, synchronous	6/6	6/6	2/2 correct; cache reuse checked	4/4 correct	PASS
DP2×TP2 SP eager, overlap	6/6	6/6	2/2 correct; cache reuse checked	4/4 correct	PASS
DP2×TP2 SP breakable graph, synchronous	6/6	6/6	2/2 correct; cache reuse checked	4/4 correct	PASS
DP2×TP2 SP breakable graph, overlap option	6/6	6/6	2/2 correct; cache reuse checked	4/4 correct	PASS

Breakable graph capture was confirmed with FULL_DECODE_ONLY, sizes [1,2,4], and Torch compile disabled. During capture the overlap option uses the main-stream fallback. The DP2×TP2 integration also included the companion SP exchange; this PR retains upstream all-gather.

Serving A/B — TP4 eager, synchronous lookup → overlap

Input → output tokens	Synchronous output tokens/s	Overlap output tokens/s	Observed throughput change	TTFT median ms, sync → overlap	TPOT median ms, sync → overlap
512 → 128	40.15	40.57	+1.04%	421.1 → 419.1	96.13 → 95.70
2048 → 128	38.07	38.38	+0.82%	764.0 → 758.1	98.03 → 97.20

Each side ran three repeats of 16 requests at concurrency 4 after warmup, with fixed inputs, prefix cache disabled, and the same default sampling settings. Values are medians across repeats; percent changes use unrounded data. These small observed changes are not a statistically established speedup. No graph-overlap speedup is claimed.

Validation / memory measurement	Result	Scope
Fresh split pre-commit and changed Python parsing	PASS; 4 changed files	This PR's isolated source
Model-loading footprint, sync → overlap	71.22 → 71.22 GiB	Rank 0 startup log; no measured reduction
Available KV budget, sync → overlap	1.96 → 1.96 GiB	Rounded startup values, not request peak memory

---

PR description checklist

• Purpose and related work described.
• Test plan and actual results stated.
• Model evaluation limitations stated.
• AI assistance disclosed.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[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.

[0z5a]

Validation follow-up for 2ae34345ee5919130a6a08fe91f969d174b8b69e: the published H200 512→128, concurrency-4 result is a valid same-head overlap-off/on screening sample (+2.21% output throughput), with matching runtime, checkpoint and KV allocation. It is a selected positive workload from one A/B pair, not evidence of a general or statistically established improvement.

A separate correctness check matched the serial and fixed-prefix probes, but 1/4 concurrent responses diverged at the fifth generated token. No A/A control established the cause, so this must remain an unresolved correctness question rather than being dismissed as rounding or claimed as a proven PR regression. Useful follow-up would be A/A plus fixed-prefix diagnostics, repeated/interleaved A/B timing across representative lengths and concurrency, and current-head graph/TP/SP coverage. A parent-vs-PR comparison with the option disabled would also check default-path regressions; the existing off/on pair cannot establish that.

Contributions from anyone with access to sufficient GPU capacity are very welcome: reproducible GPU E2E benchmarks, targeted correctness checks, and focused fixes. Please include source/model revisions, GPU topology, offload/KV settings, exact commands, and raw results; repeat timing runs when making performance claims. I am happy to review and cherry-pick fixes while preserving the original authorship and clear attribution.

[0z5a]

Resubmitted as #56436 against main, since dsv41-feat was merged into main in #56214.