[WideEP] Integrate DeepEP v2

[tlrmchlsmth]

Integrating deepseek-ai/DeepEP#605

Unit tests are passing, and it is working e2e. GSM8k appears to be good and will follow up with more thorough e2e tests.

Notes:

• I couldn't get this working on an 8xB200 system as DeepEP v2's ElasticBuffer unconditionally asserts NCCL GIN availability even for intra-node NVLink-only. This is a TODO.
• This requires NCCL 2.30.4, and PyTorch pins to NCCL 2.28.9, so for now this requires users to manually install NCCL after installing torch via uv pip install "nvidia-nccl-cu13>=2.30.4"

I'm using this Containerfile for now https://github.com/tlrmchlsmth/j-llm-d/blob/deepep-v2-dev-image/dev/Containerfile.deepep-v2

[gemini-code-assist]

Code Review

This pull request adds support for the DeepEP v2 (ElasticBuffer) all2all backend, including a new DeepEPV2PrepareAndFinalize implementation for MoE kernels, a dedicated All2AllManager, and associated configuration and environment variables. A comprehensive test suite for DeepEP v2 MoE is also introduced. Feedback identifies a critical issue where a strict bfloat16 assertion in the finalization logic would cause crashes for float16 models, recommending a cast to bfloat16 instead to maintain compatibility.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

Hi @tlrmchlsmth, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

Hi @tlrmchlsmth, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

Hi @tlrmchlsmth, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[mergify]

Hi @tlrmchlsmth, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.