[PR #370] [MERGED] [Feature] Add Elasticity Support to DeepEP for Fault-Tolerant EP Inference #538

Closed
opened 2026-10-27 21:03:47 +00:00 by bsakaguchi · 0 comments
Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/370
Author: @jammycc
Created: 8/21/2025
Status: ✅ Merged
Merged: 9/25/2025
Merged by: @sphish

Base: main ← Head: elastic


📝 Commits (10+)

  • 78349f8 support EP elastic shrink
  • 3507f2a Add test script for EP shrink failover
  • 6195f87 add coll buffer and allgather api
  • c7e11f5 remove allgather api, merge barrier and clean kernels, rename elastic to shrink
  • 0535a7b fix bug
  • 6bd18ff refine code, add shrink test script
  • ef616aa refine code
  • c683c3f refine code
  • 3e326a7 fix bug
  • c6a7030 add/remove blank lines

📊 Changes

6 files changed (+420 additions, -102 deletions)

View changed files

📝 csrc/deep_ep.cpp (+44 -2)
📝 csrc/deep_ep.hpp (+12 -1)
📝 csrc/kernels/api.cuh (+9 -0)
📝 csrc/kernels/internode_ll.cu (+259 -90)
📝 deep_ep/buffer.py (+32 -2)
📝 tests/test_low_latency.py (+64 -7)

📄 Description

Motivation and Target

This PR ensures that large-scale Expert Parallel (EP) inference remain operational in a downscaled state when encountering failures — such as individual GPU, NIC, or node outages, or even corner-case engine failures, instead of a full teardown.

This provides the following benefits:

  • The downscaled deployment can still provide degraded (lower-capacity) service

  • For multi-endpoint/API-server setups, most requests will only experience a brief stall before resuming processing, rather than being entirely aborted.

  • This graceful degradation provides sufficient time for instance-level redeployment or recovery actions.

This PR includes

  • Enable dispatch and combine operations to bypass failed ranks by introducing per-rank status information in the buffer, achieving real-time degradation.

  • Decentralized rank failure detection, with each rank monitoring communication timeout to determine failure.

  • Provide query and update APIs for the engine to query and update per-rank status.

This PR does not include (as future work)

  • Support for Prefill phase / Normal kernels

    • If a Prefill request fails, it usually needs to be fully re-executed anyway, making fault tolerance less impactful.

    • In PD-disaggregated architectures, the Prefill stage typically uses smaller EP sizes — sometimes even single-node deployments.

  • Dynamic scaling up (expansion)

    • Scaling up involves complex initialization and warmup procedures for new ranks, which must be orchestrated by the engine (e.g., model loading, expert placement, CUDA context setup). Communication connectivity is only one part of this process. Given this complexity, we leave it as future work.

DeepEP Modifications

Currently, DeepEP’s low-latency (LL) kernels hang when any rank fails, causing the entire EP inference service and request processing pipeline to stall. We aim to enhance DeepEP with failure detection, graceful exit on faults, and dynamic downscaling. To this end, all communication-related APIs — including low_latency_dispatch, low_latency_combine, and clean_low_latency_buffer — must support these capabilities.

Additionally, since other engine-level collective communications (e.g., metadata synchronization within DP groups) can also be affected by rank failures, we introduce a new low_latency_allgather API to replace these collectives.

1. Failure Detection and Downscaling via Mask Buffer

Each node maintains a mask buffer to track the status of all ranks. Once a rank failure is detected, its corresponding status in the mask buffer is set to 1, indicating that subsequent communication with this rank should be skipped.

Key Changes:

  • Add enable_elastic parameter during buffer initialization to enable elastic mode. This triggers allocation of an additional mask_buffer of size num_ranks (int32).

  • Add timeout-based failure detection and rank masking in all communication APIs (low_latency_dispatch, low_latency_combine, clean_low_latency_buffer) :

    • If a rank's data is not received within NUM_TIMEOUT_CYCLES, its status in the mask buffer is set to 1.

    • Subsequent communication will neither receive data from nor send data to masked ranks.

  • New APIs:

    • low_latency_query_mask_buffer: Query current rank mask status.

    • low_latency_update_mask_buffer: Manually update the mask status of a specific rank.

2. Synchronization Mechanism via Sync Buffer

For non-dispatch/combine communication (e.g., clean_low_latency_buffer), global synchronization is required. Previously, this relied on NVSHMEM barriers, which fail if any rank is down.

To support synchronization in downscaled scenarios:

  • Introduce a dedicated sync_buffer to track synchronization counters.

    • This is similar to the flag buffer used in dispatch/combine, but we allocate a separate buffer to avoid interference.
  • Replace global nvshmem_barrier calls in clean_low_latency_buffer with sync-buffer-based logic, enabling barrier even when some ranks are masked.

3. New Collective Communication: low_latency_allgather

We observed that non-EP collectives (e.g., metadata sync in pure DP attention) can also be blocked by rank failures. To address this, we implement a fault-tolerant low_latency_allgather API to replace these communication.

  • Add num_coll_buffer_bytes parameter during initialization to allocate space for collective communication buffers.

  • Use the sync_buffer for synchronization and skip masked ranks.

Engine Integration Example

We have validated this design by integrating DeepEP into vLLM. We briefly describe below how the engine integrates with DeepEP’s fault-tolerance features, using vLLM v1 as an example.

workflow

1. Enable Elastic Mode in DeepEP Initialization

deep_ep_buffer = Buffer(
    group,
    ...
    enable_elastic=True,
    num_coll_buffer_bytes=1024 * 1024  # For collective operations
)

2. Check Status After Each Decode Step

In gpu_worker.py, the Worker class exposes a check_step_status method, called via RPC by EngineCore to check for failures:

def check_step_status(self) -> Optional[list[int]]:
    if self.status_tensor is None:
        self.status_tensor = torch.zeros(
            get_expert_parallel_world_size(),
            dtype=torch.int32,
            device='cuda'
        )
    torch.cuda.synchronize()
    deep_ep_buffer.low_latency_query_mask_buffer(self.status_tensor)
    torch.cuda.synchronize()

    failed_ranks = torch.nonzero(self.status_tensor, as_tuple=True)[0].tolist()
    return failed_ranks

3. Handle Detected Failures

Once EngineCore detects new failed ranks:

  • Report to LLMEngine

  • Trigger expert rebalance (EPLB)

  • Recover requests (via migration or recomputation)

def handle_failed_engines(self, failed_engines: set[int]):
    # Notify core clients of failure
    combined_list = [self.engine_index] + list(failed_engines)
    msg = self.ENGINE_CORE_FAIL + bytes(f'{combined_list}', encoding='utf-8')
    self.output_queue.put_nowait(msg)

    # Expert rebalance (EPLB)
    rebalanced_experts = self._calculate_expert_mapping(self.failed_ranks)
    self.model_executor.reload_experts(rebalanced_experts)

    # Recover requests: pop last token and KV cache
    for r in self.scheduler.running:
        r.num_computed_tokens -= 1
        r.pop_output_token_ids()
        self.model_executor.pop_kv_cache(r.request_id)

4. Handling Externally Triggered Failures

In gpu_worker.py, the Worker class exposes a update_deepep_mask_buffer method, called via RPC by EngineCore to manually mask a rank:

def update_deepep_mask_buffer(self, rank_to_mask: int, mask: bool):
    deep_ep_buffer.low_latency_update_mask_buffer(rank_to_mask, mask)

5. Redundant Expert Slots for CUDAGraph

  • Reserve extra expert slots per rank to accommodate experts from failed ranks.

  • Ensures sufficient capacity for redistribution upon single-rank failure.

6. Decoupling from Other Communication Backends

  • Replace other collectives (e.g., DP-group metadata sync) with DeepEP’s extended low_latency_allgather.

  • TP-group communication is not replaced, as entire TP groups can be treated as atomic units — if one fails, the whole group is isolated.


🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepEP/pull/370 **Author:** [@jammycc](https://github.com/jammycc) **Created:** 8/21/2025 **Status:** ✅ Merged **Merged:** 9/25/2025 **Merged by:** [@sphish](https://github.com/sphish) **Base:** `main` ← **Head:** `elastic` --- ### 📝 Commits (10+) - [`78349f8`](https://github.com/deepseek-ai/DeepEP/commit/78349f8fc2c88522d62de0711ef80587b2ddda5d) support EP elastic shrink - [`3507f2a`](https://github.com/deepseek-ai/DeepEP/commit/3507f2a00a849b216ab7730852a3ac26777e49a0) Add test script for EP shrink failover - [`6195f87`](https://github.com/deepseek-ai/DeepEP/commit/6195f87505b6ae95d2c3187e7de470c124f12b94) add coll buffer and allgather api - [`c7e11f5`](https://github.com/deepseek-ai/DeepEP/commit/c7e11f534e2c5a2ca2980638571a46d593aef043) remove allgather api, merge barrier and clean kernels, rename elastic to shrink - [`0535a7b`](https://github.com/deepseek-ai/DeepEP/commit/0535a7b2aaba53abd152fd3de881d91fd22cb5dd) fix bug - [`6bd18ff`](https://github.com/deepseek-ai/DeepEP/commit/6bd18ff87697ffb244fdc3dce6fbdc53589a854a) refine code, add shrink test script - [`ef616aa`](https://github.com/deepseek-ai/DeepEP/commit/ef616aab8189f448811ab8be779c026cc72b6783) refine code - [`c683c3f`](https://github.com/deepseek-ai/DeepEP/commit/c683c3f4b28fc52e7f10944e8b00ed077c0d6be4) refine code - [`3e326a7`](https://github.com/deepseek-ai/DeepEP/commit/3e326a7c377f5fb4095cd8c4f54b3b7542505b3e) fix bug - [`c6a7030`](https://github.com/deepseek-ai/DeepEP/commit/c6a7030bd25a0b3b6fb434d438f21ba94e751b81) add/remove blank lines ### 📊 Changes **6 files changed** (+420 additions, -102 deletions) <details> <summary>View changed files</summary> 📝 `csrc/deep_ep.cpp` (+44 -2) 📝 `csrc/deep_ep.hpp` (+12 -1) 📝 `csrc/kernels/api.cuh` (+9 -0) 📝 `csrc/kernels/internode_ll.cu` (+259 -90) 📝 `deep_ep/buffer.py` (+32 -2) 📝 `tests/test_low_latency.py` (+64 -7) </details> ### 📄 Description ## Motivation and Target This PR ensures that large-scale Expert Parallel (EP) inference remain operational in a **downscaled state** when encountering failures — such as individual GPU, NIC, or node outages, or even corner-case engine failures, instead of a full teardown. This provides the following benefits: * The downscaled deployment can still provide **degraded (lower-capacity) service** * For multi-endpoint/API-server setups, most requests will only experience a **brief stall** before resuming processing, rather than being entirely aborted. * This graceful degradation provides sufficient time for instance-level redeployment or recovery actions. ### This PR includes * Enable `dispatch` and `combine` operations to bypass failed ranks by introducing per-rank status information in the buffer, achieving real-time degradation. * Decentralized rank failure detection, with each rank monitoring communication timeout to determine failure. * Provide query and update APIs for the engine to query and update per-rank status. ### This PR does not include (as future work) * **Support for Prefill phase / Normal kernels** * If a Prefill request fails, it usually needs to be fully re-executed anyway, making fault tolerance less impactful. * In PD-disaggregated architectures, the Prefill stage typically uses smaller EP sizes — sometimes even single-node deployments. * **Dynamic scaling up (expansion)** * Scaling up involves complex initialization and warmup procedures for new ranks, which must be orchestrated by the engine (e.g., model loading, expert placement, CUDA context setup). Communication connectivity is only one part of this process. Given this complexity, we leave it as future work. ## DeepEP Modifications Currently, DeepEP’s low-latency (LL) kernels hang when any rank fails, causing the entire EP inference service and request processing pipeline to stall. We aim to enhance DeepEP with **failure detection, graceful exit on faults, and dynamic downscaling**. To this end, all communication-related APIs — including `low_latency_dispatch`, `low_latency_combine`, and `clean_low_latency_buffer` — must support these capabilities. Additionally, since other engine-level collective communications (e.g., metadata synchronization within DP groups) can also be affected by rank failures, we introduce a new `low_latency_allgather` API to replace these collectives. ### 1. Failure Detection and Downscaling via Mask Buffer Each node maintains a **mask buffer** to track the status of all ranks. Once a rank failure is detected, its corresponding status in the mask buffer is set to 1, indicating that subsequent communication with this rank should be skipped. #### Key Changes: * Add `enable_elastic` parameter during buffer initialization to enable elastic mode. This triggers allocation of an additional `mask_buffer` of size `num_ranks` (int32). * Add timeout-based failure detection and rank masking in all communication APIs (`low_latency_dispatch`, `low_latency_combine`, `clean_low_latency_buffer`) : * If a rank's data is not received within `NUM_TIMEOUT_CYCLES`, its status in the mask buffer is set to 1. * Subsequent communication will neither receive data from nor send data to masked ranks. * New APIs: * `low_latency_query_mask_buffer`: Query current rank mask status. * `low_latency_update_mask_buffer`: Manually update the mask status of a specific rank. ### 2. Synchronization Mechanism via Sync Buffer For non-dispatch/combine communication (e.g., `clean_low_latency_buffer`), global synchronization is required. Previously, this relied on NVSHMEM barriers, which fail if any rank is down. To support synchronization in downscaled scenarios: * Introduce a dedicated **sync\_buffer** to track synchronization counters. * This is similar to the flag buffer used in dispatch/combine, but we allocate a separate buffer to avoid interference. * Replace global `nvshmem_barrier` calls in `clean_low_latency_buffer` with sync-buffer-based logic, enabling barrier even when some ranks are masked. ### 3. New Collective Communication: **`low_latency_allgather`** We observed that non-EP collectives (e.g., metadata sync in pure DP attention) can also be blocked by rank failures. To address this, we implement a fault-tolerant `low_latency_allgather` API to replace these communication. * Add `num_coll_buffer_bytes` parameter during initialization to allocate space for collective communication buffers. * Use the **sync\_buffer** for synchronization and skip masked ranks. ## Engine Integration Example We have validated this design by integrating DeepEP into vLLM. We briefly describe below how the engine integrates with DeepEP’s fault-tolerance features, using **vLLM v1** as an example. <img width="2864" height="2369" alt="workflow" src="https://github.com/user-attachments/assets/c6c29cbc-76e8-4a4f-9f4b-2ea2343993d5" /> ### 1. Enable Elastic Mode in DeepEP Initialization ```python deep_ep_buffer = Buffer( group, ... enable_elastic=True, num_coll_buffer_bytes=1024 * 1024 # For collective operations ) ``` ### 2. Check Status After Each Decode Step In `gpu_worker.py`, the `Worker` class exposes a `check_step_status` method, called via RPC by `EngineCore` to check for failures: ```python def check_step_status(self) -> Optional[list[int]]: if self.status_tensor is None: self.status_tensor = torch.zeros( get_expert_parallel_world_size(), dtype=torch.int32, device='cuda' ) torch.cuda.synchronize() deep_ep_buffer.low_latency_query_mask_buffer(self.status_tensor) torch.cuda.synchronize() failed_ranks = torch.nonzero(self.status_tensor, as_tuple=True)[0].tolist() return failed_ranks ``` ### 3. Handle Detected Failures Once `EngineCore` detects new failed ranks: * Report to `LLMEngine` * Trigger expert rebalance (EPLB) * Recover requests (via migration or recomputation) ```python def handle_failed_engines(self, failed_engines: set[int]): # Notify core clients of failure combined_list = [self.engine_index] + list(failed_engines) msg = self.ENGINE_CORE_FAIL + bytes(f'{combined_list}', encoding='utf-8') self.output_queue.put_nowait(msg) # Expert rebalance (EPLB) rebalanced_experts = self._calculate_expert_mapping(self.failed_ranks) self.model_executor.reload_experts(rebalanced_experts) # Recover requests: pop last token and KV cache for r in self.scheduler.running: r.num_computed_tokens -= 1 r.pop_output_token_ids() self.model_executor.pop_kv_cache(r.request_id) ``` ### 4. Handling Externally Triggered Failures In `gpu_worker.py`, the `Worker` class exposes a `update_deepep_mask_buffer` method, called via RPC by `EngineCore` to manually mask a rank: ```python def update_deepep_mask_buffer(self, rank_to_mask: int, mask: bool): deep_ep_buffer.low_latency_update_mask_buffer(rank_to_mask, mask) ``` ### 5. Redundant Expert Slots for CUDAGraph * Reserve extra expert slots per rank to accommodate experts from failed ranks. * Ensures sufficient capacity for redistribution upon single-rank failure. ### 6. Decoupling from Other Communication Backends * Replace other collectives (e.g., DP-group metadata sync) with DeepEP’s extended `low_latency_allgather`. * **TP-group communication is not replaced**, as entire TP groups can be treated as atomic units — if one fails, the whole group is isolated. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
bsakaguchi 2026-10-27 21:03:47 +00:00
Sign in to join this conversation.
No labels
pull-request
No milestone
No project
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
deepseek-ai/DeepEP#538
No description provided.