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 |

docs/DEPLOYMENT.md

@@ -581,6 +581,7 @@ 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). |
+| `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` until the upstream UCX `cuda_copy_md` length-truncation fix ships. |
 | `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. |
@@ -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:
 
@@ -684,6 +687,37 @@ operations that cross into later handles fail. See the reproducer and
 fix notes in this gist:
 <https://gist.github.com/nicolasnoble/e0e57eb5a1b902057ae3d1df59c039cf>.
 
+#### 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:
+
+```
+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.
 
@@ -431,6 +441,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/tests/test_pool_registration.py

@@ -14,6 +14,7 @@
 
 from modelexpress.accelerators import NIXL_ACCELERATOR_MEM_TYPE
 from modelexpress.nixl_transfer import (
+    _arena_single_mr_forced,
     NixlTransferManager,
     _pool_reg_enabled,
 )
@@ -222,6 +223,8 @@ def test_arena_registration_uses_vram_segment(self):
         used = tensor.numel() * tensor.element_size()
 
         class FakeArena:
+            live_allocation_count = 1
+
             def registered_range(self):
                 return base, used
 
@@ -237,6 +240,8 @@ def registered_range(self):
     def test_arena_registration_falls_back_when_tensor_uncovered(self):
         # A tensor outside [base, base+used) must not be served by the single MR.
         class FakeArena:
+            live_allocation_count = 1
+
             def registered_range(self):
                 return 0x1000, 0x2000
 
@@ -254,6 +259,8 @@ def test_arena_fallback_bypasses_pool_registration(self, monkeypatch):
         monkeypatch.setenv("MX_POOL_REG", "1")
 
         class FakeArena:
+            live_allocation_count = 1
+
             def registered_range(self):
                 return 0x1000, 0x2000
 
@@ -520,3 +527,119 @@ def test_raises_labeled_error(self):
         manager = self._manager(["ERROR"])
         with pytest.raises(RuntimeError, match="test transfer failed"):
             manager._wait_for_xfer(object(), None, "test transfer")
+
+
+class TestArenaSingleMrForced:
+    """MX_ARENA_SINGLE_MR keeps the single-MR path on a multi-allocation arena."""
+
+    def test_default_is_off(self, monkeypatch):
+        monkeypatch.delenv("MX_ARENA_SINGLE_MR", raising=False)
+        assert _arena_single_mr_forced() is False
+
+    def test_one_is_on(self, monkeypatch):
+        monkeypatch.setenv("MX_ARENA_SINGLE_MR", "1")
+        assert _arena_single_mr_forced() is True
+
+    def test_read_at_call_time(self, monkeypatch):
+        monkeypatch.setenv("MX_ARENA_SINGLE_MR", "1")
+        assert _arena_single_mr_forced() is True
+        monkeypatch.setenv("MX_ARENA_SINGLE_MR", "0")
+        assert _arena_single_mr_forced() is False
+
+
+class TestMultiAllocationArena:
+    """A single MR cannot describe an arena spanning several cuMemCreate handles.
+
+    UCX cuda_ipc resolves a region with cuMemGetAddressRange, which reports only
+    the allocation holding the base pointer, so the published rkey would cover
+    just the first chunk and the peer would read past what it mapped.
+    """
+
+    @staticmethod
+    def _make_manager() -> NixlTransferManager:
+        mgr = NixlTransferManager(agent_name="test", device_id=0)
+        mgr._agent = MagicMock()
+        mgr._agent.get_agent_metadata.return_value = b"metadata"
+        return mgr
+
+    @staticmethod
+    def _covering_arena(tensor, live_allocs):
+        base = tensor.data_ptr()
+        used = tensor.numel() * tensor.element_size()
+
+        class FakeArena:
+            live_allocation_count = live_allocs
+
+            def registered_range(self):
+                return base, used
+
+        return FakeArena(), base, used
+
+    def test_multi_allocation_arena_registers_per_tensor(self, monkeypatch):
+        monkeypatch.delenv("MX_ARENA_SINGLE_MR", raising=False)
+        tensor = torch.zeros(4, dtype=torch.float32)
+        arena, _, _ = self._covering_arena(tensor, live_allocs=1019)
+
+        mgr = self._make_manager()
+        assert mgr.register_arena(arena, {"w": tensor}) == b"metadata"
+
+        # Per-tensor, not one MR over the arena range.
+        args, kwargs = mgr._agent.register_memory.call_args
+        assert args[0][0] is tensor
+        assert kwargs == {"backends": ["UCX"]}
+
+    def test_single_allocation_arena_keeps_one_mr(self, monkeypatch):
+        monkeypatch.delenv("MX_ARENA_SINGLE_MR", raising=False)
+        tensor = torch.zeros(4, dtype=torch.float32)
+        arena, base, used = self._covering_arena(tensor, live_allocs=1)
+
+        mgr = self._make_manager()
+        assert mgr.register_arena(arena, {"w": tensor}) == b"metadata"
+
+        mgr._agent.register_memory.assert_called_once_with(
+            [(base, used, 0, "")],
+            mem_type=NIXL_ACCELERATOR_MEM_TYPE,
+            backends=["UCX"],
+        )
+
+    def test_env_override_keeps_one_mr_on_multi_allocation(self, monkeypatch):
+        # dmabuf/IB can span several handles in one registration, so deployments
+        # that validated the single-MR path there can keep it.
+        monkeypatch.setenv("MX_ARENA_SINGLE_MR", "1")
+        tensor = torch.zeros(4, dtype=torch.float32)
+        arena, base, used = self._covering_arena(tensor, live_allocs=1019)
+
+        mgr = self._make_manager()
+        assert mgr.register_arena(arena, {"w": tensor}) == b"metadata"
+
+        mgr._agent.register_memory.assert_called_once_with(
+            [(base, used, 0, "")],
+            mem_type=NIXL_ACCELERATOR_MEM_TYPE,
+            backends=["UCX"],
+        )
+
+    def test_multi_allocation_fallback_bypasses_pool_registration(self, monkeypatch):
+        # Pool reg resolves the same per-handle bounds that were insufficient.
+        monkeypatch.delenv("MX_ARENA_SINGLE_MR", raising=False)
+        monkeypatch.setenv("MX_POOL_REG", "1")
+        tensor = torch.zeros(4, dtype=torch.float32)
+        arena, _, _ = self._covering_arena(tensor, live_allocs=1019)
+
+        mgr = self._make_manager()
+        assert mgr.register_arena(arena, {"w": tensor}) == b"metadata"
+
+        args, kwargs = mgr._agent.register_memory.call_args
+        assert args[0][0] is tensor
+        assert kwargs == {"backends": ["UCX"]}
+
+    def test_warning_names_the_allocation_count(self, monkeypatch, caplog):
+        monkeypatch.delenv("MX_ARENA_SINGLE_MR", raising=False)
+        tensor = torch.zeros(4, dtype=torch.float32)
+        arena, _, _ = self._covering_arena(tensor, live_allocs=1019)
+
+        mgr = self._make_manager()
+        with caplog.at_level(logging.WARNING):
+            mgr.register_arena(arena, {"w": tensor})
+
+        assert "1019 physical allocations" in caplog.text
+        assert "MX_ARENA_SINGLE_MR=1" in caplog.text