[2/N] Elastic EP Milestone 2: Integrating NIXL-EP

[itayalroy]

This PR is a rebase of #29630 originally authored by @libertyeagle that integrates NIXL-EP kernels into vLLM.
NIXL-EP is an implementation of expert-parallel communication kernels over NIXL's device API. It provides elastic scaling capabilities, enabling dynamic addition and removal of processes (ranks) during runtime, without the need to destroy and recreate communicators during scaling up/down.

This PR also includes a few small fixes to vLLM Elastic EP (#34861) that we found while thoroughly testing vLLM with NIXL-EP.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors.

You ask your reviewers to trigger select CI tests on top of fastcheck CI.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

[mergify]

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

uv pip install pre-commit
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 or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

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

[gemini-code-assist]

Code Review

This pull request integrates NIXL-EP kernels for elastic expert parallelism, which is a significant enhancement. The changes are mostly about adding the new nixl_ep backend and its related logic. The implementation also includes some important fixes, such as properly destroying NCCL communicators to prevent resource leaks.

My review focuses on ensuring thread safety in the new NixlEPAll2AllManager. I've identified a potential race condition when accessing the shared buffer and suggested a fix using a lock to ensure robustness in a multi-threaded environment.

[gemini-code-assist]

The _buffer class attribute is a shared mutable state. The get_handle method reads and writes to this shared state without any synchronization, which can lead to a race condition if called from multiple threads concurrently. This could happen, for example, during dynamic LoRA loading, leading to incorrect behavior or crashes.

To prevent this, a lock should be used to protect access to _buffer.

First, please add a lock to the NixlEPAll2AllManager class:

class NixlEPAll2AllManager(All2AllManagerBase):
    ...
    _lock = threading.Lock()
    ...

Then, wrap the get_handle method's logic with this lock as suggested below.

    def get_handle(self, kwargs):
        with NixlEPAll2AllManager._lock:
            if (
                NixlEPAll2AllManager._buffer is not None
                and NixlEPAll2AllManager._buffer[1] == self.cpu_group.size()
            ):
                return NixlEPAll2AllManager._buffer[0]

            num_experts_per_rank = kwargs["num_global_experts"] // kwargs["num_ep_ranks"]
            nixl_kwargs = dict(
                max_num_tokens_per_dp_rank=kwargs["max_num_tokens_per_dp_rank"],
                token_hidden_size=kwargs["token_hidden_size"],
                num_experts_per_rank=num_experts_per_rank,
            )
            if NixlEPAll2AllManager._buffer is None:
                self._init_buffer(**nixl_kwargs)
            else:
                self._update_buffer()

            assert NixlEPAll2AllManager._buffer is not None
            handle = NixlEPAll2AllManager._buffer[0]
            return handle

[tlrmchlsmth]

@itayalroy could you address this comment?

[itayalroy]

I don't think this can actually happen, since get_handle() only appears to be called from a single thread during initial setup or elastic EP reconfiguration. In any case, this isn't on the data path, the cost of adding a lock here is negligible, so I added it to be safe.

[mergify]

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

uv pip install pre-commit
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 or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

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

[tlrmchlsmth]

One quesiont: Does NIXL-EP use NVLINK at all for intranode traffic? Is it suitable for MNNVL systems? And are there any de-duplication optimizations?

[tlrmchlsmth]

@itayalroy could you address this comment?

[tlrmchlsmth]

this file is extremely similar to the DeepEP LL prepare_finalize implementation. Should we consolidate these?

@bnellnm WDYT?

[bnellnm]

It might be good to factor out some of the common utilities, e.g. dequant_fp8, maybe_roundup_layer_hidden_size, (maybe _do_quant?) but I think it might be good to keep the main implementations separate in case one or the other of the backends changes their API.

[itayalroy]

We expect nixl_ep_prepare_finalize.py and deepep_ll_prepare_finalize.py to diverge pretty quickly as NIXL-EP progresses, and possibly on the DeepEP side too, so preferred to keep them separate

[tlrmchlsmth]

OK, fair enough!

[itayalroy]

One quesiont: Does NIXL-EP use NVLINK at all for intranode traffic? Is it suitable for MNNVL systems? And are there any de-duplication optimizations?

NVLink is used for intranode traffic, support for MNNVL is in review, and currently no de-duplication optimizations (although this might change, we are still evaluating the tradeoffs of that approach)

[tlrmchlsmth]

Last thing missing is better integration testing, which is tricky given the heavy dependency. From @itayalroy we should be able to pip install nixl in either the next NIXL release or the one after, which will make this easier to manage in the test image