Add DGX SPARK memory profiling fix mod and optimized recipe - #286
Open
bandageshi wants to merge 1 commit into
Open
bandageshi wants to merge 1 commit into
bandageshi wants to merge 1 commit into
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
On DGX SPARK,
torch.compileis often triggered during the vLLM initialization warmup (profile_run). The Triton/JIT compilation process allocates a massive temporary CPU heap and OS page caches (up to 30GB+). Because host RAM and GPU memory share the same physical pool on UMA, vLLM's memory snapshotting (which relies onpsutil.virtual_memory().availableon UMA) misidentifies this transient compiler peak as permanent non-torch memory, subtracting it from the KV Cache allocation. Once initialized, the compiler memory is freed, but the KV Cache size remains permanently starved (e.g., restricted to only 9.5 GiB instead of using the full available 40+ GiB with 60 GB allocation).Solution
non_torch_increase(capped at 2 GiB for standard NCCL/CUDA driver overhead) . This prevents transient host-side JIT compiler memory and file page caches from starving the KV Cache, successfully recovering the lost ~31.6 GiB of memory for inference (improving KV token capacity by 4.3x).recipes/diffusion-gemma-nvfp4-thinking-optimized.yaml) for DiffusionGemma-26B-A4B-IT-NVFP4 with thinking enabled, embedding this UMA memory profiling fix as a reusable mod.