Skip to content

Revert "feat(comm): preserve all-reduce graph VAs across checkpoint restore" #3939

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
jiahanc wants to merge 1 commit into from
Closed

Revert "feat(comm): preserve all-reduce graph VAs across checkpoint restore" #3939

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 0 additions & 25 deletions docs/api/comm.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,31 +92,6 @@ Unified AllReduce Fusion API
TRTLLMAllReduceFusionWorkspace
MNNVLAllReduceFusionWorkspace

All-reduce workspaces backed by ``SymmDeviceMemory`` preserve their CUDA
virtual addresses across process checkpoint/restore. After quiescing all
work, release the physical handles and restore them with a fresh communication
backend before replaying a captured CUDA graph:

.. code-block:: python

workspace.checkpoint_prepare()
workspace.checkpoint_restore(comm_backend)

Both methods are collective. Every rank must call them in the same order, and
``comm_backend`` must reproduce the original rank and world size. Repeated
calls are no-ops after the workspace reaches the requested state. If an
exception occurs after detach or reattach begins, do not retry or reuse the
workspace; restart the affected rank. Workspaces backed by torch symmetric
memory do not support this lifecycle.

.. autosummary::
:toctree: ../generated

TRTLLMAllReduceFusionWorkspace.checkpoint_prepare
TRTLLMAllReduceFusionWorkspace.checkpoint_restore
MNNVLAllReduceFusionWorkspace.checkpoint_prepare
MNNVLAllReduceFusionWorkspace.checkpoint_restore

vLLM AllReduce
--------------

Expand Down
59 changes: 0 additions & 59 deletions flashinfer/comm/allreduce.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,6 @@

from .trtllm_ar import trtllm_allreduce_fusion
from .trtllm_ar import trtllm_create_ipc_workspace_for_all_reduce_fusion
from .trtllm_ar import _initialize_allreduce_fusion_protocol
from .trtllm_ar import check_trtllm_allreduce_fusion_workspace_metadata
from .trtllm_ar import trtllm_moe_allreduce_fusion
from .trtllm_ar import trtllm_moe_finalize_allreduce_fusion
Expand Down Expand Up @@ -179,59 +178,6 @@ def is_buffer_size_sufficient(
logger.warning("Workspace is insufficient for problem size. %s", e)
return False

@flashinfer_api
def checkpoint_prepare(self) -> None:
"""Detach physical backing; repeated successful calls are no-ops."""
if not self.mem_handles or not all(
isinstance(handle, SymmDeviceMemory) for handle in self.mem_handles
):
raise NotImplementedError(
"Stable-VA checkpointing is unavailable for workspaces backed "
"by torch symmetric memory"
)

mapped = [handle.mapped for handle in self.mem_handles]
if not any(mapped):
return
if not all(mapped):
raise RuntimeError("TRT-LLM symmetric-memory handle state is inconsistent")

for handle in self.mem_handles:
handle._unmap_and_release_handles()
# Do not return until every rank has released all workspace handles.
self.mem_handles[0].comm_backend.barrier()

@flashinfer_api
def checkpoint_restore(self, comm_backend: CommBackend) -> None:
"""Restore physical backing; repeated successful calls are no-ops."""
if not self.mem_handles or not all(
isinstance(handle, SymmDeviceMemory) for handle in self.mem_handles
):
raise NotImplementedError(
"Stable-VA checkpointing is unavailable for workspaces backed "
"by torch symmetric memory"
)

mapped = [handle.mapped for handle in self.mem_handles]
if all(mapped):
return
if any(mapped):
raise RuntimeError("TRT-LLM symmetric-memory handle state is inconsistent")
for handle in self.mem_handles:
handle._create_and_map_handles(comm_backend)

_initialize_allreduce_fusion_protocol(
ipc_handles=self.ipc_handles,
tp_rank=self.rank,
flag_size=self.metadata["flag_size"],
lamport_buffer_size=self.metadata["lamport_buffer_size"],
lamport_comm_size=self.metadata["lamport_comm_size"],
use_fp32_lamport=self.metadata["use_fp32_lamport"],
control_flag_ptr=self.metadata["control_flag_ptr"],
)
torch.cuda.synchronize()
comm_backend.barrier()

def destroy(self) -> None:
"""Destroy workspace and free resources."""
if getattr(self, "_destroyed", False):
Expand Down Expand Up @@ -730,11 +676,6 @@ def allreduce_fusion(
# Dispatch based on workspace type
if isinstance(workspace, TRTLLMAllReduceFusionWorkspace):
# TensorRT-LLM backend implementation
if any(
isinstance(handle, SymmDeviceMemory) and not handle.mapped
for handle in workspace.mem_handles
):
raise RuntimeError("TRT-LLM symmetric-memory handles are not attached")

# ---- MOE Reduction pattern ----
if pattern == AllReduceFusionPattern.kMoEReductionARResidualRMSNorm:
Expand Down
155 changes: 52 additions & 103 deletions flashinfer/comm/mnnvl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1004,7 +1004,6 @@ def __init__(
self.signal_pad_offset = 0
self.allocation_size = 0
self.comm_backend = comm_backend_for_handle_transfer or MPIBackend()
self._enable_multicast = enable_multicast

# CUDA memory handles and pointers
self.mc_ptr = 0 # CUdeviceptr mMcPtr
Expand All @@ -1016,24 +1015,22 @@ def __init__(
self.uc_handles: List[
int
] = [] # std::vector<CUmemGenericAllocationHandle> mUcHandles
self._mapped = False

# Signal pad constants
self.SIGNAL_PAD_ALIGNMENT = 16
self.SIGNAL_PAD_SIZE = SIGNAL_PAD_SIZE

# Check if device supports multicasting
if self._enable_multicast:
multicast_supported = checkCudaErrors(
cuda.cuDeviceGetAttribute(
cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED,
device_idx,
)
multicast_supported = checkCudaErrors(
cuda.cuDeviceGetAttribute(
cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED,
device_idx,
)
)
if multicast_supported == 0:
raise RuntimeError(
"[SymmDeviceMemory] Device does not support multicasting."
)
Comment on lines +1024 to 1033

@gemini-code-assist gemini-code-assist Bot Jul 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The multicast support check is now performed unconditionally. If enable_multicast is passed as False (for example, on devices that do not support multicasting), this will raise a RuntimeError and prevent SymmDeviceMemory from being initialized at all. We should wrap this check in if enable_multicast: and also store self._enable_multicast = enable_multicast so that helper methods can check it.

Suggested change
multicast_supported = checkCudaErrors(
cuda.cuDeviceGetAttribute(
cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED,
device_idx,
)
)
if multicast_supported == 0:
raise RuntimeError(
"[SymmDeviceMemory] Device does not support multicasting."
)
self._enable_multicast = enable_multicast
# Check if device supports multicasting
if enable_multicast:
multicast_supported = checkCudaErrors(
cuda.cuDeviceGetAttribute(
cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED,
device_idx,
)
)
if multicast_supported == 0:
raise RuntimeError(
"[SymmDeviceMemory] Device does not support multicasting."
)

Comment on lines 1023 to 1033

@coderabbitai coderabbitai Bot Jul 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Map the file structure first, then inspect the relevant ranges.
ast-grep outline flashinfer/comm/mnnvl.py --view expanded || true

echo
echo '--- lines 980-1085 ---'
sed -n '980,1085p' flashinfer/comm/mnnvl.py | cat -n

echo
echo '--- lines 1210-1265 ---'
sed -n '1210,1265p' flashinfer/comm/mnnvl.py | cat -n

echo
echo '--- search for enable_multicast and multicast-specific APIs ---'
rg -n "enable_multicast|cuMulticast|multicast_supported|_mc_granularity|_allocate_unicast_buffers|_get_allocation_prop" flashinfer/comm/mnnvl.py

Repository: flashinfer-ai/flashinfer

Length of output: 14662

🏁 Script executed:

#!/bin/bash
set -euo pipefail

ast-grep outline flashinfer/comm/mnnvl.py --view expanded || true

echo
echo '--- lines 980-1085 ---'
sed -n '980,1085p' flashinfer/comm/mnnvl.py | cat -n

echo
echo '--- lines 1210-1265 ---'
sed -n '1210,1265p' flashinfer/comm/mnnvl.py | cat -n

echo
echo '--- search for enable_multicast and multicast-specific APIs ---'
rg -n "enable_multicast|cuMulticast|multicast_supported|_mc_granularity|_allocate_unicast_buffers|_get_allocation_prop" flashinfer/comm/mnnvl.py

Repository: flashinfer-ai/flashinfer

Length of output: 14662

🏁 Script executed:

#!/bin/bash
set -euo pipefail

sed -n '1000,1055p' flashinfer/comm/mnnvl.py | cat -n
echo '---'
sed -n '1218,1258p' flashinfer/comm/mnnvl.py | cat -n

Repository: flashinfer-ai/flashinfer

Length of output: 4608

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo '--- lines 1178-1368 ---'
sed -n '1178,1368p' flashinfer/comm/mnnvl.py | cat -n

echo
echo '--- lines 1368-1392 ---'
sed -n '1368,1392p' flashinfer/comm/mnnvl.py | cat -n

Repository: flashinfer-ai/flashinfer

Length of output: 10067

Gate multicast-only setup on enable_multicast
enable_multicast=False still hits cuDeviceGetAttribute(CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED) and cuMulticastGetGranularity(), so construction fails on devices without multicast support instead of falling back to unicast-only mode. Move the support check and multicast allocation setup behind the flag, or split the unicast path entirely.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@flashinfer/comm/mnnvl.py` around lines 1023 - 1033, Gate the multicast
support check and all multicast-specific allocation setup in the constructor
around enable_multicast, including cuDeviceGetAttribute and
cuMulticastGetGranularity calls. When enable_multicast is false, skip these
operations entirely and preserve a unicast-only initialization path that does
not require multicast-capable devices.

if multicast_supported == 0:
raise RuntimeError(
"[SymmDeviceMemory] Device does not support multicasting."
)

# Calculate signal pad offset with alignment (matching C++ exactly)
self.signal_pad_offset = round_up(buf_size, self.SIGNAL_PAD_ALIGNMENT)
Expand All @@ -1044,8 +1041,16 @@ def __init__(
f"Signal pad offset: {self.signal_pad_offset}"
)

self._exchanger: Optional[HandleExchanger] = None
self._create_and_map_handles(self.comm_backend)
# Create handle exchanger
if is_mnnvl_fabric_supported(device_idx):
self._exchanger: HandleExchanger = FabricHandleExchanger(
self.comm_backend, self.group_rank, self.group_size
)
else:
self._exchanger = PosixFDHandleExchanger(
self.comm_backend, self.group_rank, self.group_size
)
self._alloc_mn_mcast_mem(buf_size, enable_multicast)

if allocate_signal_pads:
# Initialize signal pads
Expand All @@ -1063,7 +1068,7 @@ def __init__(
def __del__(self):
"""Destructor - cleanup allocated memory"""

if hasattr(self, "_exchanger") and self._exchanger is not None:
if hasattr(self, "_exchanger"):
self._exchanger.close()

# Skip cleanup during Python finalization to avoid segfaults
Expand Down Expand Up @@ -1121,8 +1126,6 @@ def __del__(self):
checkCudaErrors(cuda.cuMemRelease(self.mc_handle))
except Exception as e:
logger.warning("Destructor: Failed to release MC handle: %s", e)
elif hasattr(self, "mc_ptr") and self.mc_ptr:
checkCudaErrors(cuda.cuMemAddressFree(self.mc_ptr, self.allocation_size))

def get_signal_pad_ptrs_host(self) -> List[int]:
"""Get the raw array of signal pad pointers to all ranks (including self)"""
Expand Down Expand Up @@ -1172,65 +1175,20 @@ def get_usable_buffer_size(self) -> int:
"""Get the usable buffer size (excluding signal pad)"""
return self.allocation_size - self.SIGNAL_PAD_SIZE

@property
def mapped(self) -> bool:
return self._mapped

def _create_and_map_handles(self, comm: CommBackend) -> None:
"""Create physical backing and map it at the reserved addresses."""
# Create handle exchanger
if is_mnnvl_fabric_supported(self.device_idx):
self._exchanger = FabricHandleExchanger(
comm, self.group_rank, self.group_size
)
else:
self._exchanger = PosixFDHandleExchanger(
comm, self.group_rank, self.group_size
)

def _alloc_mn_mcast_mem(self, buf_size: int, enable_multicast: bool):
"""Allocate multi-node multicast memory using MNNVL"""
self._verify_cuda_context()

# Compute allocation size and get allocation properties
allocation_prop, mc_prop = self._get_allocation_prop(self.buf_size)
allocation_prop, mc_prop = self._get_allocation_prop(buf_size)

# Allocate, exchange, and map unicast buffers
self._allocate_unicast_buffers(allocation_prop)

# Setup multicast object, exchange handles, map and bind memory
if self._enable_multicast:
if enable_multicast:
self._setup_multicast(mc_prop)

self.comm_backend = comm
self._mapped = True

def _unmap_and_release_handles(self) -> None:
"""Unmap and release physical backing while retaining reserved addresses."""
# Drain local work, then align ranks before changing shared mappings.
cuda.cuCtxSynchronize()
self.comm_backend.barrier()

if self._enable_multicast:
checkCudaErrors(
cuda.cuMulticastUnbind(
self.mc_handle, self.device_idx, 0, self.allocation_size
)
)
checkCudaErrors(cuda.cuMemUnmap(self.mc_ptr, self.allocation_size))

for ptr in self.uc_ptrs:
checkCudaErrors(cuda.cuMemUnmap(ptr, self.allocation_size))

if self._enable_multicast:
checkCudaErrors(cuda.cuMemRelease(self.mc_handle))
self.mc_handle = 0
for handle in self.uc_handles:
checkCudaErrors(cuda.cuMemRelease(handle))

self._exchanger.close()
self.uc_handles = [0] * self.group_size
self._exchanger = None
self._mapped = False

def _verify_cuda_context(self):
"""Verify CUDA context is set to the correct device."""
try:
Expand Down Expand Up @@ -1268,23 +1226,20 @@ def _get_allocation_prop(self, buf_size: int):
buf_size + self.SIGNAL_PAD_SIZE, alloc_granularity
)

self._mc_granularity = alloc_granularity
mc_prop = None
if self._enable_multicast:
# Set up multicast properties
mc_prop = cuda.CUmulticastObjectProp()
mc_prop.numDevices = self.group_size
mc_prop.size = self.allocation_size
mc_prop.handleTypes = self._exchanger.handle_type

# Get multicast granularity and adjust allocation size
self._mc_granularity = checkCudaErrors(
cuda.cuMulticastGetGranularity(
mc_prop,
cuda.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_RECOMMENDED,
)
# Set up multicast properties
mc_prop = cuda.CUmulticastObjectProp()
mc_prop.numDevices = self.group_size
mc_prop.size = self.allocation_size
mc_prop.handleTypes = self._exchanger.handle_type

# Get multicast granularity and adjust allocation size
self._mc_granularity = checkCudaErrors(
cuda.cuMulticastGetGranularity(
mc_prop,
cuda.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_RECOMMENDED,
)
self.allocation_size = round_up(self.allocation_size, self._mc_granularity)
)
self.allocation_size = round_up(self.allocation_size, self._mc_granularity)
Comment on lines +1229 to +1242

@gemini-code-assist gemini-code-assist Bot Jul 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

In _get_allocation_prop, the multicast properties are set up and cuMulticastGetGranularity is called unconditionally. If enable_multicast is False (and multicast is not supported by the device), these calls will fail. We should wrap the multicast setup and granularity query in if self._enable_multicast: to avoid executing them when multicast is disabled.

Suggested change
# Set up multicast properties
mc_prop = cuda.CUmulticastObjectProp()
mc_prop.numDevices = self.group_size
mc_prop.size = self.allocation_size
mc_prop.handleTypes = self._exchanger.handle_type
# Get multicast granularity and adjust allocation size
self._mc_granularity = checkCudaErrors(
cuda.cuMulticastGetGranularity(
mc_prop,
cuda.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_RECOMMENDED,
)
self.allocation_size = round_up(self.allocation_size, self._mc_granularity)
)
self.allocation_size = round_up(self.allocation_size, self._mc_granularity)
self._mc_granularity = alloc_granularity
mc_prop = None
if self._enable_multicast:
# Set up multicast properties
mc_prop = cuda.CUmulticastObjectProp()
mc_prop.numDevices = self.group_size
mc_prop.size = self.allocation_size
mc_prop.handleTypes = self._exchanger.handle_type
# Get multicast granularity and adjust allocation size
self._mc_granularity = checkCudaErrors(
cuda.cuMulticastGetGranularity(
mc_prop,
cuda.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_RECOMMENDED,
)
)
self.allocation_size = round_up(self.allocation_size, self._mc_granularity)


return allocation_prop, mc_prop

Expand Down Expand Up @@ -1320,24 +1275,21 @@ def _allocate_unicast_buffers(self, allocation_prop):
self._exchanger.handle_type,
)
)
self._exchanger.cleanup(all_shareable_uc_handles[p])
self._exchanger.cleanup(local_shareable_uc_handle)
self._exchanger.cleanup(all_shareable_uc_handles[p])

# Reserve address space for UC pointers
if not self.uc_ptrs:
self.uc_ptrs = [0] * self.group_size
total_uc_size = self.allocation_size * self.group_size
self.total_uc_size = total_uc_size
uc_base_ptr = checkCudaErrors(
cuda.cuMemAddressReserve(total_uc_size, self._mc_granularity, 0, 0)
)
self.uc_base_ptr = uc_base_ptr
for i in range(self.group_size):
offset = self.allocation_size * i
self.uc_ptrs[i] = int(uc_base_ptr) + offset
self.uc_ptrs = [0] * self.group_size
total_uc_size = self.allocation_size * self.group_size
self.total_uc_size = total_uc_size
uc_base_ptr = checkCudaErrors(
cuda.cuMemAddressReserve(total_uc_size, self._mc_granularity, 0, 0)
)
self.uc_base_ptr = uc_base_ptr

# Map UC memory
for i in range(self.group_size):
offset = self.allocation_size * i
self.uc_ptrs[i] = int(uc_base_ptr) + offset
checkCudaErrors(
cuda.cuMemMap(
self.uc_ptrs[i], self.allocation_size, 0, self.uc_handles[i], 0
Expand All @@ -1347,7 +1299,7 @@ def _allocate_unicast_buffers(self, allocation_prop):
# Set memory access permissions for UC
access_desc = self._get_mem_access_desc()
checkCudaErrors(
cuda.cuMemSetAccess(self.uc_base_ptr, self.total_uc_size, [access_desc], 1)
cuda.cuMemSetAccess(uc_base_ptr, total_uc_size, [access_desc], 1)
)

def _setup_multicast(self, mc_prop):
Expand Down Expand Up @@ -1377,18 +1329,15 @@ def _setup_multicast(self, mc_prop):
self._exchanger.handle_type,
)
)
self._exchanger.cleanup(shareable_mc_handle)
self._exchanger.cleanup(shareable_mc_handle)

# Add device to multicast
checkCudaErrors(cuda.cuMulticastAddDevice(self.mc_handle, self.device_idx))

# Reserve and map MC pointer
if not self.mc_ptr:
self.mc_ptr = checkCudaErrors(
cuda.cuMemAddressReserve(
self.allocation_size, self._mc_granularity, 0, 0
)
)
self.mc_ptr = checkCudaErrors(
cuda.cuMemAddressReserve(self.allocation_size, self._mc_granularity, 0, 0)
)
checkCudaErrors(
cuda.cuMemMap(self.mc_ptr, self.allocation_size, 0, self.mc_handle, 0)
)
Expand Down
Loading
Loading