[Weight Transfer] NIXL + MX Integration

[S1ro1]

Consolidates all the existing PRs for NIXl + MX into a clean implementation where I actually watched my agents

BF16 <-> BF16 + BF16 <-> FP8 Qwen3 30B on dummy task. Qwen3 235B A22B BF16 <-> BF16 transfer time comparison

---

Note

High Risk
Introduces a new RDMA-based weight transfer backend and changes orchestrator/trainer/inference coordination logic, which can affect distributed training stability and correctness. Also adds new runtime dependencies (Model Express, protobuf) and Docker-launched services that may impact deployment and CI.

Overview
Enables a new nixl_mx weight broadcast option that pushes weights trainer→inference via NIXL RDMA while using Model Express (MX) for rendezvous/metadata, wiring this through configs, orchestrator init, and inference worker extensions.

Implements the NIXL+MX transport stack: new vLLM worker (NIXLMxWeightUpdateWorker) and /init_nixl_mx endpoint, trainer-side NIXLMxWeightBroadcast with slot-based buffer allocation + conversion (bf16 passthrough / FP8 blockwise) and a TransportPlan that posts RDMA WRITEs based on MX-published tensor descriptors and expert maps.

Updates SLURM templates to optionally launch the MX + Redis stack via docker compose, sets UCX/LD paths for RDMA, and adjusts orchestrator scheduling/weight update flow to use MX status handshakes instead of filesystem markers when in nixl_mx mode. Adds unit/integration tests for MX rendezvous and conversion/slot planning, and adds modelexpress/protobuf dependencies plus a tilelang CUDA runtime preload shim.

^{Reviewed by Cursor Bugbot for commit dabaa19. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Hardcoded timeout ignores configurable timeout setting

Medium Severity

update_weights_from_path calls self.rendezvous.wait_for_all_peers_ready(timeout=1200) with a hardcoded 1200-second timeout. The trainer and orchestrator both use a configurable timeout from their respective configs. If a user configures a different timeout (e.g., for very large models needing longer), the inference worker won't respect it and could time out prematurely or wait too long.

 

^{Reviewed by Cursor Bugbot for commit 9d9b2cc. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

There are 3 total unresolved issues (including 1 from previous review).

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 75efd8c. Configure here.}

[cursor]

Scheduler filesystem probe broken for NIXL+MX mode

Medium Severity

_compute_next_ckpt_step always calls get_latest_ckpt_step on the filesystem broadcast directory, but NIXL+MX never writes checkpoint files there. This means latest_ckpt_step is always 0, forcing the scheduler into effectively strict-async behavior regardless of the strict_async_level setting. In non-strict mode with filesystem/NCCL, the orchestrator can jump ahead to the latest available checkpoint; this optimization is silently lost with NIXL+MX.

 

^{Reviewed by Cursor Bugbot for commit 75efd8c. Configure here.}

[cursor]

Non-primary HSDP ranks skip lazy_init but hit barrier deadlock risk

Medium Severity

In broadcast_weights, when dp_replicate > 1, the master rank (rank 0) calls self.rendezvous.wait_for_all_peers_ready which blocks synchronously waiting for the orchestrator, while all other ranks (including non-primary HSDP ranks) proceed directly to dist.barrier(). If the orchestrator is slow to signal, rank 0 blocks before the barrier while other ranks are already waiting at the barrier. This is fine structurally, but wait_for_all_peers_ready is a synchronous blocking call that holds the GIL and busy-polls with time.sleep(0.05), which could delay the barrier on large clusters.

 

^{Reviewed by Cursor Bugbot for commit 75efd8c. Configure here.}