[ivanium]

Overview

Targets the branch of vllm-project#63921. In the latency_specialized decode, the Engram WKV chain now runs ahead on its side stream for GPU-resident tables too, not only for host-offloaded ones. The side stream waits for the current stream, where GPU-resident rows are looked up inline. The chain keeps the NCCL gathers of Engram.forward.

Claims

• The median decode step drops from 4165 µs to 4075 µs, a 90 µs (2.2%) gain over [Perf][DSv4.1] Extend the latency_specialized decode on NVIDIA vllm-project/vllm#63921 (DeepSeek-V4.1-Flash, TP4, one request, GPU-resident tables, unprofiled).
• The output is bitwise unchanged.

Validation

Unprofiled step time, FlashInfer autotune off, on one node, 3 interleaved restarts × 3 samples per arm. Setup: TP4 on GB200, a 128k context, DSpark with 5 tokens, GPU-resident tables.

Engram path	Median step (µs)	Min (µs)
Inline (vllm-project#63921)	4165	4068
Run ahead (this PR)	4075	4046

Tests on one GB200 node (4 GPUs):

Test	Runs	Result
test_kernels.py::test_engram_wkv_prefetch_matches_inline (host-offloaded and GPU-resident tables)	1	3 passed
test_layers.py::test_specialized_decode_matches_the_model, host-offloaded tables (the test's default)	3	3 passed
The same test with cpu_offload: False	2	2 passed

test_layers.py skips off SM103, so I relaxed its skipif and the backend's SM103 gate locally to run it on GB200. No accuracy eval was run, because the output is bitwise unchanged.

Limitations:

• Measured on GB200 (SM100) only. The new code has nothing architecture-specific.
• There is no end-to-end serving benchmark.

Details

• An earlier revision of this PR also moved the chain's two all-gathers to a LamportAllGather of their own (about 17 µs more). It hung test_specialized_decode_matches_the_model (12 requests decoding together) in 2 of 3 runs, while [Perf][DSv4.1] Extend the latency_specialized decode on NVIDIA vllm-project/vllm#63921 alone passed 3 of 3, and the same tree with only those gathers on NCCL passed 3 of 3. So the gathers stay on NCCL, and the module docstring says why. The inferred mechanism: the Lamport threads spin for peer data, and on a side stream they can block the executor's kernels while a peer waits.
• No open PR changes this path. [ROCm][DSv4.1] Overlap the Engram WKV chain with the decoder layers vllm-project/vllm#63326 and [ROCm][DSv4.1] Shard the Engram gate so TP ranks gather only the value slice vllm-project/vllm#63367 change Engram on ROCm only.

AI assistance (Claude Code) was used to write this change and this description.

🤖 Generated with Claude Code

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

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors.

You ask your reviewers to trigger select CI tests on top of fastcheck CI.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

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

🚀

[yaoyaoding]

Thanks for this! I tried it on GB300 TP4 on top of the current dsv41-specialized-decode-perf (b7bb822). I used the default config (Engram cpu_offload: true, FlashInfer autotune on, latency_specialized, MTP with synthetic AL 3.51) and the same decode-step ABBA we use for the PR: batch 1 at 524K tokens, and batch 8 and 32 at 84K–648K tokens a request.

1. It hangs once a step reaches about 12 requests (2 of 2 runs). Both servers built from this branch stopped making progress while the batch-32 load ramped up. The last scheduled steps were 72 tokens (12 requests × 6) and 90 tokens (15 × 6). About five minutes later the engine died with TimeoutError: RPC call to sample_tokens timed out. The baseline server ran the same batch-32 load to completion. Batch 1 (6 tokens a step) and batch 8 (48 tokens) were fine on both trees. test_kernels.py -k engram and the new test_numerics.py pass, so the problem only shows up when serving at larger step sizes.

My guess, which I haven't verified: the chain's Lamport gathers spin on the side stream while the main stream runs its own spinning kernels (the executor's Lamport collectives and other cross-rank waits). Once the gather grid is large enough, spinning CTAs on one rank can hold SMs that a kernel on the other stream needs before it can finish. The ranks then wait on each other. NCCL on the side stream doesn't run into this. One quick test: cap the gather's grid, or fall back to NCCL above some row count, and rerun batch 32.

2. The default path only gets the Lamport half. Our recipe and benchmarks use cpu_offload: true, which is also the EngramConfig default, and with it the chain already ran ahead before this change. The new GPU-resident run-ahead doesn't apply there; only replacing the two NCCL all-gathers does.

3. Before the hang, the gain is small. Decode step in µs, from the blocks that completed (one baseline block, two blocks of this branch, two client runs each):

batch	baseline	this branch
1	3936	3868 / 3918
8	7658	7555 / 7643

That's about −1.1% at batch 1 and −0.8% at batch 8. With only one baseline block, this is within the spread we usually see between server restarts.

4. Minor: "bitwise unchanged" has one exception. The Lamport gather sends the sentinel word 0x80000000 as 0, so −0.0 arrives as +0.0. That doesn't matter for the GEMM, but the description should mention it.

Otherwise the code looks right to me. The transposes match gather_output's layout, and the chain uses its own LamportAllGather instance on a single stream, so the call-order requirement holds.

So I'm not merging this into the PR for now. I'm happy to rerun the ABBA once the hang is fixed.