@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b5b61c622c

[Core][Distributed] Add process-checkpoint lifecycle hooks for communicators (starting with Flashinfer) (#46877)

Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com>

Schwinn Saereesitthipitak committed Jul 26, 2026 at 11:47 UTC b5b61c622c941426cf78a83672e878b5b829fe90

10 files changed +185 -4

tests/distributed/test_comm_ops.py

+40

@@ -7,6 +7,7 @@ Run `pytest tests/distributed/test_comm_ops.py`.

8 from collections.abc import Callable

9 from typing import Any

10 +from unittest.mock import Mock

12 import pytest

13 import ray

@@ -19,6 +20,8 @@ from vllm.distributed import (

20 tensor_model_parallel_all_reduce,

21 tensor_model_parallel_reduce_scatter,

22 )

23 +from vllm.distributed.device_communicators import flashinfer_all_reduce

24 +from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator

25 from vllm.distributed.parallel_state import GroupCoordinator, TensorMetadata

26 from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors

@@ -278,6 +281,43 @@ def test_irecv_tensor_dict_send_allgather_postprocess_binds_keys(

281 torch.testing.assert_close(td["b"], torch.ones(4, dtype=torch.int32))

282

283

284 +@pytest.mark.parametrize("aliased", [False, True])

285 +def test_cuda_communicator_checkpoints_flashinfer_workspaces(

286 + monkeypatch: pytest.MonkeyPatch,

287 + aliased: bool,

288 +) -> None:

289 + group = object()

290 + normal_workspace = Mock()

291 + quant_workspace = normal_workspace if aliased else Mock()

292 + unique_workspaces = (

293 + [normal_workspace] if aliased else [normal_workspace, quant_workspace]

294 + )

295 +

296 + monkeypatch.setattr(flashinfer_all_reduce, "_fi_ar_workspace", normal_workspace)

297 + monkeypatch.setattr(

298 + flashinfer_all_reduce, "_fi_ar_quant_workspace", quant_workspace

299 + )

300 + monkeypatch.setattr(

301 + flashinfer_all_reduce,

302 + "_fi_ar_workspace_groups",

303 + {id(workspace): group for workspace in unique_workspaces},

304 + )

305 + monkeypatch.setattr(

306 + flashinfer_all_reduce, "TorchDistBackend", lambda group: group, raising=False

307 + )

308 +

309 + communicator = CudaCommunicator.__new__(CudaCommunicator)

310 + communicator.cpu_group = group

311 + communicator.fi_ar_comm = None

312 + communicator.all2all_manager = None

313 + communicator.checkpoint_prepare()

314 + communicator.checkpoint_restore()

315 +

316 + for workspace in unique_workspaces:

317 + workspace.checkpoint_prepare.assert_called_once_with()

318 + workspace.checkpoint_restore.assert_called_once_with(group)

319 +

320 +

321 def test_async_intermediate_tensors_lazy_wait() -> None:

322 work = _DummyWork()

323 post_calls = {"n": 0}

vllm/compilation/passes/fusion/allreduce_rms_fusion.py

+2 -2

@@ -225,7 +225,7 @@ if flashinfer_comm is not None:

225 max_token_num=max_token_num,

226 hidden_dim=hidden_size,

227 dtype=allreduce_in.dtype,

228 - group=get_tp_group().device_group,

228 + group=get_tp_group().cpu_group,

229 )

230 assert workspace is not None, (

231 "Flashinfer allreduce workspace must be initialized when using flashinfer"

@@ -996,7 +996,7 @@ class AllReduceFusionPass(VllmPatternMatcherPass):

996 )

997 return

998 self.hidden_dim = config.model_config.get_hidden_size()

999 - self.group = get_tp_group().device_group

999 + self.group = get_tp_group().cpu_group

1000 rank = get_tensor_model_parallel_rank()

1001 if flashinfer_comm is None:

1002 logger.warning(

vllm/distributed/device_communicators/all2all.py

+14

@@ -885,6 +885,20 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):

885 self.mapping = None

886 self.initialized = False

887

888 + def checkpoint_prepare(self) -> None:

889 + if self.initialized:

890 + assert self.moe_alltoall is not None

891 + self.moe_alltoall.checkpoint_prepare()

892 +

893 + def checkpoint_restore(self) -> None:

894 + if self.initialized:

895 + assert self.moe_alltoall is not None

896 + from vllm.distributed.device_communicators.mnnvl_compat import (

897 + CustomCommunicator,

898 + )

899 +

900 + self.moe_alltoall.checkpoint_restore(CustomCommunicator(self.cpu_group))

901 +

902

903 class MoriAll2AllManager(All2AllManagerBase):

904 def __init__(self, cpu_group, all2all_backend: str):

vllm/distributed/device_communicators/base_device_communicator.py

+21

@@ -7,8 +7,11 @@ import torch

7 import torch.distributed as dist

8 from torch.distributed import ProcessGroup

10 +from vllm.logger import init_logger

11 from vllm.utils import is_moe_layer

13 +logger = init_logger(__name__)

14 +

16 class Cache:

17 def __init__(self):

@@ -137,6 +140,18 @@ class All2AllManagerBase:

140 def max_sms_used(self) -> int | None:

141 return None # None means it could use the whole GPU

142

143 + def checkpoint_prepare(self) -> None:

144 + logger.warning_once(

145 + "%s.checkpoint_prepare is not implemented; skipping.",

146 + type(self).__name__,

147 + )

148 +

149 + def checkpoint_restore(self) -> None:

150 + logger.warning_once(

151 + "%s.checkpoint_restore is not implemented; skipping.",

152 + type(self).__name__,

153 + )

154 +

155 def combine(self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False):

156 raise NotImplementedError

157

@@ -216,6 +231,12 @@ class DeviceCommunicatorBase:

231 dist.all_reduce(input_, group=self.device_group)

232 return input_

233

234 + def checkpoint_prepare(self) -> None:

235 + """Prepare reclaimable communicator state for checkpoint (default: no-op)."""

236 +

237 + def checkpoint_restore(self) -> None:

238 + """Restore communicator state after checkpoint (default: no-op)."""

239 +

240 def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:

241 if dim < 0:

242 # Convert negative dim to positive.

vllm/distributed/device_communicators/cuda_communicator.py

+16

@@ -571,6 +571,22 @@ class CudaCommunicator(DeviceCommunicatorBase):

571 self.all2all_manager.destroy()

572 self.all2all_manager = None # type: ignore[assignment]

573

574 + def checkpoint_prepare(self) -> None:

575 + # Only FlashInfer all-reduce and FlashInfer all2all are supported for now.

576 + from .flashinfer_all_reduce import checkpoint_prepare_fi_ar_workspaces

577 +

578 + checkpoint_prepare_fi_ar_workspaces(self.cpu_group)

579 + if self.all2all_manager is not None:

580 + self.all2all_manager.checkpoint_prepare()

581 +

582 + def checkpoint_restore(self) -> None:

583 + # Only FlashInfer all-reduce and FlashInfer all2all are supported for now.

584 + from .flashinfer_all_reduce import checkpoint_restore_fi_ar_workspaces

585 +

586 + checkpoint_restore_fi_ar_workspaces(self.cpu_group)

587 + if self.all2all_manager is not None:

588 + self.all2all_manager.checkpoint_restore()

589 +

590 def all_gatherv(

591 self,

592 input_: torch.Tensor | list[torch.Tensor],

vllm/distributed/device_communicators/flashinfer_all_reduce.py

+41 -1

@@ -6,6 +6,7 @@ import atexit

6 import os

7 import random

8 import threading

9 +from typing import Any

11 import torch

12 import torch.distributed as dist

@@ -39,6 +40,7 @@ _fi_ar_workspace = None

40 # allreduce backend or a fallback backend when the primary workspace is not

41 # available on the current topology.

42 _fi_ar_quant_workspace = None

43 +_fi_ar_workspace_groups: dict[int, ProcessGroup] = {}

46 def _create_workspace(

@@ -81,6 +83,14 @@ def _create_workspace(

83 return None

84 finally:

85 random.setstate(rng_state)

86 + workspace_id = id(workspace)

87 + workspace_group = _fi_ar_workspace_groups.get(workspace_id)

88 + if workspace_group is not None and workspace_group is not group:

89 + raise RuntimeError(

90 + "FlashInfer returned an all-reduce workspace already associated "

91 + "with a different process group"

92 + )

93 + _fi_ar_workspace_groups[workspace_id] = group

94 logger.debug(

95 "Initialized FlashInfer All Reduce workspace: backend=%s, "

96 "world_size=%d, rank=%d, max_token_num=%d, hidden_dim=%d, dtype=%s",

@@ -111,7 +121,7 @@ def _resolve_fi_ar_backend() -> tuple[str, bool]:

121 # Default to mnnvl for both single- and multi-node setups. The mnnvl

122 # cudagraph hang that previously forced single-node to trtllm

123 # (https://github.com/vllm-project/vllm/issues/35772) was fixed upstream in

114 - # FlashInfer (>= 0.6.12, vLLM pins 0.6.13), so mnnvl is safe here. trtllm

124 + # FlashInfer (>= 0.6.12, vLLM pins 0.6.15), so mnnvl is safe here. trtllm

125 # does not support multi-node allreduce, so mnnvl is required there anyway.

126 # mnnvl needs NVSwitch multicast; on single-node topologies without it,

127 # fall back to trtllm so fused allreduce stays enabled.

@@ -268,6 +278,36 @@ def destroy_fi_ar_workspace():

278 _fi_ar_quant_workspace.destroy()

279

280 _fi_ar_workspace = _fi_ar_quant_workspace = None

281 + _fi_ar_workspace_groups.clear()

282 +

283 +

284 +def _fi_ar_workspaces_for_group(group: ProcessGroup) -> list[Any]:

285 + workspaces = [_fi_ar_workspace]

286 + if _fi_ar_quant_workspace is not _fi_ar_workspace:

287 + workspaces.append(_fi_ar_quant_workspace)

288 +

289 + group_workspaces = []

290 + for workspace in workspaces:

291 + if workspace is None:

292 + continue

293 + workspace_group = _fi_ar_workspace_groups.get(id(workspace))

294 + if workspace_group is None:

295 + raise RuntimeError(

296 + "FlashInfer all-reduce workspace process group was not retained"

297 + )

298 + if workspace_group is group:

299 + group_workspaces.append(workspace)

300 + return group_workspaces

301 +

302 +

303 +def checkpoint_prepare_fi_ar_workspaces(group: ProcessGroup) -> None:

304 + for workspace in _fi_ar_workspaces_for_group(group):

305 + workspace.checkpoint_prepare()

306 +

307 +

308 +def checkpoint_restore_fi_ar_workspaces(group: ProcessGroup) -> None:

309 + for workspace in _fi_ar_workspaces_for_group(group):

310 + workspace.checkpoint_restore(TorchDistBackend(group=group))

311

312

313 atexit.register(destroy_fi_ar_workspace)

vllm/distributed/parallel_state.py

+36

@@ -127,6 +127,28 @@ def _register_group(group: "GroupCoordinator") -> None:

127 _groups[group.unique_name] = weakref.ref(group)

128

129

130 +def _apply_to_device_comms(

131 + action: Callable[[DeviceCommunicatorBase], None],

132 +) -> None:

133 + """Apply ``action`` to every group's device communicator.

134 +

135 + Walks the registered parallel groups and skips those without a device

136 + communicator (absent at ``world_size == 1``).

137 + """

138 + comms = []

139 + for group_ref in _groups.values():

140 + group = group_ref()

141 + if group is None:

142 + continue

143 + dc = group.device_communicator

144 + if dc is None:

145 + continue

146 + comms.append(dc)

147 +

148 + for dc in comms:

149 + action(dc)

150 +

151 +

152 def all_reduce(tensor: torch.Tensor, group_name: str) -> torch.Tensor:

153 assert group_name in _groups, f"Group {group_name} is not found."

154 group = _groups[group_name]()

@@ -2028,6 +2050,20 @@ def prepare_communication_buffer_for_model(model: torch.nn.Module):

2050 _EPLB.prepare_communication_buffer_for_model(model)

2051

2052

2053 +def checkpoint_prepare_distributed_state() -> None:

2054 + """Prepare every device communicator for a process checkpoint."""

2055 + torch.accelerator.synchronize()

2056 + _apply_to_device_comms(lambda comm: comm.checkpoint_prepare())

2057 + torch.accelerator.synchronize()

2058 +

2059 +

2060 +def checkpoint_restore_distributed_state() -> None:

2061 + """Restore every device communicator after a process checkpoint."""

2062 + torch.accelerator.synchronize()

2063 + _apply_to_device_comms(lambda comm: comm.checkpoint_restore())

2064 + torch.accelerator.synchronize()

2065 +

2066 +

2067 def model_parallel_is_initialized():

2068 """Check if tensor and pipeline parallel groups are initialized."""

2069 return _TP is not None and _PP is not None

vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py

+1 -1

@@ -93,7 +93,7 @@ def _can_use_flashinfer(hidden_states: torch.Tensor, tp_size: int) -> tuple[bool

93 max_token_num=max_token_num,

94 hidden_dim=hidden_size,

95 dtype=hidden_states.dtype,

96 - group=get_tp_group().device_group,

96 + group=get_tp_group().cpu_group,

97 )

98 if workspace is None:

99 return False, 0

vllm/v1/engine/async_llm.py

@@ -943,6 +943,12 @@ class AsyncLLM(EngineClient):

943 if self.logger_manager is not None:

944 self.logger_manager.record_sleep_state(0, 0)

945

946 + async def checkpoint_prepare(self) -> None:

947 + await self.collective_rpc("checkpoint_prepare")

948 +

949 + async def checkpoint_restore(self) -> None:

950 + await self.collective_rpc("checkpoint_restore")

951 +

952 async def is_sleeping(self) -> bool:

953 return await self.engine_core.is_sleeping_async()

954

vllm/v1/worker/gpu_worker.py

@@ -41,6 +41,8 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (

41 )

42 from vllm.distributed.parallel_state import (

43 Handle,

44 + checkpoint_prepare_distributed_state,

45 + checkpoint_restore_distributed_state,

46 get_pp_group,

47 get_tp_group,

48 )

@@ -244,6 +246,12 @@ class Worker(WorkerBase):

246 if tags is None or "kv_cache" in tags:

247 self.model_runner.post_kv_cache_wake_up()

248

249 + def checkpoint_prepare(self) -> None:

250 + checkpoint_prepare_distributed_state()

251 +

252 + def checkpoint_restore(self) -> None:

253 + checkpoint_restore_distributed_state()

254 +

255 def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:

256 if (

257 current_platform.is_cuda_alike()