Skip to content

feat: add deep-ep-v2 wheel (DeepEP V2.5 as deep_ep_v2) - #18

Closed
faresobeid wants to merge 3 commits into
mainfrom
feat/deep-ep-v2
Closed

faresobeid wants to merge 3 commits into
mainfrom
feat/deep-ep-v2

Conversation

@faresobeid

@faresobeid faresobeid commented Oct 10, 2026 •

Copy link
Copy Markdown

Adds a deep-ep-v2 wheel 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:

  • Dist deep-ep-v2, import deep_ep_v2, extension deep_ep_v2._C. The script renames the Python package and its deep_ep._C imports.
  • The host extension is compiled with -Ddeep_ep=deep_ep_v2. pybind11 registers types by C++ type, and V1 already registers deep_ep::EventHandle etc., so without this import deep_ep, deep_ep_v2 fails with generic_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_ep is only stringified in error messages.
  • Unchanged: the header dir include/deep_ep/ and the namespaces in the JIT-compiled kernels. The JIT includes <deep_ep/...> relative to <package>/include.
  • Version stays upstream: 2.5.0+93eb6eb.

Build quirks:

  • The extension is host C++ only. All kernels are JIT-compiled at runtime by DeepJIT, so there is no TORCH_CUDA_ARCH_LIST and there are no cubins in the wheel. The host build works with the image's CUDA 13.0.
  • NCCL comes from the pip package 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.py bakes in an absolute RPATH to the build venv's NCCL. The script repoints it to $ORIGIN/../nvidia/nccl/lib, like install_ep_kernels.sh does for NVSHMEM.
  • DeepJIT needs elfutils/libdwfl.h (elfutils-devel added to the dnf line). libdw is dlopened at runtime and is optional (it is used for C++ stack traces).
  • setup.py hard-codes /usr/local/cuda/include/cccl, so the script sets CPATH. The DeepJIT submodule URL is ssh, so it is rewritten to https.
  • The wheel does not declare the NCCL requirement: declaring it would conflict with torch's pin. Consumers need override-dependencies with nvidia-nccl-cu13>=2.32.3. import deep_ep_v2 fails (check_nccl_so) if the NCCL that torch loads is not the one it links.
  • Runtime also needs a CUDA >= 13.1 toolkit at CUDA_HOME for the JIT (e.g. pip nvidia-cuda-nvcc/crt/nvvm/runtime/cccl 13.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):

  • The script builds the wheel (8.5 MB, one .so). RUNPATH is $ORIGIN/../nvidia/nccl/lib. NEEDED has no libcuda (the driver is loaded lazily).
  • In a fresh venv (torch + nvidia-nccl-cu13 2.32.3 + wheel), import deep_ep_v2 passes, including check_nccl_so. /proc/self/maps shows one libnccl (pip 2.32.3) and one libcudart (torch's pip nvidia/cu13).
  • Installs next to the V1 wheel deep_ep-1.2.1+29d31c0: no shared paths (deep_ep/, deep_ep_cpp*.so vs deep_ep_v2/).
  • nm -DC: 145 exported deep_ep_v2:: symbols, 0 deep_ep::; the pybind type info is now N10deep_ep_v211EventHandleE.
  • On an H200 node, in a prime-rl venv with V1 deep_ep 1.2.1 and nccl 2.32.3: import torch, deep_ep, deep_ep_v2 works, and scripts/smoke_test.py deep-ep-v2 passes (get_jit() builds the runtime even with CUDA_HOME pointing at a 12.9 toolkit).
  • Earlier trial (same commit, built by hand): upstream EP tests and an EP8/EP16 prime-rl dispatcher bench pass on H200.

Validated since (the wheel from this branch, in prime-rl PR PrimeIntellect-ai/prime-rl#3999 on H200):

  • The CI builds pass (x86_64 and aarch64, manylinux image).
  • JIT kernels compile and run from this wheel, including with the system CUDA 13.0 toolkit and an empty cache. EP8 dispatch/combine match the torch dispatcher in eager, compile and full AC.
  • GLM-4.5-Air EP8 training steps are ~5% faster than torch dispatch.

Not validated:

  • The GPU smoke job: it is disabled in this repo's CI (if: false), so it was only run by hand on H200.
  • sm_100 / sm_103: no Blackwell here. The JIT compiles for the running GPU, so the wheel itself has no per-arch code.
  • Multi-node: V2.5 is slower than V1 at EP16 on our fabric, so prime-rl rejects deepep_v2 across nodes.

prime-rl will pin deep-ep-v2 after a release that includes this PR.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant