[ROCm] Grant host access to VMM allocations in sleep-mode pool

[xiaohong42]

Summary

On large-BAR devices (MI300X), PyTorch ROCm (>=7.2) reads scalar tensors via direct host-side pointer dereference (_local_scalar_dense_cuda) instead of hipMemcpy. The sleep-mode VMM pool (hipMemCreate + hipMemMap) only grants GPU-side access, so the host dereference hits an unmapped address causing SIGSEGV that silently kills the worker with no Python traceback.

Root Cause

create_and_map() in csrc/cumem_allocator.cpp calls hipMemSetAccess with only a device descriptor. PyTorch's large-BAR fast path assumes all allocator-managed memory is host-readable, which is true for hipMalloc but not for VMM allocations without an explicit host mapping.

Fix

Add a hipMemLocationTypeHost access descriptor alongside the existing device descriptor. This gives the VMM range a host mapping, making PyTorch's fast path valid.

Only the ROCm path is modified (#ifdef USE_ROCM); the CUDA path is unchanged.

Why not fix at the read site?

We considered detecting at read time whether a given device pointer is host-accessible, but no HIP API on ROCm 6.x/7.x provides this information correctly. All candidates were tested with native HIP C++ probes (probe_discriminator.cpp, probe_memgetaccess.cpp):

Candidate API	hipMalloc (should be fast-path)	VMM device-only (should be slow-path)	Distinguishes?
hipPointerGetAttributes().hostPointer	NULL	NULL	❌
hipPointerGetAttribute(HOST_POINTER)	err=1 (not supported)	err=1 (not supported)	❌
hipPointerGetAttribute(MAPPED)	-255	-255	❌
hipMemGetAccess(location=Host)	invalid argument	invalid argument	❌ ROCm 7.14 unsupported
hipMemRetainAllocationHandle	fails	succeeds	⚠️ see below

The only API that differs (hipMemRetainAllocationHandle) would false-positive on PyTorch expandable_segments memory — which is VMM-allocated but is host-accessible. Using it as a discriminator would regress performance by ~11× on those allocations (forcing hipMemcpy instead of direct dereference).

Conclusion: ROCm 7.14 has no API to query "is this pointer host-mapped?", so the correct fix is to grant host access at mapping time rather than trying to detect it at read time.

Performance (bench_discriminator.cpp, 20000 iterations)

  1. host direct deref (current fast path)          1128.9 ns
  2. discriminator + deref (hardened fast path)     1163.9 ns
     discriminator alone (retain/release)             65.5 ns
  3. hipMemcpy D2H 4 bytes (slow path)            12787.7 ns
  fast path vs slow path: 10.99x faster

The fix adds zero overhead — it simply makes the existing fast path valid for VMM memory.

Test

• Added test_cumem_host_read_scalar in tests/basic_correctness/test_mem.py (ROCm-only):
  • Allocates tensor in cumem VMM pool → .sum() → .item() verifies no segfault
  • Repeats after sleep/wake cycle to verify remap also grants host access
• verl megatron e2e with ROLLOUT_QUANTIZATION=fp8 + --enable_sleep_mode on MI300X (ROCm 7.14): passes end-to-end.
• Standalone reproducer confirms .item() no longer segfaults on VMM memory after the fix.

Additional notes

• isLargeBar = 1 confirmed on test machine (BAR0 = 256 GiB > physical VRAM 192 GiB)
• sleep/wake cycle unaffected (physical VRAM correctly released on sleep)