Repository navigation
Nixl ep v2 - draft - #2409
Nixl ep v2 - draft#2409lishapira wants to merge 11 commits into
Conversation
|
👋 Hi lishapira! Thank you for contributing to ai-dynamo/nixl. Your PR reviewers will review your contribution then trigger the CI to test your changes. 🚀 |
|
/build |
1 similar comment
|
/build |
|
🤖 CI Triage Agent — TL;DR: The Full analysisSummary: Two of three "Build Wheel" stages (x86_64 and aarch64, node IDs 506/475) failed at Dockerfile.manylinux STEP 18/27 during Root cause: Missing include directory, not a missing file. The error is: reached via The header exists at The proxy device headers came in with the CPU-Proxy series (PR #2326 / #2327), while the EP-V2 build targets were added separately by Secondary, unrelated flake: the third wheel stage (node 462, aarch64) failed earlier in Implicated commit: [REDACTED:Hex High Entropy String] — Lior Shapira, "nixl_ep_v2: Fix build warnings and separate V2 tests" (targets introduced in b69639b, same author); the header dependency originates from Tomer Davidor's CPU-proxy work ( File: Suggested fix:
v2_inc_dirs = [
nixl_inc_dirs,
nixl_device_inc_dirs,
utils_inc_dirs, # <-- src/utils, for <device/proxy/proxy_protocol.h>
include_directories('csrc'),
include_directories('nixl_ep_v2/include'),
include_directories(v2_cccl_include_root),
torch_inc_dirs,
]
Related: PR #2326 (CPU Proxy part 5, introduces the device proxy headers), PR #2327, PR #2328; current PR #2409.
|
|
🤖 CI Triage Agent — TL;DR: The Full analysisSummary: Stage 574 ("Test CPP", debug/aarch64 variant) failed on 1 of 341 gtest cases: Root cause: A latent ordering race in the ETCD metadata backend, not a timeout.
Implicated commit: File: Suggested fix: Close the window instead of widening the poll:
Do not simply raise Related: #2349 (prior timeout bump for this same fixture's flakiness), #1930 (introduced the etcd metadata backend), #2130 (related etcd metadata-cleanup test gating) |
|
🤖 CI Triage Agent — TL;DR: The Full analysisSummary: Stage "Compiling NIXL EP Docker Image for DL" (node 161) failed: Root cause: Missing include directory, not a missing file. Both failing targets include Implicated commit: File: Suggested fix: Add v2_inc_dirs = [
nixl_inc_dirs,
nixl_device_inc_dirs,
utils_inc_dirs, # <device/proxy/proxy_protocol.h>
include_directories('csrc'),
include_directories('nixl_ep_v2/include'),
include_directories(v2_cccl_include_root),
torch_inc_dirs,
]Also extend the wheel header packaging at Related: #2326 (CPU Proxy part 5, introduced the device-side proxy include), #2225/#2226/#2223 (CPU proxy parts 1–3, added |
|
🤖 CI Triage Agent — TL;DR: The NIXL‑EP V2 CUDA targets fail to compile because Full analysisSummary: Both Root cause: Not infrastructure — a deterministic include-path gap, reproduced identically on both architectures. The include chain in the log is: That header lives at
This is a latent mismatch that surfaced when the proxy headers were pulled into Implicated commit: File: Suggested fix: Add v2_inc_dirs = [
nixl_inc_dirs,
nixl_device_inc_dirs,
utils_inc_dirs, # <device/proxy/proxy_protocol.h> lives under src/utils
include_directories('csrc'),
include_directories('nixl_ep_v2/include'),
include_directories(v2_cccl_include_root),
torch_inc_dirs,
]Two follow-ups worth doing in the same change:
Related: #2225, #2226, #2223 (CPU Proxy parts 1–3, which introduced
|
|
/build |
|
/build |
|
🤖 CI Triage Agent — TL;DR: The Full analysisSummary: Root cause: Timing-sensitive assertion against a slow etcd store, not a functional regression. Evidence from the stage log:
This is a flaky test, not a defect introduced by commit [REDACTED:Hex High Entropy String]. All other work in the build (Build, Test Python, Test Nixlbench, Test Rust, and five sibling Implicated commit: No functional regression implicated. The relevant test-timing code was last touched by File: Suggested fix:
Related: PR #2409 (
|
|
🤖 CI Triage Agent — TL;DR: The "Run CPP tests" stage aborted because Full analysisSummary: Stage 307 "Run CPP tests" failed at the final NIXL client-server test: Root cause: Not a hang or a timeout — the whole C++ stage ran continuously and this is a hard abort at the very last command. The log shows: The CI script only allocates one port (27903, via Implicated commit: unknown — latent defect, not introduced by File: Suggested fix: Stop the initiator from using the hard-coded 8888:
Related:
|
Import files byte-for-byte from DeepEP branch review/nixl_ep_v2 at
573bc1b703a849fe0dcd0a711cae4fd039b737a6, which already renames the
package and namespace deep_ep to nixl_ep_v2 and the remaining DeepEP
mentions.
Copied:
- csrc/python_api.cpp
- csrc/elastic/**
- csrc/indexing/main.cu
- csrc/jit/**
- csrc/kernels/backend/{nixl.cpp,nixl.hpp,nixl_device.cu,nixl_device.cuh}
- csrc/kernels/elastic/**
- csrc/utils/{event,format,hash,lazy_driver,lazy_init,system}.hpp
- nixl_ep_v2/**
- tests/elastic/**: the tests, the profilers and the compile probe
- tests/utils/test_gate.py
Not copied:
- CMakeLists.txt and the CMakeLists.txt files under csrc/
- csrc/utils/shared_memory.hpp
- third-party/fmt
- README.md, LICENSE, and figures/
- build.sh, develop.sh, format.sh, and install.sh
- setup.py, pyproject.toml, and requirements-lint.txt
- DeepEP repository metadata and configuration files
This commit intentionally contains no NIXL-specific adaptations or
formatting changes.
Signed-off-by: Lior Shapira <lishapira@nvidia.com>
Rename the 28 user-configurable EP_* environment variables to NIXL_EP_* across their Python and C++ reads, diagnostics, and docs. Keep compile-time EP_* macros and the EP_NIXL Store key unchanged. In particular, only the environment reads corresponding to EP_ADAPTIVE_PROFILE and EP_NUM_TOPK_IDX_BITS are renamed; their JIT -D flags and kernel macros retain the original spelling. Keep the four NIXL-specific knobs under NIXL_EP_NIXL_*: the first NIXL is the NIXL_EP component namespace, while the second identifies the NIXL transport. This avoids colliding with v1's differently configured NIXL_EP_NUM_CHANNELS. Signed-off-by: Lior Shapira <lishapira@nvidia.com>
Load the torch-versioned NIXL-EP V2 extension with the same optional RTLD_DEEPBIND handling as v1, and expose it as the package-local _C module before importing submodules that depend on it. Make all internal extension imports relative so the package loads once under its final nixl.runtime.ep name. During package initialization, select the extension matching the installed torch version, initialize the JIT, and reject a top-k index-width setting that disagrees with the compiled host extension. DeepEP formats generated JIT commands and source strings with its bundled fmt library. NIXL builds with C++20 and does not carry that third-party dependency, so expose std::format through the existing nixl_ep_v2::fmt namespace and keep the current fmt::format call sites unchanged. Signed-off-by: Lior Shapira <lishapira@nvidia.com>
find_pkgs.py: - Use the headers shipped inside the NIXL EP package by default. - Keep NIXL_EP_NIXL_ROOT_DIR and NIXL_EP_UCX_ROOT_DIR as explicit overrides. - Stop discovering arbitrary system NIXL and UCX installations, which could provide headers that do not match the wheel's libraries. - Check only the NIXL device header needed by the JIT, and warn when the system rdma-core header required for inter-node puts is missing. compiler.hpp: - Add the packaged NIXL and UCX header contents to the kernel cache key, so a wheel upgrade cannot reuse cubins built against older headers. - Add the libcu++ atomic-storage workaround required by nvcc 12.9. Signed-off-by: Lior Shapira <lishapira@nvidia.com>
meson.build: - Build runtime/ep when NIXL-EP is enabled. runtime/ep/meson.build: - Use DeepEP V2's setup.py as the reference for source files, CUDA flags, C++20, CCCL, and the NIXL and UCX dependencies. - Use V1's examples/device/ep/meson.build as the reference for PyTorch detection, Torch-versioned extension naming, rpaths, debug-build handling, the nvcc 12.9 workaround, and Python installation. - Check the required CUDA, UCX device API, CUDA IPC, PyTorch, pybind11, CCCL, and C++20 support, and skip V2 cleanly when unavailable. - Build the V2 extension and a non-installed probe that compiles all JIT kernel headers during the normal build. - Install the Python package, extension, and matching NIXL and UCX headers into the CUDA backend wheel. contrib/build-wheel.sh: - Keep the complete wheel built for the first requested Torch version. - For each additional Torch version, build a temporary wheel and copy its Torch-specific V2 extension into the complete wheel. - Keep V1 and V2 extensions in their separate package directories. - Skip the V2 merge safely when Meson did not build V2 because a required dependency was unavailable. pyproject.toml: - Add nvidia-ml-py, which ElasticBuffer uses for GPU topology discovery. runtime/ep/README.md: - Document the V2 overview, features, initialization, key APIs, requirements, build instructions, JIT behavior, environment variables, and local tests. - Mark the original nixl_ep implementation as temporarily available and planned for removal. Signed-off-by: Lior Shapira <lishapira@nvidia.com>
LICENSE and runtime/ep/LICENSE-DeepEP: - Attribute runtime/ep to deepseek-ai/DeepEP commit 01dc3aaac82068020353dce2c302e38153c0bfaa. - Include DeepEP's MIT license alongside NIXL's Apache 2.0 license. runtime/ep imported sources: - Add DeepSeek and NVIDIA copyright notices to every imported Python, C++, CUDA, and header file. - Mark the original DeepEP material as MIT licensed and NVIDIA's modifications as Apache 2.0 using the combined MIT AND Apache-2.0 SPDX identifier. CONTRIBUTING.md and the contribution guide: - State that contributions modifying the DeepEP-derived V1 and V2 implementations must be licensed under Apache 2.0. .github/workflows/clang-format.yml: - Exclude runtime/ep from NIXL's clang-format check, matching the existing V1 exclusion. - Preserve the imported formatting so upstream comparisons and future re-syncs remain focused on functional changes. This commit changes licensing, attribution, and formatting policy only; it does not change runtime or build behavior. Signed-off-by: Lior Shapira <lishapira@nvidia.com>
- Add a lazy nixl.runtime package that routes submodule imports to the selected CUDA backend, and include it in the meta wheel. - Update the V2 tests, profilers and their helpers to import the installed nixl.runtime.ep package instead of the source package nixl_ep_v2. - Run the V2 correctness, elastic-membership, and topology-mode tests in the existing NIXL-EP test job. nixl/runtime/__init__.py: - Route nixl.runtime submodule lookup to the runtime directory of the selected nixl_cu12 or nixl_cu13 backend. - Keep V2 lazy: importing nixl does not load NIXL-EP V2; it is loaded only when the user imports nixl.runtime.ep. Meta-package Meson files: - Add nixl/runtime/meson.build to copy the runtime __init__.py into the meta-wheel build directory. - Enter the runtime subdirectory from nixl/meson.build. - Add the copied runtime initializer to the meta-wheel build inputs. pyproject.toml.in: - Add nixl.runtime to the setuptools package list so the routing initializer is included in the published meta wheel. V2 tests: - Replace the nixl_ep_v2 source import with the public nixl.runtime.ep import in every Python script under runtime/ep/tests, including the helpers that other tests import (test_deterministic_sort.py, test_adaptive_modes.py, test_adaptive_elastic.py). - Update utility imports and API references to use the installed package path. .gitlab/test_ep.sh: - Run the V2 correctness, elastic-membership, and topology-mode tests through the installed backend package. - Add the tests directory to PYTHONPATH so their sibling helper imports continue to work. Signed-off-by: Lior Shapira <lishapira@nvidia.com>
runtime/ep/nixl_ep_v2: - Format the imported runtime with NIXL's black and isort settings. - Mark intentional delayed imports and public re-exports for flake8. - Add type annotations and casts for optional topology values, bootstrap state, reference tensors, profiling results, and CUDA graph state so mypy can verify the code without changing runtime behavior. - Split a tuple of list appends and remove an unused exception binding to satisfy lint checks while preserving behavior. runtime/ep/tests: - Format the V2 tests and profilers with NIXL's normal Python style. - Annotate lazy nixl.runtime.ep imports, graph lifecycle state, timing records, closure references, and a reused variable name that static analysis cannot infer. - In test_adaptive_processes.py, turn the barrier lambda into a nested function and rename the ta..tj timers to t_a..t_j, which codespell would otherwise flag. - In test_nvfp4_dispatch.py, turn the zeros lambda into a nested function and type u8's pass-through result as Any; in test_zero_copy.py, annotate the graph state dicts; in sweep_sms.py, drop an unused os import. Signed-off-by: Lior Shapira <lishapira@nvidia.com>
Suppress the two DeepEP-derived host warnings on the V2 targets so the default werror build succeeds. Keep V2 tests out of the existing V1 EP CI script; V2 validation runs separately. Signed-off-by: Lior Shapira <lishapira@nvidia.com>
test_adaptive_elastic.py: - Document how the test checks the mode. The host derives it for every committed world: direct when exactly one NVLink domain is live, hybrid otherwise. After every commit, the test checks the mode the buffer committed (get_topology()). - Accept --plan as inline JSON or a JSON file path: a list of phases of dense rank ids. tests/elastic/plans/two_domains.json: - Add a 2 x 4 plan that scales between one and two domains, so the mode switches from direct to hybrid and back. test_adaptive_processes.py: - Document the relaunch of node 0 at phase 5 in the example plan. Signed-off-by: Lior Shapira <lishapira@nvidia.com>
test_adaptive_elastic.py: - By default, every live rank times one dispatch + combine pair per phase with CUDA events, as V1's elastic.py does. It prints the pair bandwidth and the avg/min/max time. Bytes are counted as in test_ep.py: scale-up is the tokens a rank receives and returns, and scale-out is the tokens that cross NVLink domains. - Add --kineto: Kineto per-kernel timings (dispatch, copy epilogue, combine, reduce epilogue) in every phase. The existing PERF line gains scale-out and scale-up bandwidth fields after its original ones. Its pair wall time reuses the default pair timing instead of timing the pair a second time. - Remove --perf-phases. It ran the same Kineto measurement on chosen phases only; --kineto replaces it. - Add --skip-perf-test to skip the default pair timing. On its own it runs only the correctness checks, as before. With --kineto it reports only the per-kernel timings. - Add --save-output [PATH]: the rank 0 process gathers every rank's records over the test's Store cohort and writes the run configuration, the per-rank phase records and the TRANS values as one JSON file. The default path is adaptive_elastic_results.json. stdout is unchanged. - The correctness checks and the remaining flags are unchanged. README.md: - Describe test_adaptive_elastic.py (elastic, correctness and performance) and list its command first. - Add test_ep.py and test_adaptive_modes.py to the list of other tests. Signed-off-by: Lior Shapira <lishapira@nvidia.com>
|
/build |
This branch has not been deployed
2 participants
NIXL EP V2: Adaptive Expert-Parallel Communication
Overview
NIXL EP V2 is an expert-parallel communication implementation for Mixture-of-Experts (MoE) models built on top of NIXL's device API.
Compared with NIXL EP V1 (examples/device/ep), V2 adds topology-aware direct and hybrid kernels, fixed expert placement across topology changes, and CUDA graph reuse while the active world grows or shrinks.
NIXL EP V2 is located under
runtime/epand is imported throughnixl.runtime.ep. The existing NIXL EP V1 implementation remains available throughimport nixl_ep.Features
Buffer Initialization
NIXL EP V2 uses an
ElasticBuffersized for the maximum expected topology:More complex scaling operations can establish or remove peer connections separately:
Key APIs
ElasticBuffer(...): Allocate communication resources for the maximum topology and initialize the current cohort.dispatch(...): Route token activations and expert selections to the destination ranks.combine(...): Return and reduce expert outputs.connect_ranks(...): Establish NIXL connections to new peers.disconnect_ranks(...): Remove connections to departed peers.stage_topology(...): Prepare a new active topology.commit_topology(): Make the staged topology visible to subsequent kernels.set_active_ranks(...): Stage and commit a new live world in one call. The direct or hybrid mode is derived automatically.set_num_active_sms(...): Change the active kernel SM count without changing the live world.Testing
The test suite under
runtime/ep/testsvalidates fixed and elastic expert-parallel communication.Plan files define a sequence of active worlds. Each phase validates dispatch and combine correctness, active membership, and the automatically selected communication mode.
Example plan:
This plan defines:
Getting Started
Build NIXL with NIXL EP enabled:
meson setup build \ -Ducx_path=<ucx-prefix> \ -Dprefix=<nixl-prefix> \ -Dbuildtype=release \ -Dbuild_nixl_ep=true meson compile -C build meson install -C buildNIXL EP V2 requires PyTorch, the CUDA toolkit, rdma-core development headers, and UCX with its GPU device API and CUDA IPC support.
After installation:
Refer to
runtime/ep/README.mdfor build requirements, environment variables, API details, and local test commands.