刚刚，DeepSeek开源DeepEP通信库，千亿MoE训推颠覆级创新！FP8狂飙，带飞GPU

# Build and make symbolic links for SO filesNVSHMEM_DIR=/path/to/installed/nvshmem python setup.py build# You may modify the specific SO names according to your own platformln -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.py`# according to your own cluster settings, and launch into multiple nodes python tests/test_intranode.pypython tests/test_internode.pypython tests/test_low_latency.py

import torchimport torch.distributed as distfrom typing import List, Tuple, Optional, Union
from deep_ep import Buffer, EventOverlap
# Communication buffer (will allocate at runtime)_buffer: Optional[Buffer] = None
# Set the number of SMs to use# NOTES: this is a static variableBuffer.set_num_sms(24)

# You may call this function at the framework initializationdef get_buffer(group: dist.ProcessGroup, hidden_bytes: int) -> Buffer:    global _buffer
    # NOTES: you may also replace `get_*_config` with your auto-tuned results via all the tests    num_nvl_bytes, num_rdma_bytes = 0, 0    for config in (Buffer.get_dispatch_config(group.size()), Buffer.get_combine_config(group.size())):        num_nvl_bytes = max(config.get_nvl_buffer_size_hint(hidden_bytes, group.size()), num_nvl_bytes)        num_rdma_bytes = max(config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), num_rdma_bytes)
    # Allocate a buffer if not existed or not enough buffer size    # NOTES: the adaptive routing configuration of the network **must be off**    if _buffer is None or _buffer.group != group or _buffer.num_nvl_bytes < num_nvl_bytes or _buffer.num_rdma_bytes < num_rdma_bytes:        _buffer = Buffer(group, num_nvl_bytes, num_rdma_bytes)    return _buffer

def get_hidden_bytes(x: torch.Tensor) -> int:    t = x[0] if isinstance(x, tuple) else x    return t.size(1) * max(t.element_size(), 2)

def dispatch_forward(x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],                     topk_idx: torch.Tensor, topk_weights: torch.Tensor,                     num_experts: int, previous_event: Optional[EventOverlap] = None) -> \        Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], torch.Tensor, torch.Tensor, List, Tuple, EventOverlap]:    # NOTES: an optional `previous_event` means a CUDA event captured that you want to make it as a dependency     # of the dispatch kernel, it may be useful with communication-computation overlap. For more information, please    # refer to the docs of `Buffer.dispatch`    global _buffer
    # Calculate layout before actual dispatch    num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, previous_event = \        _buffer.get_dispatch_layout(topk_idx, num_experts,                                    previous_event=previous_event, async_finish=True,                                    allocate_on_comm_stream=previous_event is not None)    # Do MoE dispatch    # NOTES: the CPU will wait for GPU's signal to arrive, so this is not compatible with CUDA graph    # For more advanced usages, please refer to the docs of the `dispatch` function    recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, event = \        _buffer.dispatch(x, topk_idx=topk_idx, topk_weights=topk_weights,                         num_tokens_per_rank=num_tokens_per_rank, num_tokens_per_rdma_rank=num_tokens_per_rdma_rank,                         is_token_in_rank=is_token_in_rank, num_tokens_per_expert=num_tokens_per_expert,                         previous_event=previous_event, async_finish=True,                         allocate_on_comm_stream=True)    # For event management, please refer to the docs of the `EventOverlap` class    return recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, event

def dispatch_backward(grad_recv_x: torch.Tensor, grad_recv_topk_weights: torch.Tensor, handle: Tuple) -> \        Tuple[torch.Tensor, torch.Tensor, EventOverlap]:    global _buffer
    # The backward process of MoE dispatch is actually a combine    # For more advanced usages, please refer to the docs of the `combine` function    combined_grad_x, combined_grad_recv_topk_weights, event = \        _buffer.combine(grad_recv_x, handle, topk_weights=grad_recv_topk_weights, async_finish=True)
    # For event management, please refer to the docs of the `EventOverlap` class    return combined_grad_x, combined_grad_recv_topk_weights, event

def combine_forward(x: torch.Tensor, handle: Tuple, previous_event: Optional[EventOverlap] = None) -> \        Tuple[torch.Tensor, EventOverlap]:    global _buffer
    # Do MoE combine    # For more advanced usages, please refer to the docs of the `combine` function    combined_x, _, event = _buffer.combine(x, handle, async_finish=True, previous_event=previous_event,                                           allocate_on_comm_stream=previous_event is not None)
    # For event management, please refer to the docs of the `EventOverlap` class    return combined_x, event

def combine_backward(grad_combined_x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],                     handle: Tuple, previous_event: Optional[EventOverlap] = None) -> \        Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], EventOverlap]:    global _buffer
    # The backward process of MoE combine is actually a dispatch    # For more advanced usages, please refer to the docs of the `combine` function    grad_x, _, _, _, _, event = _buffer.dispatch(grad_combined_x, handle=handle, async_finish=True,                                                 previous_event=previous_event,                                                 allocate_on_comm_stream=previous_event is not None)
    # For event management, please refer to the docs of the `EventOverlap` class    return grad_x, event