# DeepEP Shrink Mode for Dynamic Rank Masking: Runtime Fault Tolerance in Expert Parallelism

> Explore DeepEP's shrink mode for dynamic rank masking. Temporarily disable ranks during expert-parallel communication at runtime without communicator teardown. Enhance fault tolerance.

- Repository: [DeepSeek/DeepEP](https://github.com/deepseek-ai/DeepEP)
- Tags: internals
- Published: 2026-04-25

---

**DeepEP's shrink mode enables dynamic rank masking at runtime, allowing you to temporarily disable specific ranks during low-latency expert-parallel communication without tearing down the communicator.**

DeepEP is DeepSeek's high-performance communication library designed for Mixture-of-Experts (MoE) model training and inference across NVLink and RDMA networks. When operating large-scale distributed workloads, node failures or maintenance events require graceful degradation rather than complete job restarts. The **shrink mode** feature in `deepseek-ai/DeepEP` addresses this critical need by introducing dynamic rank masking capabilities to the low-latency communication buffers in [`deep_ep/buffer.py`](https://github.com/deepseek-ai/DeepEP/blob/main/deep_ep/buffer.py).

## What Is DeepEP Shrink Mode?

Shrink mode is a runtime fault-tolerance feature that enables **dynamic rank masking** during low-latency expert-parallel (EP) communication. When activated, the runtime maintains an internal mask buffer that determines which ranks participate in collective operations.

When you initialize a `Buffer` with `enable_shrink=True`, the C++ backend allocates an additional **mask buffer** (referenced in the low-latency APIs) that records per-rank participation states. A masked rank is automatically omitted from all subsequent `dispatch`, `combine`, and `clean` kernels, effectively "shrinking" the active communication group without destroying the underlying NCCL or custom communicator.

## How Dynamic Rank Masking Works

The masking system operates through three core mechanisms implemented in [`deep_ep/buffer.py`](https://github.com/deepseek-ai/DeepEP/blob/main/deep_ep/buffer.py):

- **Mask Storage**: An internal Boolean buffer indexed by rank ID, allocated during buffer construction at lines 42-61 when `enable_shrink=True` is specified.
- **Runtime Updates**: The `low_latency_update_mask_buffer(rank, mask)` method (lines 663-672) allows dynamic toggling of rank status without synchronizing the entire group.
- **Kernel Awareness**: Low-latency kernels consult the mask buffer before executing communication, skipping masked ranks entirely in the data path.

This architecture allows applications to respond to hardware failures by masking problematic ranks while continuing computation on the remaining healthy nodes.

## Enabling Shrink Mode in deep_ep/buffer.py

To activate dynamic rank masking, initialize your low-latency buffer with the `enable_shrink` parameter set to `True`. This triggers allocation of the internal mask buffer alongside your standard RDMA and NVLink memory pools.

```python
import deep_ep
import torch.distributed as dist

# Assume distributed group already initialized

group = dist.new_group()

# Allocate low-latency buffer with shrink mode enabled

buffer = deep_ep.Buffer(
    group,
    num_nvl_bytes=0,
    num_rdma_bytes=deep_ep.Buffer.get_low_latency_rdma_size_hint(
        num_tokens=128,
        hidden=7168,
        num_ranks=dist.get_world_size(),
        num_experts=288
    ),
    low_latency_mode=True,
    enable_shrink=True,  # Activates dynamic rank masking

    explicitly_destroy=True,
)

```

The constructor validates parameters and allocates the mask buffer only when explicitly requested, ensuring zero overhead for standard communication patterns that do not require fault tolerance.

## Runtime Mask Management API

Once enabled, the buffer exposes three critical methods for mask manipulation. These APIs allow fine-grained control over rank participation during distributed expert computation.

### Updating Rank Status

Use `low_latency_update_mask_buffer` to mask or unmask specific ranks at runtime. When a rank is masked (set to `True`), it is excluded from all subsequent low-latency kernels.

```python
import torch

# Mask rank 3 - it will be skipped in all subsequent operations

buffer.low_latency_update_mask_buffer(rank_to_mask=3, mask=True)

# Later, unmask the rank to restore participation

buffer.low_latency_update_mask_buffer(rank_to_mask=3, mask=False)

```

### Querying Current State

The `low_latency_query_mask_buffer` method (implemented at lines 674-682) copies the entire mask state into a user-provided tensor for inspection.

```python

# Check which ranks are currently masked

mask_status = torch.empty(dist.get_world_size(), dtype=torch.int, device='cuda')
buffer.low_latency_query_mask_buffer(mask_status)
print(mask_status)  # tensor([0, 0, 0, 1, 0, ...]) - rank 3 is masked

```

### Clearing the Mask

To reset all ranks to active status, call `low_latency_clean_mask_buffer` (lines 684-690), which zeroes the entire mask buffer.

```python

# Reset all masks to zero (no ranks excluded)

buffer.low_latency_clean_mask_buffer()

```

## Integration with Low-Latency Kernels

When shrink mode is active, the low-latency `dispatch` and `combine` operations automatically respect the mask buffer. Masked ranks receive zero tokens and do not participate in the data exchange, allowing the remaining ranks to complete the operation without hanging or erroring.

```python

# Prepare input data

x = torch.randn((128, 7168), dtype=torch.bfloat16, device='cuda')
topk_idx = torch.randint(0, 288, (128, 8), dtype=deep_ep.topk_idx_t, device='cuda')
topk_weights = torch.rand((128, 8), dtype=torch.float32, device='cuda')

# Mask rank 2 before dispatch

buffer.low_latency_update_mask_buffer(rank_to_mask=2, mask=True)

# Dispatch automatically excludes rank 2 from the communication pattern

recv_x, recv_count, handle, event, hook = buffer.low_latency_dispatch(
    x, topk_idx,
    num_tokens=128,
    num_experts=288,
    async_finish=False,
    return_recv_hook=False,
)

# Combine also respects the mask, aggregating only from active ranks

combined_x, combine_event, _ = buffer.low_latency_combine(
    recv_x, topk_idx, topk_weights, handle, use_logfmt=False
)

```

The kernels handle the masked ranks transparently, redistributing the computational load across the remaining active participants.

## Testing Rank Failure Simulation

The shrink mode functionality is validated in [`tests/test_low_latency.py`](https://github.com/deepseek-ai/DeepEP/blob/main/tests/test_low_latency.py) using the `--shrink-test` flag (lines 326-334). This test harness simulates rank failures by randomly masking ranks, updating the buffer state, and verifying that subsequent operations complete correctly while ignoring the masked participants.

The test suite confirms that:
- Masked ranks do not contribute to `dispatch` or `combine` results
- Unmasking restores full participation without buffer reallocation
- The `clean` operation properly resets all rank states to active

This automated validation ensures reliability when shrink mode is deployed in production environments handling actual hardware failures.

## Summary

- **DeepEP shrink mode** enables runtime dynamic rank masking by allocating an internal mask buffer when `enable_shrink=True` is passed to the `Buffer` constructor in [`deep_ep/buffer.py`](https://github.com/deepseek-ai/DeepEP/blob/main/deep_ep/buffer.py).
- The **mask buffer** tracks per-rank participation states, allowing specific ranks to be excluded from `dispatch`, `combine`, and `clean` kernels without communicator teardown.
- **Three APIs** control masking behavior: `low_latency_update_mask_buffer` for toggling individual ranks, `low_latency_query_mask_buffer` for inspection, and `low_latency_clean_mask_buffer` for bulk reset.
- Masked ranks are **automatically skipped** in low-latency communication flows, providing graceful degradation during hardware failures or maintenance.
- The functionality is **validated** via the `--shrink-test` flag in [`tests/test_low_latency.py`](https://github.com/deepseek-ai/DeepEP/blob/main/tests/test_low_latency.py), ensuring correctness under simulated failure conditions.

## Frequently Asked Questions

### What happens to in-flight data when a rank is masked?

In-flight data to a newly masked rank is dropped, and the rank stops receiving new tokens in subsequent `dispatch` operations. The `combine` kernel aggregates results only from unmasked ranks, ensuring that partial gradients or activations from the masked rank do not corrupt the global state. Your application logic should handle the missing rank's data or wait for it to be unmasked before critical synchronization points.

### Can shrink mode be used with standard (non-low-latency) EP operations?

No, shrink mode currently applies only to **low-latency** communication paths. The mask buffer and its associated APIs (`low_latency_update_mask_buffer`, `low_latency_query_mask_buffer`, `low_latency_clean_mask_buffer`) are specific to the low-latency implementation in [`deep_ep/buffer.py`](https://github.com/deepseek-ai/DeepEP/blob/main/deep_ep/buffer.py). Standard EP operations require a different approach to handle rank failures, typically involving communicator destruction and reinitialization.

### How does shrink mode affect memory allocation?

Enabling shrink mode increases memory footprint by allocating an additional mask buffer on the device. This buffer size scales with the number of ranks in the communicator (typically one Boolean or integer per rank). The memory overhead is minimal—usually less than a few kilobytes for typical world sizes—but the buffer persists for the lifetime of the `Buffer` object unless explicitly cleaned.

### Is there a performance penalty when shrink mode is enabled but unused?

When enabled but unused (all ranks unmasked), the performance impact is negligible. The mask check inside the low-latency kernels adds a single comparison operation per rank, which is shadowed by the communication latency. However, the mask buffer must still be allocated and synchronized, so you should only enable shrink mode when you anticipate the need for dynamic rank masking or fault tolerance.