fix(nixl): fall back to per-tensor registration on a multi-allocation arena

docs/ARCHITECTURE.md

@@ -675,7 +675,7 @@ Manages a NIXL agent and RDMA transfers for a single GPU worker:
 |--------|---------|
 | `__init__(agent_name, device_id, listen_port, accelerator_backend)` | Create NIXL agent with UCX backend; `listen_port` enables P2P listen thread; `accelerator_backend` owns torch device operations and accelerator capability gates |
 | `register_tensors(tensors)` | Register GPU tensors for RDMA, return serialized metadata. With `MX_POOL_REG=1` on a backend that supports pool registration, registers each unique cudaMalloc allocation backing the tensors instead of registering each tensor individually |
-| `register_arena(arena, tensors)` | Register the used VMM arena range once through dmabuf when the active accelerator backend supports the VMM arena fast path, then publish every tensor descriptor against that single MR |
+| `register_arena(arena, tensors)` | Register the used VMM arena range once through dmabuf when the active accelerator backend supports the VMM arena fast path, then publish every tensor descriptor against that single MR. Falls back to per-tensor registration when a tensor lies outside the arena range, or when the arena spans several `cuMemCreate` handles (a single MR cannot be addressed by cuda_ipc then; override with `MX_ARENA_SINGLE_MR=1`) |
 | `fetch_remote_and_wait(agent_name, ip, port)` | P2P: fetch remote NIXL metadata via listen thread (polls until loaded) |
 | `receive_from_source(source_metadata, source_tensors, ..., remote_agent_name)` | Execute RDMA read transfer; `remote_agent_name` skips `add_remote_agent` (P2P) |
 | `shutdown()` | Clean up NIXL agent and resources |
@@ -848,7 +848,7 @@ After the load strategy succeeds and the engine has finished post-processing, `L
 
 The arena path does not require `MX_POOL_REG=1`. Pool-reg remains the optimization for normal cudaMalloc deployments. Arena deployments use the direct `register_arena` seam because they already know the contiguous VA range and do not need `cuMemGetAddressRange` to rediscover allocation boundaries.
 
-Empirical validation on B200 + ConnectX on 2026-05-14 showed that a dmabuf MR over a multi-handle VMM range remains valid when holes exist inside the registered VA range. That is the property that lets `process_weights_after_loading` allocate replacement tensors, free discarded tensors, and still finish with one MR for the surviving used arena range.
+Empirical validation on B200 + ConnectX showed that a dmabuf MR over a multi-handle VMM range remains valid when holes exist inside the registered VA range. That is the property that lets `process_weights_after_loading` allocate replacement tensors, free discarded tensors, and still finish with one MR for the surviving used arena range. That property is specific to the dmabuf/IB path: `cuda_ipc` cannot address a single MR spanning several `cuMemCreate` handles, so multi-allocation arenas fall back to per-tensor registration there. See [Multi-handle arenas](DEPLOYMENT.md#multi-handle-arenas).
 
 ## NIXL Integration
 
@@ -953,7 +953,7 @@ See [`metadata.md`](metadata.md) for the full storage schema and debugging guide
 | `MX_METADATA_BACKEND` | (required on server; `""` on client) | Server: `redis` or `kubernetes`. Client: `""` / `server` / `redis` / `kubernetes` (central server) or `k8s-service` (decentralized via K8s Service routing) |
 | `MX_POOL_REG` | `0` | Discover cudaMalloc allocations via `cuMemGetAddressRange` and register each as a single NIXL block instead of registering tensors individually. Reduces NIXL registration count by 80-99% on typical vLLM models, cutting `ibv_reg_mr` time and metadata blob size; transfer semantics unchanged. Not required for `MX_VMM_ARENA=1`, which registers the arena directly |
 | `MX_VMM_ARENA` | `0` | Install a `CUDAPluggableAllocator` that routes weight-loading allocations into a CUDA VMM arena, then registers the used arena range once through dmabuf at end-of-load. Reserves 16.0 TiB of VA by default and commits physical memory only for mapped allocations. See [VMM Arena](#vmm-arena-cudapluggableallocator-hook) |
-| `UCX_CUDA_COPY_REG_WHOLE_ALLOC` | (UCX default) | Set to `off` with `MX_VMM_ARENA=1` until the upstream UCX `cuda_copy_md` length-truncation fix ships. |
+| `UCX_CUDA_COPY_REG_WHOLE_ALLOC` | (UCX default) | Set to `off` with `MX_VMM_ARENA=1` on any UCX predating the `cuda_copy_md` length-truncation fix (openucx/ucx#11461). Scoped to the `cuda_copy` transport; it does not affect `cuda_ipc`. |
 | `MX_P2P_METADATA` | `1` | Enable P2P metadata exchange on source workers. Set to `0` to publish full metadata through a central-coordinator backend; ignored on decentralized backends that require P2P metadata |
 | `MX_METADATA_PORT` | `5555` | Base NIXL listen port; effective port is `MX_METADATA_PORT + device_id` |
 | `MX_WORKER_GRPC_PORT` | `6555` | Base worker gRPC port for P2P tensor and artifact manifest serving; effective port is `MX_WORKER_GRPC_PORT + device_id` |

docs/DEPLOYMENT.md

@@ -581,7 +581,8 @@ See [`K8S_SERVICE_BACKEND.md`](K8S_SERVICE_BACKEND.md) for the design rationale,
 | `MX_METRICS_SCHEME` | `""` | Optional run/scheme label added to every metric, so multiple runs compare on one dashboard. |
 | `MX_POOL_REG` | `0` | Allocation-level NIXL registration via `cuMemGetAddressRange`. Registers each unique cudaMalloc block instead of each tensor, typically 80-99% fewer registrations, without changing transfer semantics. `MX_VMM_ARENA=1` uses direct arena registration and does not require pool-reg. |
 | `MX_VMM_ARENA` | `0` | Route weight allocations into a CUDA VMM arena via PyTorch's `CUDAPluggableAllocator`, then register the used arena range as one NIXL MR with dmabuf at end-of-load. Reserves 16.0 TiB of VA by default, with no physical commit until allocations are mapped. Requires the `modelexpress.vmm._alloc_ext` C extension to have built at install time; if it did not, this flag is a no-op with a warning and the loader falls back to the pool-reg path. See [VMM Arena](#vmm-arena-single-mr-registration). |
-| `UCX_CUDA_COPY_REG_WHOLE_ALLOC` | (UCX default) | Set to `off` with `MX_VMM_ARENA=1` until the upstream UCX `cuda_copy_md` length-truncation fix ships. |
+| `MX_ARENA_SINGLE_MR` | `0` | Keep single-MR arena registration even when the arena spans several `cuMemCreate` handles. Only safe on transports that can register across handles (dmabuf/IB); cuda_ipc cannot, so the default falls back to per-tensor registration. See [VMM Arena](#vmm-arena-single-mr-registration). |
+| `UCX_CUDA_COPY_REG_WHOLE_ALLOC` | (UCX default) | Set to `off` with `MX_VMM_ARENA=1` on any UCX predating the `cuda_copy_md` length-truncation fix (openucx/ucx#11461). Scoped to the `cuda_copy` transport; it does not affect `cuda_ipc`. |
 | `MX_NIXL_BACKEND` | `UCX` | NIXL backend for GPU-to-GPU RDMA. `UCX` (default) for InfiniBand / RoCE. `LIBFABRIC` for AWS EFA — see [NIXL Backend Selection](#nixl-backend-selection). |
 | `MX_RDMA_NIC_PIN` | (unset) | Per-rank IB NIC pinning. `auto` runs a topology probe; comma-separated NIC list is an explicit override. Workaround for openucx/ucx#11259. |
 | `MX_RDMA_NIC_PIN_MIN_RATE_GBPS` | (auto, max-rate filter) | Override the auto-detect rate filter with an explicit lower bound (Gb/s). |
@@ -662,7 +663,9 @@ arena address. Frees unmap and release that handle, so replacement tensors
 created during post-processing can return physical memory before the final
 registration step. At end-of-load, ModelExpress registers the used arena
 range once through dmabuf and publishes all tensor descriptors against
-that single MR.
+that single MR — but only when the arena is backed by a single physical
+allocation. An arena that spans several `cuMemCreate` handles falls back to
+per-tensor registration; see [Multi-handle arenas](#multi-handle-arenas) below.
 
 Recommended source-worker setting:
 
@@ -677,13 +680,48 @@ arena registration bypasses the pool-reg path and calls `register_arena`
 directly. The arena produces one MR for the used range regardless of the
 pool-reg setting.
 
-Set `UCX_CUDA_COPY_REG_WHOLE_ALLOC=off` until the upstream UCX
-`cuda_copy_md` length-truncation fix ships. Without it, UCX can truncate
-a multi-handle VMM registration to the first physical handle, and RDMA
-operations that cross into later handles fail. See the reproducer and
-fix notes in this gist:
+Set `UCX_CUDA_COPY_REG_WHOLE_ALLOC=off` on any UCX predating the
+`cuda_copy_md` length-truncation fix (openucx/ucx#11461). Without it, UCX
+can truncate a multi-handle VMM registration to the first physical
+handle, and RDMA operations that cross into later handles fail. See the
+reproducer and fix notes in this gist:
 <https://gist.github.com/nicolasnoble/e0e57eb5a1b902057ae3d1df59c039cf>.
 
+That knob covers the `cuda_copy` transport only. It has no effect on
+`cuda_ipc`, which has its own multi-handle limitation described in
+[Multi-handle arenas](#multi-handle-arenas) below.
+
+#### Multi-handle arenas
+
+A CUDA fabric/IPC handle names exactly one `cuMemCreate` allocation. UCX
+cuda_ipc resolves a registered region with `cuMemRetainAllocationHandle` and
+`cuMemGetAddressRange`, both of which report the allocation holding the base
+pointer rather than the whole reserve. Registering a multi-allocation arena as
+one MR therefore publishes an rkey covering only its first chunk, and the peer
+reads past what it mapped — measured on GB200 MNNVL as a segfault in
+`cuMemcpyDtoDAsync_v2` with an arena spanning 1019 chunks.
+
+ModelExpress detects this (`live_allocation_count > 1`) and falls back to
+per-tensor registration, which is correct because each arena allocation is one
+handle, so every tensor lies wholly inside one. The log line names the count:
+
+```text
+register_arena: arena spans 1019 physical allocations; a single MR would publish
+an rkey covering only the first, which cuda_ipc cannot address. Falling back to
+per-tensor registration ...
+```
+
+`UCX_CUDA_COPY_REG_WHOLE_ALLOC=off` does not cover this case: it applies to
+cuda_copy, while the truncation above is in cuda_ipc, and UCX 1.21 has no
+cuda_ipc equivalent. Upstream fixes are in flight for both sides —
+[openucx/ucx#11283](https://github.com/openucx/ucx/pull/11283) for cuda_ipc and
+[openucx/ucx#11461](https://github.com/openucx/ucx/pull/11461) for the
+cuda_copy/dmabuf length truncation.
+
+Set `MX_ARENA_SINGLE_MR=1` to keep the single-MR path on deployments where it
+was validated, i.e. dmabuf/IB, where `ibv_reg_dmabuf_mr` does span several
+handles.
+
 ### P2P Metadata Exchange
 
 P2P metadata exchange is enabled by default. Source workers expose their own per-worker gRPC `WorkerService` (the `WorkerGrpcServer` on `MX_WORKER_GRPC_PORT`) and their NIXL agent metadata directly on the worker's NIXL listen thread (`MX_METADATA_PORT`). Targets fetch tensor manifests or artifact manifests directly from the source worker rather than pulling them through the central store. For file-backed cache artifacts, targets also call `PrepareArtifactChunk` and `ReleaseArtifactChunk` on this worker service while bytes move through NIXL into target-local staging, then install the staged artifact into the runtime cache directory. The division of responsibility depends on which metadata backend is in use:

modelexpress_client/python/modelexpress/envs.py

@@ -94,6 +94,7 @@
     MX_TRANSFER_LOG_DIR: str
     # VMM arena
     MX_VMM_ARENA: bool
+    MX_ARENA_SINGLE_MR: bool
     # Framework artifact (JIT cache) transfer
     MX_ARTIFACT_TRANSFER: bool
     MX_ARTIFACT_BUNDLE_ROOT: Optional[str]
@@ -301,6 +302,7 @@ def _env_positive_float(name: str, default: float) -> float:
     "MX_TRANSFER_LOG_DIR": lambda: os.environ.get("MX_TRANSFER_LOG_DIR", "/tmp/mx_logs"),
     # ── VMM arena ──────────────────────────────────────────────────────────
     "MX_VMM_ARENA": lambda: os.environ.get("MX_VMM_ARENA") == "1",
+    "MX_ARENA_SINGLE_MR": lambda: os.environ.get("MX_ARENA_SINGLE_MR") == "1",
     # ── Framework artifact (JIT cache) transfer ────────────────────────────
     "MX_ARTIFACT_TRANSFER": lambda: os.environ.get("MX_ARTIFACT_TRANSFER", "").strip().lower()
     in _TRUTHY,

modelexpress_client/python/modelexpress/nixl_transfer.py

@@ -71,6 +71,16 @@ def _resolve_nixl_backend() -> str:
     return raw
 
 
+def _arena_single_mr_forced() -> bool:
+    """Whether to keep the single-MR arena path on a multi-allocation arena.
+
+    MX_ARENA_SINGLE_MR=1 forces it. Only safe on transports that can span
+    several cuMemCreate handles in one registration (dmabuf/IB); cuda_ipc
+    cannot. Read at call time so tests can toggle the env var.
+    """
+    return envs.MX_ARENA_SINGLE_MR
+
+
 def _pool_reg_enabled() -> bool:
     """Whether allocation-level pool registration is enabled.
 
@@ -393,12 +403,19 @@ def register_arena(
         consumes a dmabuf via `ibv_reg_dmabuf_mr` and produces ONE
         lkey/rkey covering all live tensors.
 
-        Empirically validated on Blackwell + ConnectX over InfiniBand
-        against a CUDA VMM range with multiple cuMemCreate handles and
-        mid-range holes (chunks unmapped + released after the export):
-        registration succeeds, the dmabuf attach pins the currently-
-        mapped physical pages, and the HCA translation table survives
-        subsequent CUDA-side unmaps.
+        The multi-handle case is validated on the dmabuf/IB path only.
+        On Blackwell + ConnectX over InfiniBand, against a CUDA VMM range
+        with multiple cuMemCreate handles and mid-range holes (chunks
+        unmapped + released after the export): registration succeeds, the
+        dmabuf attach pins the currently-mapped physical pages, and the
+        HCA translation table survives subsequent CUDA-side unmaps.
+
+        It does NOT hold on UCX cuda_ipc, where a fabric handle names one
+        cuMemCreate allocation and a single MR would publish an rkey
+        covering only the first chunk. That is why this method falls back
+        to per-tensor registration when the arena spans several
+        allocations, unless MX_ARENA_SINGLE_MR overrides it. Upstream fix:
+        openucx/ucx#11283.
 
         Per-tensor descriptors are still built (tensor name -> addr,
         size, dtype) because the receiver matches by name and computes
@@ -431,6 +448,35 @@ def register_arena(
             )
             return self.register_tensors(tensors)
 
+        # A CUDA fabric/IPC handle names exactly one cuMemCreate allocation:
+        # UCX cuda_ipc resolves a region with cuMemRetainAllocationHandle and
+        # cuMemGetAddressRange, which report the FIRST allocation under the
+        # range rather than the whole reserve. Registering a multi-allocation
+        # arena as one MR therefore publishes an rkey covering only its first
+        # chunk, and the peer's cuMemcpyDtoDAsync_v2 reads past what it mapped.
+        # Measured on GB200 MNNVL: Kimi-K3 arena, 1019 chunks, segfault in
+        # uct_cuda_ipc_ep_get_zcopy. Per-tensor registration is correct because
+        # the arena does one cuMemCreate per allocation, so every tensor lies
+        # wholly inside one handle.
+        #
+        # dmabuf/IB registration does span several handles, so deployments that
+        # validated the single-MR path there can keep it with
+        # MX_ARENA_SINGLE_MR=1.
+        live_allocs = arena.live_allocation_count
+        if live_allocs > 1 and not _arena_single_mr_forced():
+            logger.warning(
+                "register_arena: arena spans %d physical allocations; a single "
+                "MR would publish an rkey covering only the first, which "
+                "cuda_ipc cannot address. Falling back to per-tensor "
+                "registration for %d tensors over [0x%x, 0x%x). Set "
+                "MX_ARENA_SINGLE_MR=1 to force single-MR (dmabuf/IB only).",
+                live_allocs,
+                len(tensor_descriptors),
+                base,
+                base + used,
+            )
+            return self.register_tensors(tensors, force_per_tensor=True)
+
         # NIXL resolves descriptors by containment, so one tensor outside
         # [base, base+used) fails prep_xfer_dlist for the whole transfer.
         uncovered = [

modelexpress_client/python/modelexpress/vmm/README.md

@@ -23,18 +23,37 @@ MX_VMM_ARENA=1
 That is the only knob. The arena enables itself on supported devices and
 falls back with a warning when the underlying C extension is missing.
 
-### Required deployment flag
+### Deployment flag for the cuda_copy path
 
 UCX's `cuda_copy_md` path probes allocations with `cuMemGetAddressRange`,
 which returns per-handle bounds rather than the full reserve when called
 against a multi-handle VMM range. Without the override, UCX truncates
-transfers to a single physical chunk. Until the upstream fix lands
-(openucx/ucx#11461), set:
+transfers to a single physical chunk. On any UCX predating the upstream
+fix (openucx/ucx#11461), set:
 
 ```bash
 UCX_CUDA_COPY_REG_WHOLE_ALLOC=off
 ```
 
+This knob is scoped to the `cuda_copy` transport and does not affect
+`cuda_ipc`. Multi-allocation arenas have a separate limitation on
+`cuda_ipc`; see Transport support below before enabling the arena on an
+NVLink or MNNVL deployment.
+
+### Transport support
+
+Single-MR arena registration is validated on the dmabuf/IB path, where
+`ibv_reg_dmabuf_mr` genuinely spans several `cuMemCreate` handles.
+
+It does not hold on `cuda_ipc`. A CUDA fabric/IPC handle names exactly
+one `cuMemCreate` allocation, so a single MR over a multi-allocation
+arena publishes an rkey covering only the first chunk and the peer reads
+past what it mapped. ModelExpress detects this and falls back to
+per-tensor registration. `MX_ARENA_SINGLE_MR=1` overrides that fallback
+and is only safe on dmabuf/IB. The upstream fix is openucx/ucx#11283.
+The mechanism and the measured failure are in the Multi-handle arenas
+section of `docs/DEPLOYMENT.md`.
+
 ## Why this exists
 
 Without an arena, NIXL registers each tensor as its own MR. A
@@ -148,12 +167,18 @@ Key invariants:
 
 Validated configurations:
 
-- Blackwell B200 + ConnectX over InfiniBand, single-pod and P2P.
+- Blackwell B200 + ConnectX over InfiniBand, single-pod and P2P. This is
+  the only transport on which single-MR registration is validated.
 - vLLM 0.x with `CUDAPluggableAllocator` (PyTorch 2.4+).
 - TP=8 ranks running concurrently against independent arenas.
 
 Known limitations:
 
+- Single-MR registration does not work on the UCX `cuda_ipc` transport
+  when the arena spans more than one `cuMemCreate` handle. The published
+  rkey covers only the first chunk. ModelExpress falls back to per-tensor
+  registration in that case; see Transport support above. Upstream fix:
+  openucx/ucx#11283.
 - VMM arena is x86_64-Linux only (CUDA driver constraint).
 - Bump pointer is monotonic. Long-lived processes that load and unload
   many models within one arena lifetime will fragment the VA but never

modelexpress_client/python/modelexpress/vmm/runtime.py

@@ -244,10 +244,12 @@ def log_arena_post_load(ctx: "LoadContext") -> None:
     ``ibv_reg_dmabuf_mr`` over ``[base, base+used_bytes)`` already ran
     inside ``LoadStrategyChain`` via
     ``NixlTransferManager.register_arena``; this hook is purely
-    diagnostic. Empirically validated on Blackwell + ConnectX over
-    InfiniBand: the registration succeeds over a VA range with
+    diagnostic. Validated on the dmabuf/IB path, Blackwell + ConnectX
+    over InfiniBand: the registration succeeds over a VA range with
     mid-range holes from prior ``cuMemUnmap`` calls, and the dmabuf pin
-    keeps live tensor pages addressable to the HCA.
+    keeps live tensor pages addressable to the HCA. On ``cuda_ipc`` a
+    multi-allocation arena cannot be covered by one MR at all, and
+    ``register_arena`` falls back to per-tensor registration.
     """
     arena = _vmm_arenas.get(ctx.device_id)
     if arena is None: