[ROCm][VMM] Selective copy-vs-remap for small command-buffer slices + auto threshold

[phambinhfin]

Dependency (read first)

This PR stacks on top of the ROCm VMM allocator work (branch phambinh/rocm-vmm-allocator) and must merge after it. The base of this PR is intentionally set to that branch, not main, so the diff here is only this change (xla/service/gpu/gpu_executable.cc). GitHub will auto-retarget to main once the VMM allocator PR lands.

It is entirely ROCm-only and does not affect NVIDIA. The whole code path lives inside ExecuteThunksWithVaRemapping, which is gated on:

dynamic_cast<se::DeviceAddressVmmAllocator*>(memory_allocator) != nullptr

On NVIDIA the active allocator is the standard BFC/stream-executor allocator, so the cast is null, use_command_buffer_va_remapping is false, the function is never entered, and the XLA_VMM_TMP_COPY_THRESHOLD env var has no effect. NVIDIA code paths, defaults, and behavior are unchanged.

Optimization

With the VMM command-buffer VA-remapping path, every command-buffer slice whose physical allocation churns each step is re-mapped via hipMemUnmap / hipMemMap / SetAccess. For the tiny scalar / scale / metric buffers that churn every step, that driver round-trip dominates and serializes across peer devices. Large tensors, by contrast, keep stable addresses across steps, so their remap is already skipped (free).

This change keeps those small churning slices in a stable shadow buffer (a fixed address baked into the command buffer once) and refreshes them with a cheap device-to-device copy each step (copy-in before exec, copy-back after, so output scalars like loss/metrics propagate correctly). Large slices stay on the remap path.

The size cutoff is chosen automatically — no manual tuning, no searching — via XLA_VMM_TMP_COPY_THRESHOLD:

value	behavior
0	disabled: remap every slice (original behavior)
1	auto: derive the cutoff from device HBM bandwidth + capacity and currently-free HBM (DeviceMemoryUsage)
>1	force that exact byte cutoff (explicit override / sweeps)

The auto selector admits the smallest slices first while bounding (a) per-slice copy time, (b) aggregate per-step copy time, and (c) the shadow footprint at a fraction of free HBM — so the shadow pool is always guaranteed to fit and never risks OOM. Default (unset) is auto.

Benchmarks (MI300X x8, llama @ 8 layers / 30 steps, median of steps ≥3)

auto + collective leg (coll_vmm), disabled vs auto — all models

Model	disabled (=0)	auto (=1)	saved	speedup	loss (both)	unmap
llama2_7b	133.5 ms	121.0 ms	12.5 ms	9.4%	0.003	0
llama3_8b	168.0 ms	156.5 ms	11.5 ms	6.8%	0.001	0
mixtral_8x1b	690.5 ms	677.5 ms	13.0 ms	1.9%	1.548	0
llama3.1-8b	713.5 ms	704.5 ms	9.0 ms	1.3%	0.001	0
gemma2-9b	1319.5 ms	1313.0 ms	6.5 ms	0.5%	0.003	0

The saving is a roughly fixed per-step cost (~6-13 ms; the work removed is the remap of the small churning slices, whose count is set by command-buffer structure, not model size). It therefore shows up as a larger percentage on short-step models and a smaller one on heavy models, but it is always a net win and never regresses. auto lands on the same optimum as a hand-tuned fixed 4 KB cutoff (llama2_7b: auto 121 ms == fixed-4K 120 ms).

vmm leg (no collectives), disabled vs auto — same pattern

Model	disabled (=0)	auto (=1)	saved	speedup
llama2_7b	138.5 ms	126.0 ms	12.5 ms	9.0%
llama3_8b	168.5 ms	158.0 ms	10.5 ms	6.2%
mixtral_8x1b / llama3.1-8b / gemma2-9b	running	running		will update

Correctness across every run: identical loss vs the disabled baseline and zero unmap errors.

Notes

• Marked [EXPERIMENT] — behind an env knob, default auto, ROCm-gated.
• vmm-leg numbers for the 3 heavier models will be appended when the in-flight sweep completes.

[phambinhfin]

Superseded by #926 (same change, head branch renamed to phambinh/balance_vmm).