Repository navigation
feat: add deep-ep-v2 wheel (DeepEP V2.5 as deep_ep_v2) - #18
Closed
faresobeid wants to merge 3 commits into
Closed
faresobeid wants to merge 3 commits into
faresobeid wants to merge 3 commits into
Conversation
…image's glibc headers
… so it imports next to deep_ep
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds a
deep-ep-v2wheel to the kernel release, built from deepseek-ai/DeepEP @ 93eb6eb (V2.5, DeepJIT submodule cba802d).V2.5 is a rewrite on top of NCCL's device/GIN API (no NVSHMEM). In an EP8 trial in prime-rl (8k tokens/rank) it cut dispatch from 3.2 to 1.6 ms and combine from 5.8 to 3.7 ms vs V1. It ships as a separate package so it can sit next to V1 while we compare them:
deep-ep-v2, importdeep_ep_v2, extensiondeep_ep_v2._C. The script renames the Python package and itsdeep_ep._Cimports.-Ddeep_ep=deep_ep_v2. pybind11 registers types by C++ type, and V1 already registersdeep_ep::EventHandleetc., so without thisimport deep_ep, deep_ep_v2fails withgeneric_type: type "EventHandle" is already registered!. The macro does not touch#include <deep_ep/...>or the JIT's kernel source strings (deep_ep::ep::...);deep_epis only stringified in error messages.include/deep_ep/and the namespaces in the JIT-compiled kernels. The JIT includes<deep_ep/...>relative to<package>/include.2.5.0+93eb6eb.Build quirks:
TORCH_CUDA_ARCH_LISTand there are no cubins in the wheel. The host build works with the image's CUDA 13.0.nvidia-nccl-cu13. V2.5 needs >= 2.32.3, but torch 2.13 pins 2.29.7, so the script upgrades it in the build venv.setup.pybakes in an absolute RPATH to the build venv's NCCL. The script repoints it to$ORIGIN/../nvidia/nccl/lib, likeinstall_ep_kernels.shdoes for NVSHMEM.elfutils/libdwfl.h(elfutils-develadded to the dnf line). libdw is dlopened at runtime and is optional (it is used for C++ stack traces).setup.pyhard-codes/usr/local/cuda/include/cccl, so the script setsCPATH. The DeepJIT submodule URL is ssh, so it is rewritten to https.override-dependencieswithnvidia-nccl-cu13>=2.32.3.import deep_ep_v2fails (check_nccl_so) if the NCCL that torch loads is not the one it links.CUDA_HOMEfor the JIT (e.g. pipnvidia-cuda-nvcc/crt/nvvm/runtime/cccl13.x installed into a dedicated dir). Cubins from nvcc 13.3 run on the 13.0 driver.Smoke test: import (this runs
check_nccl_so+init_jit), check that the loaded libnccl is >= 2.32.3, and build the JIT runtime for the GPU (_C.get_jit()). Real dispatch/combine needs an EP group and a 13.1+ toolkit, so it is out of scope. The smoke job installs nccl 2.32.3 for this kernel.Validated (built on the login node: Ubuntu 24.04, CUDA 13.0 toolkit, torch 2.13.0+cu130):
.so). RUNPATH is$ORIGIN/../nvidia/nccl/lib. NEEDED has no libcuda (the driver is loaded lazily).import deep_ep_v2passes, includingcheck_nccl_so./proc/self/mapsshows one libnccl (pip 2.32.3) and one libcudart (torch's pipnvidia/cu13).deep_ep-1.2.1+29d31c0: no shared paths (deep_ep/,deep_ep_cpp*.sovsdeep_ep_v2/).nm -DC: 145 exporteddeep_ep_v2::symbols, 0deep_ep::; the pybind type info is nowN10deep_ep_v211EventHandleE.deep_ep1.2.1 and nccl 2.32.3:import torch, deep_ep, deep_ep_v2works, andscripts/smoke_test.py deep-ep-v2passes (get_jit()builds the runtime even withCUDA_HOMEpointing at a 12.9 toolkit).Validated since (the wheel from this branch, in prime-rl PR PrimeIntellect-ai/prime-rl#3999 on H200):
Not validated:
if: false), so it was only run by hand on H200.deepep_v2across nodes.prime-rl will pin
deep-ep-v2after a release that includes this PR.