vmm support for cuda graph offload

[Xdydy]

Description

Add a fixed-VA CUDA VMM activation slot so a Graph-captured tensor address can keep its virtual address while physical backing is released and remapped outside Graph replay. This is the TE primitive behind local whole-layer CUDA Graph + fine-grained CPU activation offload.

Fixes # (issue)

Local CUDA Graph and CPU offload currently conflict on pointer semantics:

Graph replay  requires: captured tensor addresses never change
CPU offload   requires: physical device memory is released after forward

A naive empty() reload allocates a new pointer. Capture recorded src.data_ptr(), so replay either reads stale storage or silently aliases an allocator reuse of that VA. Related: NVIDIA/Megatron-LM#3697.

This PR does not implement Megatron scheduling. It exposes the VMM slot, resident workers, and deferred H2D so a caller (Megatron fine_grained_activation_offload) can:

1. Reserve a VA and map physical backing before capture.
2. Capture / replay Graphs against slot.tensor.data_ptr() (stable for the run).
3. After forward D2H, unmap/release physical pages without freeing the VA.
4. Before backward, map new backing at the same VA and H2D.
5. Optionally prefetch remap while the current Graph runs, then launch H2D at replay time.

CUDA driver names stay CUDA-shaped in source. On MUSA they map at compile time through musify.h (cuMemMap → muMemMap, cudaMemcpyAsync → musaMemcpyAsync, …). Python keeps CUDA names (CUDAActivationVMMAllocation, cuda_stream); a MUSA runtime still works when torch.cuda is aliased, and device.type in ("cuda", "musa") covers TE tests that import this module without musa_patch.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Add transformer_engine/pytorch/csrc/extensions/vmm_activation.cpp and bind it from pybind.cpp / extensions.h.
• Add Python wrapper transformer_engine/pytorch/vmm_activation.py.
• Add tests/pytorch/test_vmm_activation.py (async reload / host-callback lifetime).
• Add musify.h VMM aliases used by this file (cuMemAddressReserve, cuMemCreate, cuMemMap, cuMemSetAccess, cuMemUnmap, cuMemRelease, cuMemGetInfo, …).

Not in this PR’s intended review surface: GEMM / DPA / cpu_offload.py Megatron handler changes that happen to sit on the same working branch.

---

Why a TE primitive (not a Megatron allocator wrapper)

Graph capture records a device pointer. CPU offload must drop physical pages. Those two facts only compose if the Graph sees a VA that outlives the pages mapped under it:

Naive:  act = torch.empty(...)     # address comes from the caching allocator
VMM:    slot = VMMActivationSlot(bytes, device)
        act  = slot.tensor(shape, stride, dtype)   # act.data_ptr() == reserved VA

cuMemAddressReserve  →  fixed VA (kept for the run)
cuMemCreate          →  physical allocation handle
cuMemMap             →  map handle at that VA
cuMemSetAccess       →  RW access for the device

offload path (after D2H):
  cuMemUnmap + cuMemRelease     →  pages gone, VA still reserved

reload path (before backward):
  cuMemCreate + cuMemMap + cuMemSetAccess at the same VA
  cudaMemcpyAsync(host → VA)

The captured Graph never sees a new pointer. Physical handles may change; the VA must not.

CUDAActivationVMMAllocation constructs the slot, wraps a non-owning from_blob view, and asserts tensor.data_ptr() == info["address"] after construct / remap / adopt. Using that view while the slot is unmapped is a hard page fault — that is intentional; it turns the silent allocator-reuse bug into a crash.

---

Slot lifecycle

construct
  cuMemAddressReserve + first cuMemCreate/Map/SetAccess
  tensor view at the reserved VA

forward offload (caller D2H, then)
  release_hooks_after(slots, d2h_stream)
    host func on d2h_stream after the D2H burst
    ReleaseWorker: cuMemUnmap + cuMemRelease
    VA reservation kept

backward restore
  remap_only_slot_after(slot, pinned_host, h2d_stream)   # remap now, no H2D
  … later, at Graph replay …
  launch_remap_slot_h2d(context, slot, h2d_stream)       # H2D now
  wait_remap_slot_on_stream(context, slot, compute_stream)
  adopt_async_remap()                                    # slot owns the new handle

close
  drain in-flight remap/release, unmap, release, cuMemAddressFree

Synchronous unmap_and_release() / create_and_remap() remain for tests and fallback. Training uses the async path.

Address stability check after every remap / adopt:

info["address"] == original_va and tensor.data_ptr() == original_va

---

Python API

transformer_engine.pytorch.vmm_activation:

Call	Role
CUDAActivationVMMAllocation(shape, stride, dtype, device)	reserve VA, map first backing, non-owning tensor view
release_hooks_after(allocs, stream)	one host func after D2H; worker unmaps the batch
remap_hooks_after(allocs, stream)	remap only, no copy
remap_and_copy_after / remap_and_copy_slot_after	worker remaps and submits H2D
remap_only_slot_after	worker remaps; H2D deferred
launch_remap_slot_h2d(ctx, i, stream)	wait remapped, submit H2D on stream
wait_remap_slot_on_stream(ctx, i, stream)	wait event_recorded, stream waits on slot event
wait_until_remap_slot_submitted	host wait until that slot’s H2D event is recorded
wait_remap_copy_on_stream / enqueue_remap_copy_wait	whole-context copy event (non-deferred path)
vmm_driver_memory_info()	cuMemGetInfo free/total (caching-allocator stats do not cover cuMemCreate pages)
vmm_enable_trace / vmm_initialize_workers / vmm_set_serial_driver_workers	diagnostics and worker policy

Stream handles: stream.cuda_stream, with musa_stream fallback.

Safety checks on remap submit: unique slots, reserved VA, no in-flight remap, mapped slot must already have a pending release, host tensor pinned / offset 0 / large enough. Adopted mappings are never unmapped by context cleanup.

Debug: MEGATRON_VMM_REMAP_DEBUG=1 prints per-call remap/release stages (cuMemUnmap, cuMemSetAccess, remap wait, H2D launch wait).

---

Files

transformer_engine/pytorch/csrc/extensions/vmm_activation.cpp
transformer_engine/pytorch/csrc/extensions.h
transformer_engine/pytorch/csrc/extensions/pybind.cpp
transformer_engine/pytorch/vmm_activation.py
transformer_engine/common/include/transformer_engine/musify.h   # VMM aliases
tests/pytorch/test_vmm_activation.py

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

Suggested coverage:

• Address stability: data_ptr() unchanged across unmap/remap and deferred H2D.
• release_hooks_after + remap_and_copy_after repeated reload (test_repeated_async_reload_releases_context_off_host_callback): no host-callback deadlock, values match, context error==0.
• Deferred path: remap_only_slot_after does not submit H2D; launch_remap_slot_h2d does; compute stream waits on the slot event.
• Negative: using the tensor view after unmap and before remap must fault, not silently alias a caching-allocator reuse of the captured VA.
• cuMemGetInfo (not PyTorch memory_allocated) when claiming physical pages returned.

[greptile-apps]

Confidence Score: 0/5

This PR is not safe to merge because it breaks supported CUDA/core/JAX builds and imports, removes still-used native APIs, forces unsupported attention paths, and has a multi-device VMM synchronization defect.

Findings

1.  MUSA Replaces Supported Builds ▶
2.  Root Import Requires MUSA ▶
3.  Import Globally Rewrites PyTorch ▶
4.  Active Extension Bindings Removed ▶
5.  FlashAttention Is Always Forced ▶
6.  FlashAttention Detection Is Fabricated ▶
7.  Events Use Ambient Device ▶

Summary

This PR adds a fixed-address VMM activation allocation with asynchronous release, remapping, deferred H2D copies, Python wrappers, and one repeated-reload test. It also contains a repository-wide MUSA port and unrelated attention, offload, GEMM, permutation, build, and package-initialization changes that materially alter existing CUDA behavior.

• Adds resident release/remap workers and fixed-VA activation slots.
• Adds Python APIs for synchronous and asynchronous remap/reload sequencing.
• Replaces the common CUDA build and runtime initialization with unconditional MUSA behavior.
• Removes active attention and KV-cache extension bindings and overrides backend selection.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  C[Captured graph uses fixed VA] --> D[D2H activation copy]
  D --> R[Release host callback]
  R --> RW[Release worker]
  RW --> U[Unmap and release physical pages]
  U --> MW[Remap worker]
  MW --> M[Map new pages at same VA]
  M --> H[Deferred or immediate H2D]
  H --> E[Record completion event]
  E --> W[Compute stream waits]
  W --> G[Graph replay reads restored activation]

Loading

_{Reviews (1) · Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..."}

[greptile-apps]

MUSA Replaces Supported Builds

The common build now always requires MUSA, MUSAToolkit, and MCCL from a hard-coded /usr/local/musa installation. Because setup.py adds this common extension for core-only, CUDA PyTorch, and JAX installations, supported non-MUSA builds fail during CMake configuration instead of using the existing CUDA build path.

Knowledge Base Used: Build, extensions, and packaging

[greptile-apps]

Root Import Requires MUSA

Root initialization now imports torch_musa before the existing optional-framework error handling. In core-only, JAX-only, or ordinary CUDA installations without torch_musa, import transformer_engine therefore fails even though MUSA is not a common runtime dependency.

Knowledge Base Used:

• PyTorch runtime and public API
• Restore Eager Framework Imports

[greptile-apps]

Import Globally Rewrites PyTorch

Importing Transformer Engine now replaces process-wide PyTorch CUDA, tensor-construction, distributed, device, and autocast APIs with MUSA implementations and immediately initializes a MUSA device. As a result, unrelated PyTorch code in the same process receives altered CUDA behavior merely because the root package was imported.

Knowledge Base Used: PyTorch runtime and public API

[greptile-apps]

Active Extension Bindings Removed

These bindings are removed while active Python paths still call them. KV-cache inference calls tex.copy_to_kv_cache, THD attention calls the conversion functions, and interleaved FlashAttention calls the preparation functions. Those paths now raise AttributeError instead of running the required native operation.

Knowledge Base Used: PyTorch runtime and public API

[greptile-apps]

FlashAttention Is Always Forced

Setting use_flash_attention after backend selection discards all capability and configuration checks. When FlashAttention is missing, disabled, unsupported for the inputs, or intentionally superseded by fused or unfused attention, execution still enters self.flash_attention instead of the selected fallback, causing unsupported execution or a runtime failure.

Knowledge Base Used: PyTorch runtime and public API

[greptile-apps]

FlashAttention Detection Is Fabricated

Hard-coding the detected FlashAttention version to 2.5.0 can mark it installed even when the package is absent or a different version is present. On a compatible device, initialization then imports flash_attn.flash_attn_interface unconditionally, so an environment that should fall back cleanly can fail during import or call APIs incompatible with its actual installation.

Knowledge Base Used: PyTorch runtime and public API

[greptile-apps]

Events Use Ambient Device

In a multi-device process where the current device, slot device, or stream device differ, these completion events are created on the calling thread's ambient device before slot devices are examined. The worker later switches to each slot's device and records the events on the supplied stream, causing cross-context event recording or waiting to fail and leaving the activation reload unsynchronized. Events must be created for the slot device, and mixed-device batches or mismatched streams should be rejected.