deepseek-ai / deepseek-ai/DeepEP
[Proposal] Integrate some enhancements from TensorRT-LLM
- Dominant language
- Cuda
- Stars
- 10.1k
- Forks
- 1.4k
- Avg merge
- 4d 1h
- Merged PRs (30d)
- 2
Description
Hi, DeepEP team,
I'm a collaborator to the TensorRT-LLM project. First, thank you for creating such a performant library. We have successfully integrated it into our project. To better align with TensorRT-LLM's data format, we've made a few modifications to DeepEP. The specific commit we are using is [515a311](https://github.com/deepseek-ai/DeepEP/commits/515a311f290eb6d9592fcccfcc80c40f5123ca72/). The main changes are as follows:
(1) Besides `torch.distributed.ProcessGroup`, support for initializing buffers with `mpi4py.MPI.Comm` ( @yuantailing )
(2) Change the dtype of `topk_idx` from `int64_t` to `int` and add an offset to expert id ( @yifeizhang-c )
(3) Avoid cleaning after each change in hidden_size/token_num (Already merged in #313, @yilin-void )
(4) Add an FP4 low-latency dispatch kernel, where the tensor is quantized outside of the kernel ( @yilin-void )
(5) Add an FP4 low-latency combine kernel, where the tensor is quant/dequant inside of the kernel ( @yilin-void )
**Before we create pull requests, we would like to discuss whether the community would be open to these changes and how we can make them more general, especially for (2), (4), and (5).**
For point (2), our proposal is to add a conditional type definition, such as `using topk_idx_t = int64_t` or `int` which depends on a macro.
For point (4), the FP4 dispatch is currently a separate kernel, but we could merge it into the original dispatch kernel by adding an input tensor to indicate that the scale is quantized outside of the kernel.
For point (5), the FP4 combine is currently a separate kernel, but we could merge it into the original dispatch kernel by adding an template argument to indicate whether quant/dequant should be performance before/after the communication.
We hope these contributions will broaden DeepEP's applications and allow us to stay aligned with the main branch of your library.
Thanks
Tailing
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.