### Build NCCL from Source

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/errors.md

Install NCCL from source if pip installation is not suitable. Ensure CUDA_HOME is correctly set.

```bash
wget https://developer.nvidia.com/nccl
tar -xf nccl_*.tar.gz
cd nccl_*
make install CUDA_HOME=/usr/local/cuda
```

--------------------------------

### Install NCCL Library

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/errors.md

Install the NCCL library using pip if it's missing. Ensure the version is compatible with your CUDA installation.

```bash
pip install "nvidia-nccl-cu13>=2.30.4" --no-deps
```

--------------------------------

### Configuration Parameters

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/COMPLETION_SUMMARY.txt

Details on configuration parameters available for tuning and setup.

```APIDOC
## Configuration

### Description
Provides details on various configuration parameters for the library.

### Constructor Parameters
- 15+ parameters available for class constructors.

### Dispatch Parameters
- 20+ parameters available for dispatch operations.

### Environment Variables
- 30+ environment variables for system-level tuning.

### Configuration Scenarios
- 5+ documented scenarios illustrating configuration usage.
```

--------------------------------

### Verify NVSHMEM Installation

Source: https://github.com/deepseek-ai/deepep/blob/main/docs/nvshmem.md

Run the nvshmem-info command to confirm that NVSHMEM is installed correctly and to display its details.

```bash
nvshmem-info -a # Should display details of nvshmem
```

--------------------------------

### Install DeepEP

Source: https://github.com/deepseek-ai/deepep/blob/main/README.md

Installs the DeepEP package. After installation, import 'deep_ep' in your Python projects.

```bash
python setup.py install
```

--------------------------------

### Example Usage of Config Class

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/types.md

Shows how to obtain a recommended configuration or create a custom Config object for use in buffer dispatch operations.

```python
from deep_ep import Config, Buffer

# Get recommended config
config = Buffer.get_dispatch_config(num_ranks=8)

# Or create custom config
custom_config = Config(
    num_sms=20,
    param1=6,
    param2=256,
    param3=6,
    param4=128
)

# Use in dispatch
buffer.dispatch(
    x, handle=None, topk_idx=topk_idx,
    config=custom_config
)
```

--------------------------------

### Example Usage of topk_idx_t

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/types.md

Demonstrates how to retrieve the current topk_idx_t and allocate expert indices with the specified dtype and device.

```python
import deep_ep

# Get the current topk_idx_t type
topk_dtype = deep_ep.topk_idx_t

# Allocate expert indices
topk_idx = torch.randint(0, 256, (8192, 8), dtype=topk_dtype, device='cuda')
```

--------------------------------

### Install NVSHMEM via pip

Source: https://github.com/deepseek-ai/deepep/blob/main/docs/nvshmem.md

Install NVSHMEM using pip, compatible with CUDA 12.

```bash
pip install nvidia-nvshmem-cu12
```

--------------------------------

### Example: Forward dispatch with async execution

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/elastic-buffer.md

Demonstrates how to dispatch tokens asynchronously and wait for communication completion. Ensure to handle the event overlap for correct synchronization.

```python
# Forward dispatch
x = torch.randn(8192, 7168, dtype=torch.bfloat16, device='cuda')
topk_idx = torch.randint(0, 256, (8192, 8), device='cuda')
topk_weights = torch.rand(8192, 8, device='cuda')

recv_x, recv_topk_idx, recv_topk_weights, handle, event = buffer.dispatch(
    x,
    topk_idx=topk_idx,
    topk_weights=topk_weights,
    num_experts=256,
    num_max_tokens_per_rank=8192,
    async_with_compute_stream=True
)

# Wait for communication if overlapping
with event:
    # Do local expert computation
    pass

```

--------------------------------

### Get Physical Domain Sizes

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/elastic-buffer.md

Obtain the physical domain sizes, which correspond to the number of RDMA and NVLink ranks in the current distributed setup.

```python
num_rdma_ranks, num_nvlink_ranks = buffer.get_physical_domain_size()
print(f"RDMA ranks: {num_rdma_ranks}, NVLink ranks: {num_nvlink_ranks}")
```

--------------------------------

### API Reference - Utilities

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/COMPLETION_SUMMARY.txt

Documentation for all utility modules and functions within DeepEP, covering environment setup, NCCL management, bandwidth detection, hardware capability checking, and FP8 casting.

```APIDOC
## Utilities Module

### Description
This module encompasses various utility functions and tools that support the core functionalities of DeepEP, including environment configuration, communication management, and data type conversions.

### Key Utilities
- **Environment Setup**: Functions for configuring the DeepEP environment.
- **NCCL Management**: Utilities for managing NVIDIA Collective Communications Library (NCCL) operations.
- **Bandwidth Detection**: Functions to measure and report network bandwidth.
- **Hardware Capability Checking**: Tools to assess system hardware capabilities.
- **FP8 Casting**: Utilities for casting data to and from FP8 formats.

### Environment Variables
- Documented over 30 environment variables that influence DeepEP's behavior and performance.
```

--------------------------------

### Get Buffer Size Hint

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/elastic-buffer.md

Calculates the recommended buffer size in bytes without constructing the buffer. Use this to determine an appropriate size for general buffer needs, ensuring it's 2 MB-aligned.

```python
from deepspeed import dist

# Assuming ElasticBuffer is defined elsewhere and group is a valid ProcessGroup
# Example usage:
# size = ElasticBuffer.get_buffer_size_hint(
#     group, num_max_tokens_per_rank=8192, hidden=7168, num_topk=8
# )
```

--------------------------------

### Set NCCL Root Directory

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/errors.md

If NCCL is installed in a non-standard location, set the EP_NCCL_ROOT_DIR environment variable before building.

```bash
export EP_NCCL_ROOT_DIR=/opt/nccl
python setup.py build
```

--------------------------------

### Minimal ElasticBuffer Construction

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/configuration.md

Use this for a basic setup where buffer size is automatically calculated based on MoE settings. Ensure the `group` object is defined in your scope.

```python
from deep_ep import ElasticBuffer

buffer = ElasticBuffer(
    group,
    num_max_tokens_per_rank=8192,
    hidden=7168,
    num_topk=8
)
```

--------------------------------

### Get Combine Performance Config

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/legacy-buffer.md

Get a recommended performance configuration for combine operations, similar to dispatch configurations, based on the number of ranks.

```python
@staticmethod
def get_combine_config(num_ranks: int) -> Config
```

--------------------------------

### Get Pipeline-Parallel Buffer Size Hint

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/elastic-buffer.md

Calculates the minimum buffer size for pipeline-parallel operations. Returns the recommended size in bytes, aligned to 2 MB. This is an experimental feature.

```python
# Assuming ElasticBuffer is defined elsewhere
# Example usage:
# pp_size = ElasticBuffer.get_pp_buffer_size_hint(
#     num_max_tensor_bytes=1024*1024, num_max_inflight_tensors=4
# )
```

--------------------------------

### Get Logical Domain Sizes

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/elastic-buffer.md

Retrieve the logical domain sizes, representing the number of scaleout and scaleup ranks. This is relevant for hierarchical communication configurations.

```python
num_scaleout_ranks, num_scaleup_ranks = buffer.get_logical_domain_size()
```

--------------------------------

### Hybrid Mode Configuration for Multi-Node

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/README.md

Shows how to configure ElasticBuffer for hybrid mode in a multi-node setup, enabling hierarchical communication. It also demonstrates how to retrieve logical and physical domain sizes.

```python
buffer = ElasticBuffer(
    group, num_max_tokens_per_rank=8192, hidden=7168, num_topk=8,
    allow_hybrid_mode=True,  # Enable hierarchical communication
    allow_multiple_reduction=True
)

num_scaleout, num_scaleup = buffer.get_logical_domain_size()
num_rdma, num_nvlink = buffer.get_physical_domain_size()
print(f"Topology: {num_scaleout} x {num_scaleup} (RDMA={num_rdma}, NVLink={num_nvlink})")
```

--------------------------------

### Get Low Latency RDMA Size Hint

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/legacy-buffer.md

Obtain a recommended buffer size in bytes for low-latency RDMA operations. This calculation depends on dispatch tokens, hidden dimension size, number of ranks, and number of experts.

```python
size = Buffer.get_low_latency_rdma_size_hint(
    num_max_dispatch_tokens_per_rank=8192,
    hidden=7168,
    num_ranks=8,
    num_experts=256
)
```

--------------------------------

### Initialize Legacy Buffer in Low-Latency Mode

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/legacy-buffer.md

Initializes the legacy Buffer in low-latency mode, which uses RDMA for intranode communication. Specify the RDMA buffer size and the number of QPs per rank, which must match the number of local experts.

```python
# Low-latency mode
buffer_ll = Buffer(
    group,
    num_rdma_bytes=4 * 1024 * 1024 * 1024,
    low_latency_mode=True,
    num_qps_per_rank=8  # Must equal number of local experts
)
```

--------------------------------

### Initialize Distributed Training (8 GPUs, NVLink)

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/configuration.md

Initializes the distributed environment and configures the ElasticBuffer for training on a single node with 8 GPUs using NVLink.

```python
import torch.distributed as dist
from deep_ep import ElasticBuffer
from deep_ep.utils.envs import init_dist

rank, world_size, group = init_dist(
    local_rank=torch.cuda.current_device(),
    num_local_ranks=8,
    seed=42
)

buffer = ElasticBuffer(
    group,
    num_max_tokens_per_rank=8192,
    hidden=7168,
    num_topk=8,
    use_fp8_dispatch=False,
    allow_hybrid_mode=False,  # Single node
    allow_multiple_reduction=True,
    prefer_overlap_with_compute=False  # Use more SMs
)
```

--------------------------------

### Set NVSHMEM Root Directory

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/errors.md

If NVSHMEM is installed in a non-standard location, set the EP_NVSHMEM_ROOT_DIR environment variable before building.

```bash
export EP_NVSHMEM_ROOT_DIR=/path/to/nvshmem
python setup.py build
```

--------------------------------

### init_dist

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/utilities.md

Initializes the distributed environment using the NCCL backend. It configures distributed training settings based on the provided ranks and seed, and returns the global rank, world size, and process group.

```APIDOC
## init_dist()

### Description
Initialize distributed environment with NCCL backend, configurable per cluster.

### Signature
```python
def init_dist(local_rank: int, num_local_ranks: int, seed: int = 0) -> Tuple[int, int, dist.ProcessGroup]
```

### Parameters
#### Path Parameters
- **local_rank** (int) - Required - Local rank of this process
- **num_local_ranks** (int) - Required - Number of local ranks (GPUs per node)
- **seed** (int) - Optional - Global random seed (default: 0)

### Returns
- **rank** (int) - Global rank of the process.
- **world_size** (int) - Total number of processes in the distributed environment.
- **group** (dist.ProcessGroup) - The process group object.

### Environment Variables Used
- `MASTER_ADDR` (default: `127.0.0.1`)
- `MASTER_PORT` (default: `8361`)
- `WORLD_SIZE` (default: `1`)
- `RANK` (default: `0`)

### Example
```python
from deep_ep.utils.envs import init_dist, dist_print

rank, world_size, group = init_dist(
    local_rank=0,
    num_local_ranks=8,
    seed=42
)

dist_print(f"Rank {rank}/{world_size}", once_in_node=True)
```
```

--------------------------------

### Get Communication Stream

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/elastic-buffer.md

Retrieve the internal communication stream used by the ElasticBuffer. This can be useful for custom stream management.

```python
comm_stream = buffer.get_comm_stream()
# Use for custom stream management
```

--------------------------------

### Build and Develop DeepEP

Source: https://github.com/deepseek-ai/deepep/blob/main/README.md

Builds the project and creates symbolic links for SO files. Modify SO names and cluster settings as needed for testing.

```bash
python setup.py build
# You may modify the specific SO names according to your own platform
ln -s build/lib.linux-x86_64-cpython-38/deep_ep_cpp.cpython-38-x86_64-linux-gnu.so

# Run test cases
# NOTES: you may modify the `init_dist` function in `tests/utils/envs.py`
# according to your own cluster settings, and launch into multiple nodes
python tests/elastic/test_ep.py
python tests/elastic/test_agrs.py
python tests/elastic/test_engram.py
python tests/elastic/test_pp.py
```

--------------------------------

### get_local_buffer_tensor()

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/legacy-buffer.md

Gets the raw buffer as a PyTorch tensor, with options for slicing, data type, and using RDMA buffers.

```APIDOC
## get_local_buffer_tensor()

### Description

Get raw buffer as a PyTorch tensor (with optional slicing).

### Method

`get_local_buffer_tensor(dtype: torch.dtype, size: Optional[torch.Size] = None, offset: int = 0, use_rdma_buffer: bool = False)`

### Parameters

#### Query Parameters

- **dtype** (`torch.dtype`) - Required - Data type for the returned tensor
- **size** (`Optional[torch.Size]`) - Optional - Slice size in elements; if None, returns full buffer
- **offset** (`int`) - Optional - Starting element offset
- **use_rdma_buffer** (`bool`) - Optional - Use RDMA buffer instead of NVLink

### Returns

`torch.Tensor` — Buffer tensor or slice

### Example

```python
buffer_tensor = buffer.get_local_buffer_tensor(
    dtype=torch.bfloat16,
    size=torch.Size([8192, 7168]),
    use_rdma_buffer=False
)
```
```

--------------------------------

### Debugging Configuration

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/utilities.md

Enable detailed debugging output for buffer initialization and NCCL checks by setting environment variables. This is useful for troubleshooting.

```python
import os

os.environ['EP_BUFFER_DEBUG'] = '1'
os.environ['EP_JIT_DEBUG'] = '1'
os.environ['EP_SUPPRESS_NCCL_CHECK'] = '1'

# Now buffer initialization prints detailed info
buffer = ElasticBuffer(group, num_bytes=size)
```

--------------------------------

### Get Communication Stream from Legacy Buffer

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/legacy-buffer.md

Retrieves the internal communication stream used by the buffer. This stream is of type `torch.Stream`.

```python
buffer.get_comm_stream()
```

--------------------------------

### Get RDMA Bandwidth

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/utilities.md

Retrieves the RDMA bandwidth in GB/s for a specified NIC. This function is cached and can be configured via environment variables.

```python
from deep_ep.utils.envs import get_rdma_gbs

rdma_bw = get_rdma_gbs(nic_name='mlx5_0')
```

--------------------------------

### Get NVLink Bandwidth

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/utilities.md

Retrieves the total NVLink bandwidth in GB/s. This function is cached for performance and allows for an optional efficiency factor.

```python
from deep_ep.utils.envs import get_nvlink_gbs

nvlink_bw = get_nvlink_gbs(factor=0.9)
```

--------------------------------

### Instantiate EventOverlap

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/event-overlap.md

Instantiate EventOverlap by capturing an event or getting it directly from buffer dispatch/combine operations. The `extra_tensors` parameter is for backward compatibility.

```python
from deep_ep.utils.event import EventOverlap
from deep_ep import ElasticBuffer

buffer = ElasticBuffer(group, num_bytes=size)

# Capture event
event = ElasticBuffer.capture()
event_overlap = EventOverlap(event)

# Or get directly from dispatch/combine
recv_x, _, _, _, event_overlap = buffer.dispatch(...)
```

--------------------------------

### Buffer Constructor: Low-latency Mode

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/configuration.md

Enable low-latency mode for RDMA communication within a node. Ensure num_qps_per_rank matches num_local_experts.

```python
# Low-latency mode
buffer = Buffer(
    group,
    num_rdma_bytes=4 * 1024**3,
    low_latency_mode=True,
    num_qps_per_rank=8  # = num_local_experts
)
```

--------------------------------

### Initialize ElasticBuffer with MoE Settings

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/elastic-buffer.md

Initialize the ElasticBuffer using expert-parallel (MoE) settings. Ensure the distributed environment is initialized before creating the buffer. The buffer size is automatically calculated based on MoE parameters.

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

# Initialize distributed environment
dist.init_process_group(backend='nccl')
group = dist.new_group(list(range(dist.get_world_size())))

# Initialize buffer with MoE settings
buffer = ElasticBuffer(
    group,
    num_max_tokens_per_rank=8192,
    hidden=7168,
    num_topk=8,
    use_fp8_dispatch=True,
    allow_hybrid_mode=True
)
```

--------------------------------

### Class and Method Documentation

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/COMPLETION_SUMMARY.txt

Overview of documented classes and their methods/attributes.

```APIDOC
## ElasticBuffer Class

### Description
Provides functionality for elastic buffering with 15 methods and 6 static methods.

### Methods
- 15 documented methods.
- 6 static methods.

### Attributes
None specified.

## Buffer Class

### Description
Core buffer class with 11 methods and 3 static methods.

### Methods
- 11 documented methods.
- 3 static methods.

### Attributes
None specified.

## EPHandle Class

### Description
Represents an Expert Parallelism handle with 9 attributes.

### Methods
None specified.

### Attributes
- 9 documented attributes.

## EventOverlap Class

### Description
Manages event overlap with 1 method and context manager support.

### Methods
- 1 documented method.
- Supports context manager protocol.

### Attributes
None specified.

## Utility Functions

### Description
Over 10 utility functions are exported per module.

### Methods
- Numerous utility functions.

### Attributes
None specified.

## Other Documented Types

### Description
Includes NCCLCommHandle, Config, and other essential types.

### Methods
None specified.

### Attributes
None specified.
```

--------------------------------

### Suppress NCCL Version Check

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/errors.md

Allows suppressing the NCCL version check by setting an environment variable. This is not recommended and should be used only when a matching NCCL installation is not feasible.

```python
import os

# Suppress check (not recommended)
os.environ['EP_SUPPRESS_NCCL_CHECK'] = '1'

# Better: Install matching NCCL
# pip install "nvidia-nccl-cu13>=2.30.4" --no-deps
```

--------------------------------

### Initialize Legacy Buffer in High-Throughput Mode

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/legacy-buffer.md

Initializes the legacy Buffer for high-throughput communication. Configure NVLink and RDMA buffer sizes according to your needs. Ensure PyTorch distributed group is initialized.

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

# High-throughput mode
dist.init_process_group(backend='nccl')
group = dist.new_group(list(range(dist.get_world_size())))

buffer = Buffer(
    group,
    num_nvl_bytes=2 * 1024 * 1024 * 1024,  # 2 GB for NVLink
    num_rdma_bytes=4 * 1024 * 1024 * 1024  # 4 GB for RDMA
)
```

--------------------------------

### Initialize Distributed Environment

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/utilities.md

Initializes the distributed environment using the NCCL backend. This function is essential for setting up multi-GPU or multi-node training. It configures communication groups and returns global rank, world size, and the process group object.

```python
from deep_ep.utils.envs import init_dist, dist_print

rank, world_size, group = init_dist(
    local_rank=0,
    num_local_ranks=8,
    seed=42
)

dist_print(f"Rank {rank}/{world_size}", once_in_node=True)
```

--------------------------------

### pp_set_config()

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/elastic-buffer.md

Configures pipeline-parallel parameters. This is an experimental method.

```APIDOC
## pp_set_config()

### Description
Configure pipeline-parallel parameters (experimental).

### Method
`pp_set_config(self, num_max_tensor_bytes: int, num_max_inflight_tensors: int)
`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **num_max_tensor_bytes** (`int`) - Required - Description not specified in source.
- **num_max_inflight_tensors** (`int`) - Required - Description not specified in source.

### Request Example
```python
# Example usage (assuming 'self' is an instance of the ElasticBuffer class)
# self.pp_set_config(num_max_tensor_bytes=1024, num_max_inflight_tensors=5)
```

### Response
#### Success Response (None)
This method does not return any value.

#### Response Example
None
```

--------------------------------

### Training Forward Pass with ElasticBuffer

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/README.md

Demonstrates the typical forward pass in a training scenario using ElasticBuffer. It shows how to dispatch data, overlap computation with communication, and combine results.

```python
from deep_ep import ElasticBuffer

buffer = ElasticBuffer(group, num_max_tokens_per_rank=8192, hidden=7168, num_topk=8)

# Dispatch
recv_x, recv_topk_idx, recv_topk_weights, handle, dispatch_event = buffer.dispatch(
    x, topk_idx=topk_idx, topk_weights=topk_weights,
    num_experts=256, num_max_tokens_per_rank=8192,
    async_with_compute_stream=True
)

# Overlap computation
with dispatch_event:
    expert_output = expert_layer(recv_x)

# Combine
combined_x, _, combine_event = buffer.combine(
    expert_output, handle=handle,
    topk_weights=recv_topk_weights,
    async_with_compute_stream=True
)

with combine_event:
    final_output = post_process(combined_x)
```

--------------------------------

### Set NVSHMEM Environment Variables

Source: https://github.com/deepseek-ai/deepep/blob/main/docs/nvshmem.md

Set essential environment variables for NVSHMEM when not using RPM or deb packages. This ensures DeepEP can locate and use the NVSHMEM installation.

```bash
export NVSHMEM_DIR=/path/to/your/dir/to/install  # Use for DeepEP installation
export LD_LIBRARY_PATH="${NVSHMEM_DIR}/lib:$LD_LIBRARY_PATH"
export PATH="${NVSHMEM_DIR}/bin:$PATH"
```

--------------------------------

### Computing with Handle Metadata

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/ep-handle.md

Utilize handle metadata to dynamically allocate memory for expert GEMM output and process tokens per expert. This allows for efficient computation by leveraging information about token distribution and expand mode.

```python
# Use handle information to allocate expert GEMM output
num_recv_tokens = handle.psum_num_recv_tokens_per_scaleup_rank[-1].item()
expert_output = torch.empty(num_recv_tokens, hidden, dtype=torch.bfloat16, device='cuda')

# Process tokens per expert using metadata
for expert_id in range(num_local_experts):
    if expert_id == 0:
        start = 0
    else:
        # Recover real count from expand mode metadata
        start = handle.psum_num_recv_tokens_per_expert[expert_id - 1].item()
    
    end = handle.psum_num_recv_tokens_per_expert[expert_id].item()
    
    # Process tokens for this expert
    if end > start:
        tokens = recv_x[start:end]
        expert_output[start:end] = expert_layer(tokens)
```

--------------------------------

### Include Directories Configuration

Source: https://github.com/deepseek-ai/deepep/blob/main/CMakeLists.txt

Adds necessary include directories for project headers, third-party libraries, CUDA, PyTorch, NVSHMEM, and NCCL.

```cmake
include_directories(deep_ep/include third-party/fmt/include)
include_directories(${CUDA_TOOLKIT_ROOT_DIR}/include ${CUDA_TOOLKIT_ROOT_DIR}/include/cccl ${TORCH_INCLUDE_DIRS} ${PYTHON_INCLUDE_DIRS} ${NVSHMEM_INCLUDE_DIR} ${NCCL_ROOT_DIR}/include .)
```

--------------------------------

### Get Dispatch Performance Config

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/legacy-buffer.md

Retrieve a recommended performance configuration for dispatch operations based on the number of ranks. The configuration includes parameters like num_sms, and other tuning values.

```python
@staticmethod
def get_dispatch_config(num_ranks: int) -> Config
```

--------------------------------

### Get Local Buffer Tensor from Legacy Buffer

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/legacy-buffer.md

Obtain the raw buffer as a PyTorch tensor. You can specify the data type, size, offset, and whether to use an RDMA buffer.

```python
buffer_tensor = buffer.get_local_buffer_tensor(
    dtype=torch.bfloat16,
    size=torch.Size([8192, 7168]),
    use_rdma_buffer=False
)
```

--------------------------------

### Provide ProcessGroup to Buffer

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/errors.md

Initialize a `Buffer` by providing a PyTorch distributed `ProcessGroup`. This is required when neither a group nor a communicator is specified, particularly for legacy Buffer usage.

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

# Provide ProcessGroup
dist.init_process_group(backend='nccl')
group = dist.new_group()
buffer = Buffer(group=group, num_nvl_bytes=size)
```

--------------------------------

### Get NCCL Communicator Handle

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/utilities.md

Retrieves a cached NCCL communicator handle for a given process group. It prioritizes reusing existing communicators and PyTorch's NCCL comm if configured.

```python
from deep_ep.utils.comm import get_nccl_comm_handle

handle = get_nccl_comm_handle(group)
nccl_comm = handle.get()
```

--------------------------------

### Provide MPI Communicator to Buffer

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/errors.md

Initialize a `Buffer` by providing an MPI communicator. This is an alternative to providing a `ProcessGroup` when neither is specified, especially for legacy Buffer usage.

```python
from deep_ep import Buffer
from mpi4py import MPI

# Or provide MPI Comm
buffer = Buffer(comm=MPI.COMM_WORLD, num_nvl_bytes=size)
```

--------------------------------

### Basic Dispatch-Compute-Combine Pattern

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/event-overlap.md

Demonstrates overlapping communication dispatch with computation and combining results. Use `async_with_compute_stream=True` to prevent blocking.

```python
# Forward pass with overlap
recv_x, recv_topk_idx, recv_topk_weights, handle, event = buffer.dispatch(
    x=input_tokens,
    topk_idx=topk_idx,
    topk_weights=topk_weights,
    num_experts=256,
    num_max_tokens_per_rank=8192,
    async_with_compute_stream=True  # Don't wait in dispatch
)

# Overlap computation with communication
with event:
    # This runs while dispatch communication is in-flight
    expert_output = expert_layer(recv_x)

# By now, communication is guaranteed finished
combined_x, _, event = buffer.combine(
    expert_output,
    handle=handle,
    async_with_compute_stream=True
)

# Overlap final reduction with next layer
with event:
    output = post_processing(combined_x)
```

--------------------------------

### Get Engram Storage Size Hint

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/elastic-buffer.md

Determines the minimum buffer size for Engram (remote KV cache) storage. Returns recommended GPU and CPU bytes, both 2 MB-aligned. This is an experimental feature.

```python
import torch
from typing import Tuple

# Assuming ElasticBuffer is defined elsewhere
# Example usage:
# gpu_bytes, cpu_bytes = ElasticBuffer.get_engram_storage_size_hint(
#     num_entries=1000000, hidden=7168, num_max_tokens_per_rank=8192
# )
```

--------------------------------

### Config

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/types.md

Configuration object for tuning dispatch/combine kernel performance, specifically for the legacy Buffer dispatch/combine functionality.

```APIDOC
## Config

### Description
Performance tuning configuration for legacy Buffer dispatch/combine.

### Class Definition
```python
class Config:
    """Performance tuning configuration for legacy Buffer dispatch/combine."""
    
    def __init__(self, num_sms: int, param1: int, param2: int, param3: int, param4: int):
        self.num_sms = num_sms
        self.param1 = param1
        self.param2 = param2
        self.param3 = param3
        self.param4 = param4
```

### Module
`deep_ep._C`

### Fields
| Field | Type | Description |
|-------|------|-------------|
| num_sms | `int` | Number of SMs to use |
| param1 | `int` | Warp blocks per SM |
| param2 | `int` | Threads per block |
| param3 | `int` | Block size |
| param4 | `int` | Unused/reserved |

### Usage
Used by: `Buffer.get_dispatch_config()`, `Buffer.get_combine_config()`

### Example
```python
from deep_ep import Config, Buffer

# Get recommended config
config = Buffer.get_dispatch_config(num_ranks=8)

# Or create custom config
custom_config = Config(
    num_sms=20,
    param1=6,
    param2=256,
    param3=6,
    param4=128
)

# Use in dispatch (assuming 'buffer' and 'x', 'topk_idx' are defined)
# buffer.dispatch(
#     x, handle=None, topk_idx=topk_idx,
#     config=custom_config
# )
```
```

--------------------------------

### dispatch()

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/elastic-buffer.md

Routes input tokens to different ranks based on expert indices. This method is crucial for distributed expert parallelism, allowing tokens to be sent to the appropriate experts on other ranks.

```APIDOC
## dispatch()

### Description
Routes tokens to different ranks based on expert indices.

### Method
`dispatch`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

#### Method Parameters
- **x** (`Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]`) - Required - Input tokens `[num_tokens, hidden]` or FP8 tuple `(data, scales)`
- **topk_idx** (`Optional[torch.Tensor]`) - Optional - Expert indices `[num_tokens, num_topk]` with type `topk_idx_t`
- **topk_weights** (`Optional[torch.Tensor]`) - Optional - Expert weights `[num_tokens, num_topk]` with type `torch.float`
- **cumulative_local_expert_recv_stats** (`Optional[torch.Tensor]`) - Optional - Cumulative expert count for statistics `[num_local_experts]`
- **num_experts** (`Optional[int]`) - Optional - Total number of experts (inferred from handle if provided)
- **num_max_tokens_per_rank** (`Optional[int]`) - Optional - Maximum tokens per rank (from constructor default if None)
- **expert_alignment** (`Optional[int]`) - Optional - Token count alignment per expert
- **num_sms** (`int`) - Optional - Number of SMs (0 for automatic)
- **num_qps** (`int`) - Optional - Number of QPs (0 for automatic)
- **previous_event** (`Optional[EventHandle]`) - Optional - Event to wait before execution
- **async_with_compute_stream** (`bool`) - Optional - Don't wait for communication to finish
- **allocate_on_comm_stream** (`bool`) - Optional - Allocate output tensors on comm stream
- **handle** (`Optional[EPHandle]`) - Optional - Cached handle from previous dispatch (reuses layout)
- **do_handle_copy** (`bool`) - Optional - Clone `topk_idx` in returned handle
- **do_cpu_sync** (`Optional[bool]`) - Optional - Synchronize with CPU for exact token counts (default True)
- **do_expand** (`bool`) - Optional - Use expanding layout (one slot per token per expert)
- **use_tma_aligned_col_major_sf** (`bool`) - Optional - Use TMA-aligned column-major layout for scale factors

### Returns
- **recv_x** (`Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]`) - Received tokens (same type as input)
- **recv_topk_idx** (`Optional[torch.Tensor]`) - Received expert indices
- **recv_topk_weights** (`Optional[torch.Tensor]`) - Received expert weights (None if not provided)
- **handle** (`EPHandle`) - Communication handle for combine (cacheable)
- **event_overlap** (`EventOverlap`) - EventOverlap for async synchronization

### Example
```python
# Forward dispatch
x = torch.randn(8192, 7168, dtype=torch.bfloat16, device='cuda')
topk_idx = torch.randint(0, 256, (8192, 8), device='cuda')
topk_weights = torch.rand(8192, 8, device='cuda')

recv_x, recv_topk_idx, recv_topk_weights, handle, event = buffer.dispatch(
    x,
    topk_idx=topk_idx,
    topk_weights=topk_weights,
    num_experts=256,
    num_max_tokens_per_rank=8192,
    async_with_compute_stream=True
)

# Wait for communication if overlapping
with event:
    # Do local expert computation
    pass
```
```

--------------------------------

### get_low_latency_rdma_size_hint()

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/legacy-buffer.md

Provides a recommended RDMA buffer size for low-latency mode based on specified parameters.

```APIDOC
## get_low_latency_rdma_size_hint(num_max_dispatch_tokens_per_rank: int, hidden: int, num_ranks: int, num_experts: int)

### Description
Get recommended RDMA buffer size for low-latency mode.

### Method
`@staticmethod`

### Parameters
#### Path Parameters
- **num_max_dispatch_tokens_per_rank** (int) - Required - The maximum number of dispatch tokens per rank.
- **hidden** (int) - Required - The hidden dimension size.
- **num_ranks** (int) - Required - The total number of ranks.
- **num_experts** (int) - Required - The number of experts.

### Returns
`int` - Recommended buffer size in bytes.

### Request Example
```python
size = Buffer.get_low_latency_rdma_size_hint(
    num_max_dispatch_tokens_per_rank=8192,
    hidden=7168,
    num_ranks=8,
    num_experts=256
)
```
```

--------------------------------

### Set Even SM Count for Legacy Buffer

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/errors.md

Example of correctly setting the number of SMs for a legacy `Buffer` instance, ensuring the count is even to avoid the 'SM Count Must Be Even' assertion error. This constraint does not apply to `ElasticBuffer`.

```python
from deep_ep import Buffer

# Use even SM count only
Buffer.set_num_sms(20)  # OK
Buffer.set_num_sms(21)  # ERROR!
```

--------------------------------

### Trace Synchronization Points with EventOverlap

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/event-overlap.md

Shows how to use an EventOverlap object to synchronize streams and perform computations after communication is complete. It highlights the synchronization behavior with `with event:` and `torch.cuda.synchronize()`.

```python
from deep_ep import ElasticBuffer

buffer = ElasticBuffer(group, num_bytes=size)

recv_x, _, _, _, event = buffer.dispatch(
    x, topk_idx, topk_weights, num_experts,
    async_with_compute_stream=True
)

# Current stream is at point A (ahead of communication stream)

with event:
    # Point B: current stream is synchronized
    torch.cuda.synchronize()  # Redundant - already synced by event.current_stream_wait()
    expert_output = expert_layer(recv_x)

# Can immediately use expert_output
```

--------------------------------

### Buffer Constructor

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/legacy-buffer.md

Initializes a legacy communication buffer for expert-parallel operations. Supports configuration for NVLink and RDMA communication, low-latency mode, and other advanced settings.

```APIDOC
## Constructor: Buffer

### Description
Initializes a legacy communication buffer supporting high-throughput and low-latency expert-parallel operations.

### Method Signature
```python
def __init__(self,
             group: Optional[dist.ProcessGroup],
             num_nvl_bytes: int = 0,
             num_rdma_bytes: int = 0,
             low_latency_mode: bool = False,
             num_qps_per_rank: int = 24,
             allow_nvlink_for_low_latency_mode: bool = True,
             allow_mnnvl: bool = False,
             explicitly_destroy: bool = False,
             enable_shrink: bool = False,
             comm: Optional["mpi4py.MPI.Comm"] = None) -> None
```

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

### Parameters
- **group** (`Optional[dist.ProcessGroup]`) - Optional - PyTorch distributed process group (required if `comm` is None).
- **num_nvl_bytes** (`int`) - Optional - NVLink communication buffer size in bytes. Default: 0.
- **num_rdma_bytes** (`int`) - Optional - RDMA communication buffer size in bytes. Default: 0.
- **low_latency_mode** (`bool`) - Optional - Enable low-latency mode (uses RDMA instead of NVLink for intranode). Default: False.
- **num_qps_per_rank** (`int`) - Optional - Number of QPs per rank (must equal num_local_experts in low-latency mode). Default: 24.
- **allow_nvlink_for_low_latency_mode** (`bool`) - Optional - Allow NVLink traffic in low-latency mode. Default: True.
- **allow_mnnvl** (`bool`) - Optional - Allow multi-node NVLink. Default: False.
- **explicitly_destroy** (`bool`) - Optional - Require explicit `destroy()` call instead of destructor cleanup. Default: False.
- **enable_shrink** (`bool`) - Optional - Enable shrink mode with dynamic rank masking. Default: False.
- **comm** (`Optional[mpi4py.MPI.Comm]`) - Optional - MPI communicator (alternative to group).

### Request Example
```python
import torch.distributed as dist
from deep_ep import Buffer

# High-throughput mode
dist.init_process_group(backend='nccl')
group = dist.new_group(list(range(dist.get_world_size())))

buffer = Buffer(
    group,
    num_nvl_bytes=2 * 1024 * 1024 * 1024,  # 2 GB for NVLink
    num_rdma_bytes=4 * 1024 * 1024 * 1024  # 4 GB for RDMA
)

# Low-latency mode
buffer_ll = Buffer(
    group,
    num_rdma_bytes=4 * 1024 * 1024 * 1024,
    low_latency_mode=True,
    num_qps_per_rank=8  # Must equal number of local experts
)
```

### Response
None (Constructor does not return a value)

```

--------------------------------

### Prefix Sum Interpretation for Expert Token Counts

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/ep-handle.md

This code snippet shows how to calculate the token count for a specific expert based on whether the handle is in expand mode or non-expand mode. It uses prefix sums to determine the start and end points of an expert's tokens.

```python
if handle.do_expand:
    # Expand mode: psum[i] points to the END of expert i's tokens
    start = 0 if i == 0 else \
        align(handle.psum_num_recv_tokens_per_expert[i-1], expert_alignment)
    end = handle.psum_num_recv_tokens_per_expert[i]
    count = end - start
else:
    # Non-expand mode: standard inclusive prefix sum
    start = 0 if i == 0 else handle.psum_num_recv_tokens_per_expert[i-1]
    end = handle.psum_num_recv_tokens_per_expert[i]
    count = end - start
```

--------------------------------

### API Reference - EPHandle

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/COMPLETION_SUMMARY.txt

Documentation for EPHandle, which provides routing metadata for expert parallelism. It details attribute documentation, usage patterns for caching and debugging, and explanations of expand vs. non-expand modes.

```APIDOC
## EPHandle (Routing Metadata)

### Description
EPHandle is a component used for managing routing metadata within the expert parallelism framework. It facilitates efficient data dispatch and aggregation by providing essential routing information.

### Attributes
- **attribute1** (type) - Description of attribute 1
- **attribute2** (type) - Description of attribute 2
... (all attributes documented)

### Usage Patterns
- Caching strategies for frequently accessed metadata.
- Debugging techniques using EPHandle information.
- Explanation of 'expand' vs. 'non-expand' modes for routing.
- Guide to interpreting and utilizing metadata.
```

--------------------------------

### Optimize ElasticBuffer Configuration

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/errors.md

Optimize ElasticBuffer configuration by increasing `num_allocated_qps` for more parallelism or enabling `allow_hybrid_mode`.

```python
buffer = ElasticBuffer(
    group,
    num_bytes=size,
    num_allocated_qps=256,  # More parallelism
    allow_hybrid_mode=True
)
```

--------------------------------

### Initialize Default PyTorch Distributed Process Group

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/types.md

Initializes a default PyTorch distributed process group using the 'nccl' backend, typically involving all available ranks.

```python
# Default group (all ranks)
dist.init_process_group(backend='nccl')
```

--------------------------------

### Create Python Module with Pybind11

Source: https://github.com/deepseek-ai/deepep/blob/main/CMakeLists.txt

Builds a Python extension module named '_C' using pybind11, linking it with necessary libraries.

```cmake
# Link CPP and CUDA together
pybind11_add_module(_C csrc/python_api.cpp)
target_link_libraries(_C PRIVATE nccl ${RUNTIME_CUDA_LIBRARIES} ${LEGACY_CUDA_LIBRARIES} ${TORCH_LIBRARIES} torch_python)
```

--------------------------------

### Initialize Random Seeds

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/utilities.md

Initializes random seeds for all processes to ensure reproducibility. The local seed for each process is derived from the global seed and its rank.

```python
from deep_ep.utils.envs import init_seed

init_seed(42)
```

--------------------------------

### Library Directories Configuration

Source: https://github.com/deepseek-ai/deepep/blob/main/CMakeLists.txt

Specifies directories where libraries like PyTorch, CUDA, NVSHMEM, and NCCL can be found.

```cmake
link_directories(${TORCH_INSTALL_PREFIX}/lib ${CUDA_TOOLKIT_ROOT_DIR}/lib ${NVSHMEM_LIB_DIR} ${NCCL_ROOT_DIR}/lib)
```

--------------------------------

### Configure Pipeline-Parallel Parameters

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/elastic-buffer.md

Configure pipeline-parallel parameters using this experimental method. Specify the maximum number of tensor bytes and inflight tensors.

```python
def pp_set_config(self, num_max_tensor_bytes: int, 
                  num_max_inflight_tensors: int):
    pass
```

--------------------------------

### Define Config Class for Performance Tuning

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/types.md

Configuration object for tuning dispatch/combine kernel performance for legacy Buffers. Requires specifying number of SMs and other tuning parameters.

```python
class Config:
    """Performance tuning configuration for legacy Buffer dispatch/combine."""
    
    def __init__(self, num_sms: int, param1: int, param2: int, param3: int, param4: int):
        self.num_sms = num_sms
        self.param1 = param1
        self.param2 = param2
        self.param3 = param3
        self.param4 = param4
```

--------------------------------

### Check PyTorch Deterministic Algorithms and Memory Initialization

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/errors.md

Check the status of PyTorch's deterministic algorithms and the `fill_uninitialized_memory` setting. This code snippet is used to diagnose potential race conditions when both are enabled.

```python
# Check fails if both are True:
torch.are_deterministic_algorithms_enabled()
torch.utils.deterministic.fill_uninitialized_memory
```

--------------------------------

### Manual Synchronization with CUDA Streams

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/api-reference/event-overlap.md

Demonstrates manual synchronization using `event.current_stream_wait()` to ensure communication completion before proceeding. This is useful when operations need to be explicitly synchronized or moved to different CUDA streams.

```python
recv_x, _, _, _, event_overlap = buffer.dispatch(...)

# Do some computation in parallel
intermediate = some_compute(recv_x)

# Explicitly wait when needed
event_overlap.current_stream_wait()

# Now safe to use recv_x in non-recorded streams
other_stream = torch.cuda.Stream()
with torch.cuda.stream(other_stream):
    result = process_in_parallel(intermediate)
```

--------------------------------

### API Reference - LegacyBuffer V1

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/COMPLETION_SUMMARY.txt

Documentation for the LegacyBuffer V1 class, maintained for backward compatibility. It details its configuration, high-throughput/low-latency modes, and guidance for migrating to V2.

```APIDOC
## LegacyBuffer V1 (Backward Compatibility)

### Description
LegacyBuffer V1 is provided for backward compatibility with older versions of DeepEP. It offers similar functionality to the modern ElasticBuffer but may have different performance characteristics and configuration options.

### Features
- Full backward compatibility documentation.
- Support for high-throughput and low-latency modes.
- Guidance on migrating from V1 to V2.
- Explanation of configuration differences between V1 and V2.
```

--------------------------------

### Calculate Theoretical SM and QP for ElasticBuffer

Source: https://github.com/deepseek-ai/deepep/blob/main/_autodocs/README.md

Use ElasticBuffer methods to calculate the theoretical number of Streaming Multiprocessors (SMs) and Queue Pairs (QPs) based on buffer size and other parameters. This helps in capacity planning and performance estimation.

```python
buffer = ElasticBuffer(group, num_bytes=size)
num_sms = buffer.get_theoretical_num_sms(256, 8)
num_qps = buffer.get_theoretical_num_qps(num_sms)
```

--------------------------------

### MoE Dispatch and Combine for Training/Inference with ElasticBuffer

Source: https://github.com/deepseek-ai/deepep/blob/main/README.md

These functions demonstrate the dispatch and combine operations using ElasticBuffer for MoE models. They support both forward and backward passes and handle different input types like BF16 and FP8. Ensure ElasticBuffer and communication parameters are globally configured.

```python
import torch
import torch.distributed as dist
from typing import Tuple, Union

from deep_ep import ElasticBuffer, EPHandle, EventOverlap


def dispatch_forward(x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
                     topk_idx: torch.Tensor, topk_weights: torch.Tensor,
                     num_experts: int,
                     num_max_tokens_per_rank: int,
                     expert_alignment: int = 1) -> \
        Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
              torch.Tensor, torch.Tensor, EPHandle, EventOverlap]:
    """
    MoE dispatch: route tokens to the corresponding experts across all ranks.
    Supports both BF16 and FP8 (x as a tuple of [data, scale_factors]) inputs.
    """
    global _buffer, _num_comm_sms

    recv_x, recv_topk_idx, recv_topk_weights, handle, event = _buffer.dispatch(
        x,
        topk_idx=topk_idx,
        topk_weights=topk_weights,
        num_experts=num_experts,
        num_max_tokens_per_rank=num_max_tokens_per_rank,
        expert_alignment=expert_alignment,
        num_sms=_num_comm_sms,
        async_with_compute_stream=True,
    )

    # `handle` contains routing metadata for the subsequent combine call
    # `handle.num_recv_tokens_per_expert_list` provides per-expert token counts for GEMM
    # Use `event.current_stream_wait()` to synchronize the compute stream before using results
    return recv_x, recv_topk_idx, recv_topk_weights, handle, event


def dispatch_backward(grad_recv_x: torch.Tensor,
                      grad_recv_topk_weights: torch.Tensor,
                      handle: EPHandle) -> Tuple[torch.Tensor, torch.Tensor, EventOverlap]:
    """The backward pass of MoE dispatch is actually a combine."""
    global _buffer, _num_comm_sms

    combined_grad_x, combined_grad_topk_weights, event = _buffer.combine(
        grad_recv_x,
        handle=handle,
        topk_weights=grad_recv_topk_weights,
        num_sms=_num_comm_sms,
        async_with_compute_stream=True,
    )

    return combined_grad_x, combined_grad_topk_weights, event


def combine_forward(x: torch.Tensor,
                    handle: EPHandle) -> Tuple[torch.Tensor, EventOverlap]:
    """MoE combine: reduce expert outputs back to their original ranks."""
    global _buffer, _num_comm_sms

    combined_x, _, event = _buffer.combine(
        x,
        handle=handle,
        num_sms=_num_comm_sms,
        async_with_compute_stream=True,
    )

    return combined_x, event


def combine_backward(grad_combined_x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
                     handle: EPHandle) -> \
        Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], EventOverlap]:
    """The backward pass of MoE combine is actually a dispatch."""
    global _buffer, _num_comm_sms

    grad_x, _, _, _, event = _buffer.dispatch(
        grad_combined_x,
        handle=handle,
        num_sms=_num_comm_sms,
        async_with_compute_stream=True,
    )

    return grad_x, event

```