@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 245888ff77

[Feature] Detect all2all peer fault with fault tolerance backend and prevent corrupted output (#43637)

Signed-off-by: fangyuchu <fangyuchu@qq.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

fangyuchu committed Jul 1, 2026 at 00:00 UTC 245888ff77d4754a76bc899727fcf4290733dfd4

5 files changed +82 -4

vllm/distributed/device_communicators/all2all.py

+46

@@ -261,8 +261,13 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):

261 All2All communication based on DeepEP Low-Latency kernels.

262 """

263

264 + _buffer: Any = None

265 + _mask: torch.Tensor | None = None

266 + _last_mask: torch.Tensor | None = None

267 +

268 def __init__(self, cpu_group, tcp_store_group=None):

269 super().__init__(cpu_group, tcp_store_group)

270 + self.support_fault_tolerance = False # TODO: set to True when FT is supported.

271

272 def _make_all2all_kwargs(

273 self,

@@ -304,6 +309,7 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):

309 allow_nvlink_for_low_latency_mode=True,

310 allow_mnnvl=envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL,

311 explicitly_destroy=True,

312 + enable_shrink=self.support_fault_tolerance,

313 )

314 return kwargs

315

@@ -319,12 +325,30 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):

325 handle: deep_ep.Buffer = self.handle_cache.get_or_create(

326 buffer_kwargs, deep_ep.Buffer

327 )

328 + DeepEPLLAll2AllManager._buffer = handle

329 return handle

330

331 # DeepEP LL uses RDMA so no SMs are used for communication

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

333 return 0

334

335 + def query_active_mask(self) -> torch.Tensor:

336 + buf = DeepEPLLAll2AllManager._buffer

337 + assert buf is not None

338 + if DeepEPLLAll2AllManager._mask is None:

339 + DeepEPLLAll2AllManager._mask = torch.zeros(

340 + self.world_size, device="cuda", dtype=torch.int32

341 + )

342 + buf.low_latency_query_mask_buffer(DeepEPLLAll2AllManager._mask)

343 + return DeepEPLLAll2AllManager._mask

344 +

345 + def query_fault(self) -> torch.Tensor:

346 + current = self.query_active_mask()

347 + if DeepEPLLAll2AllManager._last_mask is None:

348 + DeepEPLLAll2AllManager._last_mask = torch.zeros_like(current)

349 + has_fault = (current != DeepEPLLAll2AllManager._last_mask).any()

350 + return has_fault

351 +

352

353 @dataclass

354 class _NixlEPBufferState:

@@ -341,6 +365,8 @@ class NixlEPAll2AllManager(All2AllManagerBase):

365

366 _buffer: _NixlEPBufferState | None = None

367 _lock = threading.RLock()

368 + _mask: torch.Tensor | None = None

369 + _last_mask: torch.Tensor | None = None

370

371 def __init__(self, cpu_group, tcp_store_group=None):

372 if tcp_store_group is None:

@@ -350,6 +376,7 @@ class NixlEPAll2AllManager(All2AllManagerBase):

376 store=dist.PrefixStore("nixl_ep", cpu_group.get_group_store()),

377 )

378 super().__init__(cpu_group, tcp_store_group)

379 + self.support_fault_tolerance = True

380

381 self.max_num_ep_ranks = envs.VLLM_NIXL_EP_MAX_NUM_RANKS

382

@@ -508,6 +535,25 @@ class NixlEPAll2AllManager(All2AllManagerBase):

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

536 return 0

537

538 + def query_active_mask(self) -> torch.Tensor:

539 + state = NixlEPAll2AllManager._buffer

540 + assert state is not None

541 + if NixlEPAll2AllManager._mask is None:

542 + NixlEPAll2AllManager._mask = torch.zeros(

543 + self.max_num_ep_ranks, device="cuda", dtype=torch.int32

544 + )

545 + state.buffer.query_mask_buffer(NixlEPAll2AllManager._mask)

546 + return NixlEPAll2AllManager._mask[: state.active_ep_size]

547 +

548 + def query_fault(self) -> torch.Tensor:

549 + current = self.query_active_mask()

550 + last = NixlEPAll2AllManager._last_mask

551 + if last is None or last.shape != current.shape:

552 + NixlEPAll2AllManager._last_mask = torch.zeros_like(current)

553 + last = NixlEPAll2AllManager._last_mask

554 + has_fault = (current != last).any()

555 + return has_fault

556 +

557

558 class FlashInferNVLinkTwoSidedManager(All2AllManagerBase):

559 """

vllm/distributed/device_communicators/base_device_communicator.py

@@ -62,6 +62,8 @@ class All2AllManagerBase:

62 in_the_same_node_as(tcp_store_group, source_rank=0)

63 )

65 + self.support_fault_tolerance = False

66 +

67 def get_handle(self, kwargs):

68 # get a handle for the all2all communication,

69 # based on the kwargs.

@@ -102,6 +104,13 @@ class All2AllManagerBase:

104 # - raise a clear error if extra_tensors is not supported.

105 raise NotImplementedError

106

107 + def query_active_mask(self) -> torch.Tensor:

108 + raise NotImplementedError

109 +

110 + def query_fault(self) -> torch.Tensor:

111 + """Returns has_fault scalar."""

112 + raise NotImplementedError

113 +

114 def set_num_sms(self, num_sms: int):

115 pass

116

vllm/model_executor/layers/fused_moe/all2all_utils.py

+3 -3

@@ -56,7 +56,7 @@ if current_platform.is_cuda_alike():

56 )

59 -def _get_ep_all2all_manager(eep_stage: bool = False) -> Any:

59 +def get_ep_all2all_manager(eep_stage: bool = False) -> Any:

60 if eep_stage:

61 from vllm.distributed.elastic_ep.standby_state import get_standby_ep_group

@@ -146,7 +146,7 @@ def maybe_make_prepare_finalize(

146 "Detected DP deployment with no --enable-expert-parallel. "

147 "Falling back to AllGather+ReduceScatter dispatch/combine."

148 )

149 - all2all_manager = _get_ep_all2all_manager(eep_stage)

149 + all2all_manager = get_ep_all2all_manager(eep_stage)

150 return make_moe_prepare_and_finalize_naive_dp_ep(

151 is_sequence_parallel=moe.moe_parallel_config.is_sequence_parallel,

152 num_dispatchers=all2all_manager.world_size,

@@ -155,7 +155,7 @@ def maybe_make_prepare_finalize(

155 else:

156 return make_moe_prepare_and_finalize_no_dp_ep(use_monolithic)

157

158 - all2all_manager = _get_ep_all2all_manager(eep_stage)

158 + all2all_manager = get_ep_all2all_manager(eep_stage)

159

160 prepare_finalize: FusedMoEPrepareAndFinalize | None = None

161

vllm/v1/executor/multiproc_executor.py

+5 -1

@@ -939,7 +939,11 @@ class WorkerProc:

939 converted to a FAILURE response.

940 """

941 if isinstance(output, AsyncModelRunnerOutput):

942 - output = output.get_output()

942 + try:

943 + output = output.get_output()

944 + except Exception as e:

945 + logger.exception("Error getting async model runner output")

946 + output = e

947

948 if isinstance(output, Exception):

949 result = (WorkerProc.ResponseStatus.FAILURE, str(output))

vllm/v1/worker/gpu_model_runner.py

+19

@@ -58,6 +58,7 @@ from vllm.logger import init_logger

58 from vllm.lora.layers import LoRAMapping, LoRAMappingType

59 from vllm.model_executor.layers.attention import Attention, MLAAttention

60 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

61 +from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager

62 from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (

63 RoutedExpertsCapturer,

64 )

@@ -249,6 +250,7 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

250 async_output_copy_stream: torch.cuda.Stream,

251 vocab_size: int,

252 routed_experts: RoutedExpertsTensors | None = None,

253 + check_ep_fault: bool = False,

254 ):

255 self._model_runner_output = model_runner_output

256 self._invalid_req_indices = invalid_req_indices

@@ -262,6 +264,7 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

264 self.vocab_size = vocab_size

265 self._logprobs_tensors = logprobs_tensors

266 self._routed_experts = routed_experts

267 + self._has_fault: torch.Tensor | None = None

268

269 # Initiate the copy on a separate stream, but do not synchronize it.

270 default_stream = torch.cuda.current_stream()

@@ -280,6 +283,9 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

283 if self._routed_experts is not None

284 else None

285 )

286 + if check_ep_fault:

287 + has_fault = get_ep_all2all_manager().query_fault()

288 + self._has_fault = has_fault.to("cpu", non_blocking=True)

289 self.async_copy_ready_event.record()

290

291 def get_output(self) -> ModelRunnerOutput:

@@ -316,6 +322,14 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

322 output.routed_experts = self._routed_experts_cpu.tolists()

323 del self._routed_experts

324

325 + if self._has_fault is not None and self._has_fault.item():

326 + mask = get_ep_all2all_manager().query_active_mask()

327 + raise RuntimeError(

328 + "Fault detected in EP all2all communication: "

329 + "one or more ranks timed out during dispatch/combine. "

330 + f"Mask: {mask.cpu().tolist()}"

331 + )

332 +

333 return output

334

335

@@ -445,6 +459,10 @@ class GPUModelRunner(

459 self.device = device

460 self.dtype = self.model_config.dtype

461

462 + self.check_ep_fault = False

463 + if parallel_config.data_parallel_size > 1 and self.model_config.is_moe:

464 + self.check_ep_fault = get_ep_all2all_manager().support_fault_tolerance

465 +

466 self.kv_cache_dtype = kv_cache_dtype_str_to_dtype(

467 cache_config.cache_dtype, self.model_config

468 )

@@ -4695,6 +4713,7 @@ class GPUModelRunner(

4713 async_output_copy_stream=self._get_or_create_async_output_copy_stream(),

4714 vocab_size=self.input_batch.vocab_size,

4715 routed_experts=routed_experts_snapshot,

4716 + check_ep_fault=self.check_ep_fault,

4717 )

4718 with record_function_or_nullcontext(

4719 "gpu_model_runner: set_async_sampled_token_ids"