Tyler Michael SmithandClaude 5b4f6d5284 Use cumem allocator for the KV cache by default ...

Enable the cumem (CUDA VMM) allocator by default on CUDA and ROCm
platforms so users get stable physical pages for KV cache without
needing to set --enable-cumem-allocator explicitly. This is required
for MNNVL KV transfers.

Changes:
- ModelConfig.enable_cumem_allocator defaults to None, resolved to
  True when the cumem C extension is available
- Cached fabric handle probe in csrc/cumem_allocator.cpp uses a real
  cuMemCreate probe instead of trusting cuDeviceGetAttribute, with
  fallback to POSIX FD when fabric handles aren't available
- Safety-net fallback: if fabric was probed as available but a real
  allocation still fails, updates the cache and retries with POSIX FD
- Weight loading bypasses cumem pool when sleep mode is off (cumem is
  only needed for KV cache stability in that case)
- Guard against stale engine allocations corrupting the singleton

Co-authored-by: Claude <[email protected]>
Signed-off-by: Tyler Michael Smith <[email protected]>

2026-06-10 15:25:35 -04:00

..

__init__.py

[CI/Build] Move test_utils.py to tests/utils.py (#4425)

2024-05-13 23:50:09 +09:00

test_basic_correctness.py

Make Gemma and Gemma 2 accept inputs_embeds like Gemma 3 (#36787)

2026-03-11 18:12:43 +00:00

test_cpu_offload.py

[Feature] Support CPU Offloading without Pytorch Pinned Memory that leads to doubled allocation (#32993)

2026-02-13 08:11:26 -08:00

test_mem.py

Use cumem allocator for the KV cache by default

2026-06-10 15:25:35 -04:00

test_prefetch_offload.py

[offloader] v2: Hide weight onloading latency via prefetching (#29941)

2026-02-25 17:20:59 -08:00