refactor(reshard): route device work through the backend

docs/ARCHITECTURE.md

@@ -640,7 +640,7 @@ Loading precedence: CLI args > environment variables > config file > defaults.
 |--------|---------|
 | `__init__.py` | Package init, exports `register_modelexpress_loaders()` for callers to register the `modelexpress` and `mx` loaders with vLLM |
 | `client.py` | `MxClient` - gRPC client wrapping `PublishMetadata`, `ListSources`, `GetMetadata`, and `UpdateStatus` RPCs |
-| `accelerators/` | `AcceleratorBackend` boundary for accelerator-specific torch device control and fast-path capability gates, split into `base.py` (protocol), `cuda.py` (`CudaAcceleratorBackend`), and `xpu.py` (`XpuAcceleratorBackend`). CUDA and XPU are implemented backends; XPU keeps CUDA-only fast paths (pool registration, VMM arena, GDS) disabled and falls back to generic per-tensor NIXL registration. Further backends can be added behind the same interface |
+| `accelerators/` | `AcceleratorBackend` boundary for accelerator-specific torch device control and fast-path capability gates, split into `base.py` (protocol), `cuda.py` (`CudaAcceleratorBackend`), and `xpu.py` (`XpuAcceleratorBackend`). CUDA and XPU are implemented backends; XPU keeps CUDA-only fast paths (pool registration, VMM arena, GDS) disabled, falls back to generic per-tensor NIXL registration, and requires no classic-allocator pool for registered buffers. Further backends can be added behind the same interface |
 | `nixl_transfer.py` | `NixlTransferManager` - NIXL agent lifecycle, tensor registration, RDMA transfers |
 | `refit/` | Engine-agnostic live-refit primitives. `RefitTimingRecorder` provides normalized stage timing; `reshard/` provides loader-observed geometry capture, slice/transfer planning, rendezvous, and transport abstractions |
 | `gds_transfer.py` | GPUDirect Storage availability check and transfer utilities |
@@ -742,6 +742,35 @@ STALE. Long-lived framework integrations still need to call `close()` from
 their lifecycle; SIGKILL and mid-transfer failure recovery remain follow-up
 work.
 
+The receiver's device-specific work goes through `AcceleratorBackend` rather than
+`torch.cuda` directly: per-stage synchronization, the backend handed to
+`NixlTransferManager`, and the allocation scope for its receive and staging
+buffers. `registered_buffer_alloc_scope()` selects that scope from
+`requires_classic_alloc_pool()`, which states a requirement rather than a
+capability. CUDA requires it and XPU does not: CUDA uses the
+classic-`cudaMalloc` pool in `reshard/cuda_pool.py` because its caching allocator
+under `expandable_segments:True` can return VMM ranges that register successfully
+but fail during RDMA WRITE when `nvidia_peermem` cannot pin the underlying pages
+(that module's docstring records the exact failure signature). XPU uses its normal
+allocator because no equivalent hazard is known or has been observed on XPU.
+Successful XPU registration does not prove the WRITE-time hazard absent, so ruling
+it out would take a real RDMA write into an XPU buffer allocated under an
+expandable-segment equivalent. `torch.xpu` does expose `MemPool` and
+`XPUPluggableAllocator`, so an alternate XPU pool could be implemented if that
+test ever says one is needed.
+
+The receiver applies no publisher/target accelerator compatibility policy, and
+this is a known gap rather than a decision that cross-family refit is safe. The
+rendezvous identity deliberately carries no accelerator (the receiver builds it
+to *discover* the trainer, before it knows anything the trainer served) and the
+shard table carries no publisher family, so `accelerators_compatible` has nothing
+to compare on this path. No pairing is rejected on accelerator-family grounds;
+other NIXL, fabric, or model-geometry constraints may still prevent transfer.
+Closing the gap means
+publishing the source family in the shard table and comparing both endpoints
+through that gate; a target-only check would be unable to express the constraint,
+since compatibility is a property of the source-target pair.
+
 `engines/vllm/refit/receiver.py` supplies the vLLM-specific boundaries: capture
 on an unquantized meta twin, then installation through vLLM's layerwise reload
 and `process_weights_after_loading` path. Unsupported loader operations fail

modelexpress_client/python/modelexpress/accelerators/base.py

@@ -70,3 +70,15 @@ def supports_vmm(self) -> bool:
     def supports_gds(self) -> bool:
         """Return whether GPUDirect Storage loading is supported."""
         ...
+
+    def requires_classic_alloc_pool(self) -> bool:
+        """Return whether NIXL-registered buffers must come from a separate
+        classic-allocation pool on this backend.
+
+        A requirement, not a capability: True where this backend's default torch
+        allocator can hand out a range the HCA cannot pin, so registered buffers
+        have to be scoped into a classic-allocation pool instead (see
+        ``refit/reshard/cuda_pool.py`` for the hazard and the CUDA
+        implementation). False means the default allocator is already suitable.
+        """
+        ...

modelexpress_client/python/modelexpress/accelerators/cuda.py

@@ -60,3 +60,6 @@ def supports_vmm(self) -> bool:
 
     def supports_gds(self) -> bool:
         return True
+
+    def requires_classic_alloc_pool(self) -> bool:
+        return True

modelexpress_client/python/modelexpress/accelerators/xpu.py

@@ -78,3 +78,11 @@ def supports_vmm(self) -> bool:
 
     def supports_gds(self) -> bool:
         return False
+
+    def requires_classic_alloc_pool(self) -> bool:
+        # The classic pool works around a CUDA-specific interaction between
+        # expandable-segment allocations and nvidia_peermem pinning; no equivalent
+        # hazard is known or has been observed on XPU. Not a proof of absence -
+        # refit/reshard/cuda_pool.py has the failure signature and
+        # docs/ARCHITECTURE.md what would settle it.
+        return False

modelexpress_client/python/modelexpress/refit/README.md

@@ -110,7 +110,9 @@ Refit has two independent optimization surfaces:
 
 [`MdlLoader`](../engines/vllm/refit/installer.py) is a separate experimental vLLM installer called Mapped Direct Load (MDL). It caches direct, fused, and expert destination views so warm updates can copy into known slots instead of repeating general loader dispatch. MDL can consume partial input batches, but the reshard transport in this package does not yet expose a selector that reduces wire bytes for partial updates. The two features must not be treated as one end-to-end partial-refit path until that selector is wired and validated.
 
-Each receiver keeps one load-time receive buffer per captured destination, plus a source-dtype conversion staging buffer for any parameter whose served dtype differs from its load-time dtype. Both come from classic CUDA allocations and stay registered for the receiver's lifetime, because re-registering per refit is what the cached plan exists to avoid.
+Each receiver keeps one load-time receive buffer per captured destination, plus a source-dtype conversion staging buffer for any parameter whose served dtype differs from its load-time dtype. Both stay registered for the receiver's lifetime, because re-registering per refit is what the cached plan exists to avoid. These buffers use a backend-selected allocation scope: CUDA uses a classic `cudaMalloc`-backed pool because its caching allocator under `expandable_segments` can return VMM ranges that register successfully but fail during RDMA WRITE when `nvidia_peermem` cannot pin the underlying pages. XPU uses its normal allocator because no equivalent hazard is known or has been observed on XPU; successful XPU registration does not prove the WRITE-time hazard absent. `torch.xpu` does expose `MemPool` and `XPUPluggableAllocator`, so an alternate XPU pool could be implemented if one is ever needed. The selection is `AcceleratorBackend.requires_classic_alloc_pool()`.
+
+No publisher/target accelerator compatibility is checked on this path: neither the rendezvous identity nor the shard table carries the publisher's family, so `metadata/payload.py` has nothing to compare. No pairing is rejected on accelerator-family grounds; other NIXL, fabric, or model-geometry constraints may still prevent transfer. This is documented as a gap, not as a judgement that cross-family refit is validated; closing it means publishing the source family and comparing both endpoints.
 
 That buffer shape is why the vLLM receiver installs through `process_weights_after_loading` (PWAL) rather than MDL. The receiver reconstructs *load-time* tensors, and a quantized model still needs the engine's post-load processing to derive its runtime representation from them. MDL is appropriate only when the incoming tensors already match the validated runtime representation, which is why it is a separate opt-in path rather than the default.
 

modelexpress_client/python/modelexpress/refit/reshard/__init__.py

@@ -45,6 +45,7 @@
     ReadDescriptor,
     Transport,
 )
+from modelexpress.refit.reshard.alloc_scope import registered_buffer_alloc_scope
 from modelexpress.refit.reshard.cuda_pool import classic_cuda_alloc
 from modelexpress.refit.reshard.receiver import ReshardReceiver
 from modelexpress.refit.reshard.rendezvous import (
@@ -86,6 +87,7 @@
     "paired_runs",
     "plan_pull",
     "plan_transfer",
+    "registered_buffer_alloc_scope",
     "shard_region",
     "tensor_digest",
     "wrap_rendezvous_blob",

modelexpress_client/python/modelexpress/refit/reshard/alloc_scope.py

@@ -0,0 +1,50 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+"""Where the reshard receiver's receive and staging buffers are allocated from.
+
+Registering a buffer with NIXL constrains how it may be allocated, and the
+constraint is a property of the accelerator rather than of the refit. This module
+holds the family-agnostic selection so callers do not have to name a specific
+allocator; :mod:`modelexpress.refit.reshard.cuda_pool` holds the one
+accelerator-specific implementation that exists.
+"""
+
+from __future__ import annotations
+
+from contextlib import nullcontext
+from typing import TYPE_CHECKING, ContextManager
+
+if TYPE_CHECKING:
+    from modelexpress.accelerators import AcceleratorBackend
+
+
+def registered_buffer_alloc_scope(
+    backend: "AcceleratorBackend",
+) -> ContextManager[None]:
+    """Return the allocation scope ``backend`` needs for NIXL-registered buffers.
+
+    A backend that reports ``requires_classic_alloc_pool()`` gets
+    :func:`~modelexpress.refit.reshard.cuda_pool.classic_cuda_alloc`; every other
+    backend gets a no-op and allocates normally, which is the correct answer here
+    rather than a fallback. CUDA is the only dedicated pool implementation; fail
+    clearly if another backend requests one instead of silently using CUDA code.
+    """
+    if backend.requires_classic_alloc_pool():
+        if backend.name != "cuda":
+            raise NotImplementedError(
+                f"No registered-buffer allocation pool is implemented for "
+                f"backend {backend.name!r}"
+            )
+        # Imported per call rather than at module scope so this selection carries
+        # no import-time dependency on the one accelerator-specific
+        # implementation. Note this does not by itself keep cuda_pool out of a
+        # non-CUDA process: the package __init__ re-exports classic_cuda_alloc,
+        # so importing anything under refit.reshard loads that module anyway.
+        from modelexpress.refit.reshard.cuda_pool import classic_cuda_alloc
+
+        return classic_cuda_alloc()
+    return nullcontext()
+
+
+__all__ = ["registered_buffer_alloc_scope"]

modelexpress_client/python/modelexpress/refit/reshard/receiver.py

@@ -32,9 +32,10 @@
 import torch
 
 from modelexpress import envs
+from modelexpress.accelerators import accelerator_backend_for
 from modelexpress.client import MxClient
 from modelexpress.nixl_transfer import NixlTransferManager
-from modelexpress.refit.reshard.cuda_pool import classic_cuda_alloc
+from modelexpress.refit.reshard.alloc_scope import registered_buffer_alloc_scope
 from modelexpress.refit.reshard.rendezvous import gather_sources
 from modelexpress.refit.reshard.transfer_plan import (
     exact_descriptors,
@@ -327,13 +328,24 @@ def __init__(
         self._num_trainer_sources = num_trainer_sources
         self._timeout = timeout
         self._global_rank = global_rank
+        self._backend = accelerator_backend_for(device)
+
+        # TODO(publisher-accelerator): nothing here checks the publisher's
+        # accelerator family, because the rendezvous identity and the shard table
+        # do not carry one. Both same-family and cross-family pairings are
+        # unchecked by accelerator family. Publish the source family in the shard
+        # table and run it through
+        # metadata/payload.py::accelerators_compatible.
 
         # TODO(transport-agnostic): the receiver is engine-agnostic but still
         # transport-bound to NIXL (this manager, NixlReshardTransport, and the
         # fetch_remote_and_wait P2P handshake in _prepare). Abstract these behind
         # a transport interface so non-NIXL backends can plug in.
         self._manager = NixlTransferManager(
-            agent_name=agent_name, device_id=local_rank, listen_port=listen_port
+            agent_name=agent_name,
+            device_id=local_rank,
+            listen_port=listen_port,
+            accelerator_backend=self._backend,
         )
         self._manager.initialize()
         self._mx_client = MxClient(server_url=mx_server)
@@ -484,21 +496,21 @@ def _prepare(self, timeout: float) -> None:
 
         # dtype-mismatched sources (e.g. a bf16-served router for an fp32 dest):
         # one persistent bf16 STAGING buffer per convert param, registered as an
-        # RDMA target (classic cudaMalloc so the HCA can RDMA into it); each refit
-        # we RDMA into staging then cast staging -> the (load-time) receive buffer.
+        # RDMA target using the backend-selected allocation scope; each refit we
+        # RDMA into staging then cast staging -> the (load-time) receive buffer.
         # Allocation and registration happen in three places below (convert
         # staging, full-pull staging, receive buffers). They are accumulated into
         # one figure each rather than reported per buffer class, because the
-        # actionable question is how much of a cold start is cudaMalloc versus HCA
-        # registration.
+        # actionable question is how much of a cold start is allocation versus
+        # HCA registration.
         alloc_s = 0.0
         register_s = 0.0
 
         self._staging = {}
         self._staging_ptr = {}
         if plan.converts:
             _t = time.perf_counter()
-            with classic_cuda_alloc():
+            with registered_buffer_alloc_scope(self._backend):
                 self._staging = {
                     c.param_name: torch.empty(
                         c.dest_shape, dtype=c.src_dtype, device=self._device
@@ -520,7 +532,7 @@ def _prepare(self, timeout: float) -> None:
         self._full_staging_ptr = {}
         if plan.full_pulls:
             _t = time.perf_counter()
-            with classic_cuda_alloc():
+            with registered_buffer_alloc_scope(self._backend):
                 self._full_staging = {
                     full_pull.src_name: torch.empty(
                         full_pull.global_shape,
@@ -543,15 +555,16 @@ def _prepare(self, timeout: float) -> None:
             }
 
         # Receive buffers: one per captured param at its CAPTURED (load-time)
-        # shape/dtype, classic cudaMalloc, registered once. The live params are
-        # NOT RDMA targets; _install() writes the buffers into the live params.
+        # shape/dtype, allocated through the backend-selected scope and registered
+        # once. The live params are NOT RDMA targets; _install() writes the buffers
+        # into the live params.
         # Segment params (captured == served) are the RDMA targets - register them
         # + point _param_ptr at them. Convert params (router) are captured fp32 ->
         # their bf16 staging is the RDMA target and the refit casts into the buffer.
         seg_params = {seg.param_name for seg in plan.segments}
         self._recv_buffers = {}
         _t = time.perf_counter()
-        with classic_cuda_alloc():
+        with registered_buffer_alloc_scope(self._backend):
             for name in all_params:
                 shape, dtype = param_layout[name]
                 self._recv_buffers[name] = torch.empty(
@@ -806,7 +819,7 @@ def update_weights(self, step: int, *, timeout: float | None = None) -> dict:
                         receive_buffer.storage_offset() + copy.dest_offset,
                     )
                     destination.copy_(source_view)
-            torch.cuda.synchronize(self._device)
+            self._backend.synchronize(self._device.index)
             stages["reslice_s"] = time.perf_counter() - _t
 
         # Cast the served bf16 staging into the (fp32) receive buffer - a torch
@@ -817,12 +830,12 @@ def update_weights(self, step: int, *, timeout: float | None = None) -> dict:
                 self._recv_buffers[convert.param_name].copy_(
                     self._staging[convert.param_name]
                 )
-            torch.cuda.synchronize(self._device)
+            self._backend.synchronize(self._device.index)
             stages["convert_s"] = time.perf_counter() - _t
 
         _t = time.perf_counter()
         self._install(self._recv_buffers)
-        torch.cuda.synchronize(self._device)
+        self._backend.synchronize(self._device.index)
         stages["install_s"] = time.perf_counter() - _t
 
         metrics = {

modelexpress_client/python/tests/conftest.py

@@ -29,6 +29,7 @@ class MockAcceleratorBackend:
     pool_reg: bool = False
     vmm: bool = False
     gds: bool = False
+    classic_alloc_pool: bool = False
     set_device_calls: list[int] = field(default_factory=list)
     synchronize_calls: list[int | None] = field(default_factory=list)
     empty_cache_calls: int = 0
@@ -63,6 +64,9 @@ def supports_vmm(self) -> bool:
     def supports_gds(self) -> bool:
         return self.gds
 
+    def requires_classic_alloc_pool(self) -> bool:
+        return self.classic_alloc_pool
+
 
 @pytest.fixture
 def mock_accelerator_backend_cls():