Skip to content

Add DGX SPARK memory profiling fix mod and optimized recipe - #286

Open
bandageshi wants to merge 1 commit into
eugr:mainfrom
bandageshi:fix-diffusiongemma-memory-allocation-problem
Open

bandageshi wants to merge 1 commit into
eugr:mainfrom
bandageshi:fix-diffusiongemma-memory-allocation-problem

Conversation

@bandageshi

Copy link
Copy Markdown

Problem

On DGX SPARK, torch.compile is 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 on psutil.virtual_memory().available on 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

  1. I introduce a cap on 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).
  2. Added a new optimized serving recipe (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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant