Repository navigation
[DeepSeek-V4.1] Enable Engram disk offload for memory-constrained nodes - #62238
abdmalik-amd wants to merge 1 commit into
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: 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. 🚀 |
|
Documentation preview: https://vllm--62238.org.readthedocs.build/en/62238/ |
mergify
Bot
added
documentation
|
This pull request has merge conflicts that must be resolved before it can be |
Add EngramConfig.disk_offload for DeepSeek V4.1. The Engram tables stay in the checkpoint's safetensors files: after loading, each rank memory-maps its own row range, and each step gathers the distinct rows it needs through the page cache into fixed GPU buffers, where the existing lookup kernel (SORTED mode) reads them. Results are bit-identical to the in-memory table. The lookup is an eager break, so it runs eagerly or under breakable PIECEWISE CUDA graphs; other graph modes are rejected at startup. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Abdullah Malik <abdmalik@amd.com>
9320979 to
93bb234
Compare
This branch has not been deployed
1 participant
Purpose
DeepSeek-V4.1-Flash's Engram tables are about 197 GB, so they fit neither in GPU memory nor, with
cpu_offload, in host memory on a 128 GB node. At TP=4, each rank owns about 50 GB of table rows on top of about 78 GB of weights. On an APU (Strix Halo, DGX Spark), host and GPU memory are the same RAM.This PR adds
--engram-config '{"disk_offload": true}'. The tables stay in the checkpoint's safetensors files:_engram_lookup_kernel(SORTED mode) runs over that compact table, so results are bit-identical to the in-memory table.The gather needs the hash ids on the host, so the lookup is an
@eager_break_during_capturesegment. It runs with--enforce-eager, or withVLLM_USE_BREAKABLE_CUDAGRAPH=1 -cc.cudagraph_mode=PIECEWISE. Startup rejects other graph modes, the host-memory options, and models other than DeepSeek V4.1.The change is in the shared Engram code and is not tied to a platform. It was tested on gfx1151 and should equally serve CUDA unified-memory machines such as DGX Spark (not tested there).
Duplicate-work check
host_file_gather) #58815 / [Qwen4Exp] Checkpoint-mapped PLE storage for unified-memory GPUs (DGX Spark) #58439 add file-backed PLE tables for Qwen4Exp on CUDA. This PR is an independent implementation for DeepSeek V4.1 and uses the samedisk_offloadname, so the two compose.Test Plan
New tests check bit-exact lookups at TP 1/2/4/8 for single-file and sharded checkpoints, and a breakable-graph capture plus replay.
End-to-end:
deepseek-ai/DeepSeek-V4.1-Flash-1017on 4x Strix Halo (gfx1151, 128 GB), TP=4 over Ray,--enforce-eager.Test Result
On
mainthe model does not load on these nodes: the Ray worker dies pinning its ~47 GiB of table rows. Withdisk_offload:Graph mode is covered by the unit test only. On this cluster (RCCL over TCP), graph mode hangs on the first request in a collective; the same hang occurred with DeepSeek-V4-Flash on stock kernels.
AI assistance
AI assistance was used for this change. The submitting human has reviewed every changed line and ran the tests reported above.