vllm.distributed.weight_transfer.sharded_rdt_engine ¶

class ShardedRDTWeightTransferEngine(
    WeightTransferEngine[
        ShardedRDTWeightTransferInitInfo,
        ShardedRDTWeightTransferUpdateInfo,
    ]
):
    """Pull-based RDT/NIXL backend that transports only the slice each worker
    consumes.

    Requires ``distributed_executor_backend="ray"``, ``nixl`` in the shared env,
    a named trainer actor exposing a ``@ray.method(tensor_transport="nixl")``
    producer, and weight loaders that stay inside ``SUPPORTED_OPS`` — anything
    needing real data (``.to``, ``.item``, arithmetic, bool-mask indexing)
    raises during the bake.

    The plan is baked once at ``init_transfer_engine`` into one scatter list
    per fully-loaded leaf module, indexed by source name; every
    ``update_weights`` replays the modules its gathered names cover.
    """

    init_info_cls = ShardedRDTWeightTransferInitInfo
    update_info_cls = ShardedRDTWeightTransferUpdateInfo
    # receive_weights pulls synchronously but defers GPU post-processing to a
    # background thread so it overlaps the next chunk's pull, so ``update_weights``
    # skips the base's device sync and ``finish_weight_update`` drains the
    # deferred work before finalize.
    defers_processing = True
    # The baked replay plan is a function of one concrete model's parameter
    # layout, so a separate draft model cannot reuse it.
    supports_draft_weight_update = False

    def __init__(
        self,
        config: WeightTransferConfig,
        vllm_config: "VllmConfig",
        device: torch.device,
        model: torch.nn.Module,
    ) -> None:
        super().__init__(config, vllm_config, device, model)
        # This engine's data plane is Ray's: the consumer resolves each trainer's
        # serve actor with ``ray.get_actor`` and pulls over NIXL through the
        # workers' tensor transport, neither of which exists unless the workers
        # ARE Ray actors. Caught here, at construction, because the alternative
        # is an opaque failure during the first handshake. A single-worker
        # deployment defaults to "uni", so "ray" has to be asked for explicitly.
        backend = vllm_config.parallel_config.distributed_executor_backend
        # A custom Executor subclass is a deliberate override we cannot judge;
        # only reject the named backends that are definitely not Ray.
        if isinstance(backend, str) and backend != "ray":
            raise ValueError(
                f"The 'sharded_rdt' weight transfer backend requires "
                f"distributed_executor_backend='ray', got {backend!r}. Pass "
                f"--distributed-executor-backend ray (and keep "
                f"VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1, the default, so the "
                f"workers are RayExecutorV2 actors with tensor transport)."
            )
        # Ownership table, routing rules, and (once bound) every producer's
        # handle — pulls go to one owner per chunk, but free_group fans out to
        # every owner of the group, so all handles are bound.
        self._router: RdtRouter | None = None
        # Driver-supplied total consumer count (init_info.num_consumers); 0 => infer
        # from parallel_config. See _num_consumers.
        self._num_consumers_override: int = 0
        # Source name -> the plan for the module consuming it. Several names of a
        # fused module share one entry; replay dedups. A live name absent here
        # fails the plan build.
        self._name_to_plan: dict[str, list[_Scatter]] = {}
        # name -> (dtype_name, shape) for every init name; the bake builds its
        # lazies from this.
        self._name_meta: dict[str, tuple[str, list[int]]] = {}
        # Names whose copy_ fired during the bake. Unbaked names not here never
        # move data (e.g. non-local EP experts), so receive_weights skips them.
        self._live_names: set[str] = set()

        # ---- Consumer-side pre-registered receive buffers ----------------------
        # One persistent 1-D buffer per dtype per ring slot, STORAGE-registered with
        # NIXL once, so every per-slice view hits Ray's registration cache and the
        # recv path does zero registration in steady state. See the doc's "Sizing
        # buffers once matters beyond throughput" for why they rarely grow.
        #
        # Ring depth is both the slot count and the in-flight pull count, one
        # quantity from ``num_rdt_buffers``. Decoupling them with a spare slot
        # measurably slows the sync.
        self._ring_depth = 1
        self._dest_buffers: list[dict[torch.dtype, torch.Tensor]] = [{}]
        self._slot_read_done: list[Any] = []  # one torch.cuda.Event per slot
        # Generation counters guarding the events (see _ensure_proc_worker).
        self._slot_queued: list[int] = []
        self._slot_done: list[int] = []
        self._slot_cv: Any = None
        self._pull_slot = 0
        # (id(chunk), slot) -> (buffer data_ptr, per-key views).
        # Lives on the engine, not on the _Chunk: the plan is static and shared
        # across syncs, so a per-slot runtime cache has no business inside it.
        self._targets_cache: dict[tuple[int, int], tuple[int, list[torch.Tensor]]] = {}
        self._buffer_presize = 0  # bytes; set from init_info
        self._pending_frees: list[Any] = []  # free_group signal refs, drained per sync
        # The STATIC plan: built once — at init from init_info.group_lens, else
        # lazily on the first update_weights — and reused. Non-None means
        # update_weights ignores its per-sync names.
        self._cached_plan: _CallPlan | None = None
        # Completed syncs. The chunk pipeline runs SERIAL during sync 0, when both
        # sides still register buffers, and pipelines from sync 1. See the doc.
        self._completed_syncs = 0

        # ---- Background post-processing thread (pull/process pipelining) -------
        # The RPC thread hands pulled buffer views to this worker, which runs
        # materialize/scatter/quant/kernel_copy on its own stream while the next
        # pull proceeds. drain_pending() joins it before finalize.
        self._proc_queue: Any | None = None
        self._proc_thread: Any | None = None
        self._proc_stream: Any | None = None
        # Dedicated quant thread (pass 2): own queue + stream so quant never
        # sits between two items' scatters on the scatter thread (measured to
        # stall the RPC thread's slot handshake ~0.5-0.75s/iter).
        self._quant_queue: Any | None = None
        self._quant_thread: Any | None = None
        self._quant_stream: Any | None = None
        self._proc_error: BaseException | None = None

    def init_transfer_engine(self, init_info: ShardedRDTWeightTransferInitInfo) -> None:
        """Configure the ring, bind the producers, bake the replay plan, and
        pre-register every NIXL buffer -- in that order, because each step depends
        on the previous one.

        The bake drives ``model.load_weights`` and the pre-registration blocks on
        RPCs to the producers, so this is a heavyweight one-off; every later
        ``update_weights`` is pure replay.
        """
        # Read the attribute directly: if a vLLM bump renames it, this guard
        # must fail loudly at init rather than silently disappear.
        if self.parallel_config.enable_eplb:
            raise RuntimeError(
                "sharded_rdt does not support EPLB (enable_eplb=true): dynamic "
                "expert rearrangement invalidates the baked replay plan, which "
                "records each expert's destination slot once at init — after a "
                "rearrangement the replay would silently load weights into the "
                "wrong expert slots. Disable EPLB or use another weight-transfer "
                "backend."
            )
        # The worker is often a different install from whatever spawned it, so
        # check here too and not only where the actor options are set.
        check_ray_rdt_version()
        self._configure_ring(init_info)
        self._resolve_producers(init_info)

        # A pure dry run: the trainer's gather cache is empty at init, so nothing
        # can (or does) get pulled -- we only record how each slice is fetched and
        # where it lands, then restore the model.
        self._bake(init_info)
        self._build_static_plan(init_info)
        # Register ALL NIXL memory now, while the fabric is idle: dma-buf GPUDirect
        # registration concurrent with in-flight RDMA intermittently fails
        # (ibv_reg_mr 'Bad address'), which bites under M:N fan-in. Sizes come from
        # the static plan, so this is exact.
        self._preregister_at_init()
        # Start the background post-processing worker (pull/process pipelining).
        self._ensure_proc_worker()

    def _configure_ring(self, init_info: ShardedRDTWeightTransferInitInfo) -> None:
        """Ring depth K.

        Must run before ``_ensure_proc_worker`` creates the per-slot events and
        counters, and before any buffer is grown (both happen on the first pull).
        """
        self._num_consumers_override = int(init_info.num_consumers or 0)
        k = max(1, int(init_info.num_rdt_buffers))
        self._ring_depth = k
        self._dest_buffers = [{} for _ in range(k)]
        self._targets_cache = {}
        self._buffer_presize = int(float(init_info.buffer_presize_gb) * (1 << 30))
        logger.info(
            "[RDT-RING] active_pulls=%d slots=%d presize=%.2fGiB",
            k,
            k,
            self._buffer_presize / (1 << 30),
        )

    def _resolve_producers(self, init_info: ShardedRDTWeightTransferInitInfo) -> None:
        """Work out this worker's consumer identity, build the router, and bind
        EVERY producer actor.

        Pull routing is M:N: each chunk goes to ONE producer holding every name
        in it (see ``RdtRouter``), so one producer serves a whole pull. All
        producers are bound regardless, because the per-group ``free_group``
        signal fans out to every owner of a group, including producers this
        worker never pulls from.
        """
        try:
            import ray
        except ImportError as e:
            raise RuntimeError(
                "Ray is required for the 'sharded_rdt' weight transfer "
                "backend. Install Ray and run workers as Ray actors "
                "(distributed_executor_backend='ray')."
            ) from e

        producer_names = list(init_info.trainer_actor_names)
        if not producer_names:
            raise RuntimeError(
                "Sharded RDT engine requires a trainer producer: set "
                "init_info.trainer_actor_names."
            )

        consumer_id = self._resolve_consumer_id(init_info)
        name_owner_class = list(init_info.name_owner_class or [])
        if name_owner_class and len(name_owner_class) != len(init_info.names):
            raise RuntimeError(
                f"Sharded RDT engine: {len(name_owner_class)} owner-class entries "
                f"for {len(init_info.names)} names."
            )
        router = RdtRouter(
            len(producer_names),
            self._num_consumers(),
            list(init_info.owner_sets) or None,
            name_owner_class or None,
            list(init_info.names),
            list(init_info.group_lens),
            # Carve the block per deployment so every replica's worker w pulls
            # each name from the same producer, which is what lets that producer
            # serve all of them from one slot.
            workers_per_replica=self._workers_per_replica(init_info),
        )
        router.validate()

        # A fresh router per init, never appended to: every owner index is a
        # position in these lists, so a rejoining engine's second
        # init_transfer_engine must not shift them.
        actors: list[Any] = []
        methods: list[Any] = []

        for chosen_name in producer_names:
            try:
                actor = ray.get_actor(
                    chosen_name,
                    namespace=init_info.trainer_actor_namespace,
                )
            except ValueError as e:
                raise RuntimeError(
                    f"Sharded RDT engine could not find trainer actor "
                    f"{chosen_name!r} (namespace="
                    f"{init_info.trainer_actor_namespace!r})."
                ) from e
            # ``ray.get_actor`` resolves by NAME, in a process that never
            # imported the producer's class, and ``enable_tensor_transport``
            # lives in that class's metadata: the creating process passes
            # ``meta.enable_tensor_transport``, and Ray infers it from the
            # class's ``@ray.method(tensor_transport=...)`` decorators. A
            # name-resolved handle therefore always reports False and the
            # dispatch guard rejects the pull, even though the trainer set the
            # option. Inherent to resolving by name across processes, not a
            # version bug; it fails on Ray 2.56.0.
            #
            # Forcing it skips no validation. Ray's next guard,
            # ``actor_has_tensor_transport``, asks the LIVE actor whether it can
            # build a NIXL agent and still runs, and ``_spawn_server`` hardcodes
            # the actor option, so a misconfigured producer cannot reach here.
            actor._ray_enable_tensor_transport = True
            actors.append(actor)
            methods.append(getattr(actor, init_info.produce_method_name))
        router.bind(actors, methods, consumer_id)
        self._router = router
        logger.info(
            "Sharded RDT engine (consumer %d) bound to all %d producers "
            "(batched method %r, %d owner class(es))",
            consumer_id,
            len(producer_names),
            init_info.produce_method_name,
            len(init_info.owner_sets) or 1,
        )

    def _workers_per_replica(self, init_info: ShardedRDTWeightTransferInitInfo) -> int:
        """Consumers per inference deployment, assuming a uniform fleet.

        Read by this worker's consumer id and by the router's block carve, which
        must not disagree, so it is derived once here.
        """
        return max(1, self._num_consumers() // max(1, int(init_info.num_replicas or 1)))

    def _resolve_consumer_id(self, init_info: ShardedRDTWeightTransferInitInfo) -> int:
        """This worker's DISTINCT index in 0..C-1 across the whole fleet.

        Within one engine that is ``_global_worker_index()``. But a fleet of
        INDEPENDENT engines (each with its own parallel config) restarts that
        index at 0 per engine, so each engine offsets into its own range using
        ``replica_rank``: with a uniform fleet,
        ``workers_per_replica = C // num_replicas``. ``num_replicas`` defaults to
        1 (offset 0), preserving single-engine and single-DP-deployment behaviour.
        """
        replica_rank = max(0, int(init_info.replica_rank or 0))
        return (
            replica_rank * self._workers_per_replica(init_info)
            + self._global_worker_index()
        )

    def _build_static_plan(self, init_info: ShardedRDTWeightTransferInitInfo) -> None:
        """Build the chunk/free plan once. It never changes across syncs, so
        ``update_weights`` needs no per-sync names."""
        if not init_info.group_lens:
            raise ValueError(
                "Sharded RDT engine requires init_info.group_lens (the gather-"
                "group partition of init_info.names)."
            )
        if sum(init_info.group_lens) != len(init_info.names):
            raise ValueError(
                f"init_info.group_lens sums to {sum(init_info.group_lens)} "
                f"but {len(init_info.names)} names were given."
            )
        self._cached_plan = self._build_call_plan(init_info.names, init_info.group_lens)
        logger.info(
            "[RDT-PLAN] pre-built static call plan at init: %d chunks",
            len(self._cached_plan.chunks),
        )

    def _preregister_at_init(self) -> None:
        """Register every NIXL buffer this worker will use at init, before any
        transfer runs, so nothing registers during the sync-0 RDMA churn.

        Both sides are sized from the static plan: receive buffers are
        ``ring_depth`` slots at the largest chunk's ``pack_bytes``, and each bound
        producer is asked to pre-register a serve ring at the max bytes this
        consumer will pull from it. A no-op when this worker has no chunks."""
        plan = self._cached_plan
        if plan is None or not plan.chunks:
            return
        from ray.experimental import register_nixl_memory

        # (a) consumer receive buffers — one per ring slot, at the largest chunk.
        max_pack = max(c.pack_bytes for c in plan.chunks)
        alloc = buffer_alloc_bytes(max_pack, self._buffer_presize)
        for slot in range(self._ring_depth):
            buffer = self._dest_buffers[slot].get(torch.uint8)
            if buffer is None or buffer.numel() < alloc:
                buffer = torch.empty(alloc, dtype=torch.uint8, device=self.device)
                register_nixl_memory(buffer)
                self._dest_buffers[slot][torch.uint8] = buffer

        # (b) producer serve rings — max bytes this consumer pulls from each
        # bound producer, plus the per-producer plan digest a producer sharing
        # one ring across deployments checks them against.
        assert self._router is not None
        serve_bytes = [0] * self._router.num_producers
        per_owner_keys: list[list] = [[] for _ in range(self._router.num_producers)]
        for c in plan.chunks:
            serve_bytes[c.owner] = max(serve_bytes[c.owner], c.pack_bytes)
            per_owner_keys[c.owner].append(c.keys)
        digests = [_plan_digest(keys) for keys in per_owner_keys]
        import ray

        refs = self._router.reserve_serve_buffers(serve_bytes, digests)
        if refs:
            ray.get(refs)  # block until every serve ring is registered
        logger.info(
            "[RDT-PLAN] pre-registered %d receive slots (%.0f MiB each) + serve "
            "rings on %d producer(s) %s",
            self._ring_depth,
            alloc / (1 << 20),
            len(refs),
            [nb // (1 << 20) for nb in serve_bytes],
        )

    def _global_worker_index(self) -> int:
        """This worker's stable, distinct global index across the inference fleet:
        ``data_parallel_index * world_size + rank`` over the TP*PP world.

        ``data_parallel_index``, not ``data_parallel_rank``: vLLM resets the
        latter to 0 in a dense worker but keeps the former as the distinct global
        DP rank. Same formula as the sibling ``nccl_engine``, so dense-via-TP and
        MoE-via-DP+EP both yield distinct 0..C-1."""
        pc = self.parallel_config
        return pc.data_parallel_index * pc.world_size + pc.rank  # world_size = TP*PP

    def _num_consumers(self) -> int:
        """Total inference-worker count. Prefers the driver-supplied
        ``init_info.num_consumers`` (authoritative -- the driver knows the whole
        fleet); else ``world_size_across_dp``, the same stride
        ``_global_worker_index`` indexes with, so the two agree at any pp."""
        if self._num_consumers_override > 0:
            return self._num_consumers_override
        return self.parallel_config.world_size_across_dp

    def start_weight_update(self) -> None:
        """Put the model's params on meta so layerwise reload streams them in
        as each layer's slices land. Baked replay uses checkpoint format."""
        from vllm.model_executor.model_loader.reload import (
            initialize_layerwise_reload,
        )

        initialize_layerwise_reload(self.model)

    def finish_weight_update(self) -> None:
        """Drain the deferred pull/process pipeline (so every layer is fully
        loaded) before finalizing the layerwise reload."""
        from vllm.model_executor.model_loader.reload import (
            finalize_layerwise_reload,
        )

        self.drain_pending()
        finalize_layerwise_reload(self.model, self.model_config)

    def update_weights(self, update_info: dict[str, Any]) -> None:
        """Receive one update. Unlike the base, does NOT issue a per-update
        device sync: post-processing is deferred to background threads and a
        sync here would block on them and serialize the pull/process pipeline.
        Completion is guaranteed by ``drain_pending`` in
        ``finish_weight_update``."""
        self.receive_weights(self.parse_update_info(update_info))

    def receive_weights(
        self,
        update_info: ShardedRDTWeightTransferUpdateInfo,
    ) -> None:
        """Pull + replay the baked leaf modules the sync covers.

        The chunk/free plan is STATIC across syncs — a pure function of the baked
        plan and the driver's group partition — so it was built once at init and
        every sync just re-runs the pipeline over its self-describing chunks, with
        no per-sync bookkeeping and an empty ``update_info``.

        Assumes each baked module's source names fall within one gather group,
        which the per-layer / pre / post partition guarantees (a leaf module's
        sources all live in one decoder layer). A module that did span groups
        would be planned once per group and could pull a name whose group the
        pipeline already freed, which parks the pull until the producer's stall
        watchdog fires.
        """
        del update_info  # the plan is static; nothing arrives per sync
        if self._router is None:
            raise RuntimeError(
                "Sharded RDT engine not initialized. Call init_transfer_engine() first."
            )
        # Surface any error the background thread hit on a prior item promptly.
        self._raise_proc_error()
        if self._cached_plan is None:
            raise RuntimeError(
                "Sharded RDT engine has no call plan: init_info.group_lens must "
                "be supplied at init_transfer_engine()."
            )
        self._run_chunk_pipeline(self._cached_plan)

    def _build_fake_weights(
        self,
        names: list[str],
        sink: "BakeSink",
        device: torch.device,
    ) -> list[tuple[str, torch.Tensor]]:
        """Zero-storage lazies for ``names``, dtype/shape from the init metadata,
        all feeding the bake's recording sink."""
        return [
            (
                name,
                FakeRDTTensor(
                    name=name,
                    shape=torch.Size(self._name_meta[name][1]),
                    dtype=_dtype_from_name(self._name_meta[name][0]),
                    device=device,
                    sink=sink,
                ),
            )
            for name in names
        ]

    def _issue_pull(self, chunk: "_Chunk", slot: int) -> "_PendingPull":
        """Reserve ``slot``, lay the targets out in its buffer, dispatch the
        produce RPC and point the transfer at the buffer — WITHOUT the blocking
        ``ray.get`` (that is ``_complete_pull``). The chunked pipeline issues
        chunk i+1 before completing chunk i, so the producer serves the next
        chunk while the in-flight RDMA streams.

        Slot-reuse guard, both stages required: a generation wait (the CUDA event
        binds only to its LAST record, so synchronizing before the background
        thread recorded this item's event passes silently — observed as
        nondeterministic weight corruption), then the event synchronize. It must
        precede ``set_target_for_ref``, not just the get: the transfer may start
        any time after the metadata push. See the doc's "slot generation
        handshake".
        """
        from ray.experimental import register_nixl_memory, set_target_for_ref

        if self._slot_read_done:
            with self._slot_cv:
                while self._slot_done[slot] < self._slot_queued[slot]:
                    self._slot_cv.wait(timeout=1.0)
            self._slot_read_done[slot].synchronize()

        # Every slice is byte-packed into ONE uint8 buffer, 16B-aligned in
        # keys order, under one NIXL descriptor. ``chunk.pack_layout`` mirrors the
        # producer's rule byte-exactly; ``targets`` carves the dtype views back out.
        keys = chunk.keys
        cur = chunk.pack_bytes
        buffers = self._dest_buffers[slot]
        buffer = buffers.get(torch.uint8)
        if buffer is None or buffer.numel() < cur:
            # Size ONCE with headroom: Ray's desc cache is keyed by data_ptr and
            # outlives its tensor, so a regrowth can false-hit a recycled pointer
            # and skip registering the new extent -> NIXL_ERR_NOT_FOUND.
            alloc = buffer_alloc_bytes(cur, self._buffer_presize)
            buffer = torch.empty(alloc, dtype=torch.uint8, device=self.device)
            register_nixl_memory(buffer)
            buffers[torch.uint8] = buffer
        targets = self._pull_targets(chunk, slot, buffer)
        # One produce RPC serves the whole packed chunk, since the group's owner
        # holds every slice. The buffer view must stay strongly referenced through
        # the get: set_target_for_ref stores WEAKREFS, and a dropped target
        # reroutes the transfer into a fallback buffer.
        blob = buffer[:cur]
        assert self._router is not None
        ref = self._router.pull(chunk.owner, keys, chunk.seq)
        set_target_for_ref(ref, [blob])
        return _PendingPull(
            ref=ref,
            keys=keys,
            targets=targets,
            blob=blob,
            slot=slot,
        )

    def _pull_targets(
        self, chunk: "_Chunk", slot: int, buffer: torch.Tensor
    ) -> "list[torch.Tensor]":
        """Per-key dtype views into ``slot``'s buffer.

        The packed layout is static, so the views are built once per (chunk, slot)
        rather than once per pull -- rebuilding them cost ~1150 Python ops per pull
        at 235B. Keyed on the buffer pointer as well, so a regrow invalidates
        instead of handing back views into a freed buffer.
        """
        cached = self._targets_cache.get((id(chunk), slot))
        if cached is not None and cached[0] == buffer.data_ptr():
            return cached[1]
        targets = [
            buffer[off : off + n * dt.itemsize].view(dt).reshape(shape)
            for off, dt, n, shape in chunk.pack_layout
        ]
        self._targets_cache[(id(chunk), slot)] = (buffer.data_ptr(), targets)
        return targets

    def _complete_pull(self, pending: "_PendingPull") -> "dict[FetchKey, torch.Tensor]":
        """Blocking half of a pull: the NIXL read lands during this ``ray.get``."""
        import ray

        ray.get(pending.ref)
        return dict(zip(pending.keys, pending.targets))

    # ---------------- Bake (dry run, at init) / replay ----------------

    def _bake(self, init_info: ShardedRDTWeightTransferInitInfo) -> None:
        """Bake the replay plan once, as a self-driven meta dry run.

        Puts the params on meta, then drives ``model.load_weights`` over
        ``init_info.names`` through the model's ORIGINAL loaders (the stamps
        bypass ``online_process_loader``, so ``_layerwise_process`` is never in
        the path). Nothing materializes or pulls; the fake's ``copy_`` records
        the source op chain and the meta destination's geometry. Afterwards one
        scatter list per fully-loaded leaf module (copied numel == loadable
        size) is indexed by source name; a partial or unrecordable module fails
        the plan build. The model is restored.

        This leans on layerwise internals a public API should expose first-class:
        a currently-loading hook instead of monkeypatched stamps, a dry-run mode
        instead of bypassing ``online_process_loader``, and an
        ``abort_layerwise_reload`` instead of ``_restore_after_dry_run``.
        """
        from vllm.model_executor.model_loader.reload.layerwise import (
            initialize_layerwise_reload,
        )
        from vllm.model_executor.model_loader.reload.utils import get_layer_size

        names, dtype_names, shapes = (
            init_info.names,
            init_info.dtype_names,
            init_info.shapes,
        )
        self._name_meta = {n: (d, s) for n, d, s in zip(names, dtype_names, shapes)}
        if not names:
            return

        model = self.model
        recorder = BakeSink()

        _t0 = time.perf_counter()
        with torch.device(self.device):
            # Meta-restore params + save kernel tensors (we bypass the loader
            # wrapping it installs, below).
            initialize_layerwise_reload(model)
            # Stamp the *original* loaders (bypassing online_process_loader), so
            # the single load pass runs the loaders on meta and records via the
            # fake's copy_ — with no inline _layerwise_process, no deferral.
            self._install_recording_stamps(model, recorder)
            model.load_weights(self._build_fake_weights(names, recorder, self.device))
            # Keep only fully-loaded modules (copied numel >= loadable size, the
            # test online_process_loader uses): a partial module leaves unwritten
            # regions that finalize would init, so baking it scatters garbage.
            for module, recorded in recorder.copies_by_layer.items():
                if not recorded or any(c is None for c in recorded):
                    continue  # unrecordable copy_ -> slow path
                # Guard above guarantees every entry is a real _Scatter.
                copies = cast("list[_Scatter]", recorded)
                copied = sum(prod(c.shape) for c in copies)
                if copied < get_layer_size(module):
                    continue  # partial -> slow path
                for c in copies:
                    # Every name of the module shares ONE list, so identity
                    # dedups them back to one module at plan time.
                    self._name_to_plan[c.src[0]] = copies
            self._restore_after_dry_run(model)

        # Names whose copy_ fired during the bake. Names not in here no-op for
        # this worker (e.g. foreign-EP experts) and are skipped; live names that
        # did not bake fail the plan build.
        self._live_names = set(recorder.copied_names)

        n_groups = len({id(g) for g in self._name_to_plan.values()})
        logger.info(
            "Sharded RDT dry-run baked %d/%d names into %d leaf modules "
            "(%d live) in %.3fs",
            len(self._name_to_plan),
            len(names),
            n_groups,
            len(self._live_names),
            time.perf_counter() - _t0,
        )

    def _install_recording_stamps(
        self, model: torch.nn.Module, recorder: "BakeSink"
    ) -> None:
        """Wrap each loadable param's ``weight_loader`` to stamp
        ``recorder.current = (leaf_module, param_name)`` before delegating to the
        original loader, so the fake's ``copy_`` can attribute each recorded copy.
        ``functools.wraps`` keeps the loader's real signature (so vLLM's
        ``_layerwise_process`` ``param`` redirect still works if a stamp leaks),
        and ``_rdt_stamp_inner`` tags it so ``_restore_after_dry_run`` can unwrap it.
        """
        import functools

        from vllm.model_executor.model_loader.reload.layerwise import (
            _get_original_loader,
        )
        from vllm.model_executor.model_loader.reload.utils import get_layer_tensors
        from vllm.model_executor.model_loader.weight_utils import default_weight_loader

        def _make_stamp(layer, name, inner, added=False):
            @functools.wraps(inner)  # keep ``inner``'s signature (incl. ``param``)
            def stamp(*args, **kwargs):
                recorder.current = (layer, name)
                try:
                    return inner(*args, **kwargs)
                finally:
                    recorder.current = None

            # Tag so _restore_after_dry_run can detect and unwrap leaked stamps,
            # and so a second bake doesn't double-wrap.
            stamp._rdt_stamp_inner = inner  # type: ignore[attr-defined]
            stamp._rdt_stamp_added = added  # type: ignore[attr-defined]
            return stamp

        for module in model.modules():
            for name, tensor in get_layer_tensors(module).items():
                if getattr(tensor, "weight_loader", None) is None:
                    # A param with NO loader (e.g. GLM's router bias, a plain
                    # nn.Parameter) is still loaded by the model's load_weights
                    # through the getattr(param, "weight_loader",
                    # default_weight_loader) fallback — unstamped, its bake copy
                    # would be unattributable, failing the module's coverage
                    # gate and the plan build. Stamp the same default loader
                    # the fallback would pick; the restore deletes it again.
                    tensor.weight_loader = _make_stamp(
                        module, name, default_weight_loader, added=True
                    )
                    continue
                # Bypass online_process_loader: stamp the *original* loader.
                original = _get_original_loader(tensor)
                tensor.weight_loader = _make_stamp(module, name, original)

    def _restore_after_dry_run(self, model: torch.nn.Module) -> None:
        """Restore each layerwise layer's saved kernel tensors without pulling
        (a real ``finalize_layerwise_reload`` would materialize/load) and reset
        its info. Also unwrap any recording ``stamp`` left on the params, since a
        leaked stamp would sit under the next sync's ``online_process_loader`` and
        silently break ``_layerwise_process``'s ``param`` redirect.
        """
        from vllm.model_executor.model_loader.reload.layerwise import (
            LAYERWISE_INFO,
            _place_kernel_tensors,
        )
        from vllm.model_executor.model_loader.reload.utils import get_layer_tensors

        for layer in model.modules():
            info = LAYERWISE_INFO.get(layer)
            if info is not None and info.can_load():
                if info.kernel_tensors is not None:
                    _place_kernel_tensors(layer, info)
                info.reset()
        # Unwrap any recording stamps left on the (now-restored) params so they
        # never leak into a later update_weights. ``_rdt_stamp_inner`` is set by
        # ``_install_recording_stamps``; unwrap repeatedly in case of nesting.
        for module in model.modules():
            for _name, tensor in get_layer_tensors(module).items():
                loader = getattr(tensor, "weight_loader", None)
                added = False
                while loader is not None and hasattr(loader, "_rdt_stamp_inner"):
                    added = added or getattr(loader, "_rdt_stamp_added", False)
                    loader = loader._rdt_stamp_inner
                    tensor.weight_loader = loader
                if added:
                    # The stamp was ATTACHED to a param that had no loader
                    # (see _install_recording_stamps); leave none behind.
                    del tensor.weight_loader
        if hasattr(model, "_original_do_torchao_reload"):
            model._do_torchao_reload = model._original_do_torchao_reload

    # ---------------- Background post-processing (pull/process pipeline) -------

    def _ensure_proc_worker(self) -> None:
        """Lazily create the per-slot events, the background CUDA stream, the work
        queue, and the single processing thread. Idempotent."""
        if self._proc_thread is not None:
            return
        import queue
        import threading

        self._slot_read_done = [torch.cuda.Event() for _ in range(self._ring_depth)]
        # Generation handshake for the events above: the RPC thread counts items
        # queued per slot, this thread counts records, and a pull may synchronize()
        # only once done[slot] has caught up with queued[slot]. Without it the
        # synchronize can bind to a stale record and pass, letting the next RDMA
        # overwrite a slot under a pending scatter. See the doc's "The slot
        # generation handshake".
        self._slot_queued = [0] * self._ring_depth
        self._slot_done = [0] * self._ring_depth
        self._slot_cv = threading.Condition()
        self._proc_stream = torch.cuda.Stream(device=self.device)
        self._proc_queue = queue.Queue()
        self._proc_error = None
        t = threading.Thread(
            target=self._proc_worker_loop, name="rdt-postprocess", daemon=True
        )
        self._proc_thread = t
        t.start()
        self._quant_stream = torch.cuda.Stream(device=self.device)
        self._quant_queue = queue.Queue()
        qt = threading.Thread(
            target=self._quant_worker_loop, name="rdt-quant", daemon=True
        )
        self._quant_thread = qt
        qt.start()

    def _proc_worker_loop(self) -> None:
        """Single persistent thread: run each queued item's process phase on the
        background stream. Exits on the ``None`` sentinel (shutdown). An item that
        raises is recorded in ``_proc_error`` and re-raised on the RPC thread /
        at drain, so a failed sync fails loudly rather than corrupting silently.
        """
        torch.cuda.set_device(self.device)
        q = self._proc_queue
        assert q is not None
        while True:
            item = q.get()
            try:
                if item is None:
                    return
                # _process_item publishes the slot generation internally (in a
                # finally around its scatter pass), so an error here still
                # unblocks the RPC thread's generation wait.
                self._process_item(item)
            except BaseException as e:  # noqa: BLE001 - surfaced on the RPC thread
                self._proc_error = e
                logger.exception("RDT background post-processing failed")
            finally:
                q.task_done()

    def _mark_slot_done(self, slot: int) -> None:
        """Publish that a queued item's read-done event has been recorded (or the
        item failed) so a pull waiting to reuse ``slot`` can proceed to its
        CUDA-event synchronize."""
        with self._slot_cv:
            self._slot_done[slot] += 1
            self._slot_cv.notify_all()

    def _raise_proc_error(self) -> None:
        """Re-raise (once) any error captured by the background thread."""
        if self._proc_error is not None:
            err = self._proc_error
            self._proc_error = None
            raise RuntimeError("RDT background post-processing failed") from err

    def drain_pending(self) -> None:
        """Block until the background thread has processed every queued item and
        its stream work is complete, then re-raise any error it hit. Called from
        the worker's ``finish_weight_update`` before ``finalize_layerwise_reload``
        so every baked layer is fully loaded (and ``info.reset()``-ed) first."""
        if self._proc_queue is not None:
            self._proc_queue.join()  # every put() item task_done()'d
        # The scatter thread feeds the quant thread, so join it SECOND (all
        # completed-group batches have been put by now), then sync both streams
        # so finalize sees fully-materialized, quanted, reset layers.
        if self._quant_queue is not None:
            self._quant_queue.join()
        if self._proc_stream is not None:
            self._proc_stream.synchronize()
        if self._quant_stream is not None:
            self._quant_stream.synchronize()
        self._raise_proc_error()
        # Ensure every fired free_group signal has EXECUTED on
        # the producer before the sync ends: ``begin_sync`` resets the producer's
        # per-group signal counts, so a signal landing after the next sync
        # started would credit a group it does not belong to, over-crediting the
        # gather-lookahead backpressure (extra gather resident -> trainer OOM
        # risk).
        if self._pending_frees:
            import ray

            try:
                ray.get(self._pending_frees)
            finally:
                self._pending_frees.clear()
        # One sync iteration fully drained; buffers/registrations are at (or
        # nearer) high-water — the chunk pipeline may issue ahead from now on.
        self._completed_syncs += 1

    def _dispatch_item(self, item: "_ProcItem") -> None:
        """Hand one chunk item to the background scatter thread.

        Counts the item against its slot BEFORE dispatch: the next pull into
        that slot must wait until the background thread has processed (and
        RECORDED the read-done event for) every item ever queued on it.
        """
        with self._slot_cv:
            self._slot_queued[item.slot] += 1
        assert self._proc_queue is not None
        self._proc_queue.put(item)

    def _chunk_module_scatters(
        self, modules: "list[list[_Scatter]]"
    ) -> "list[tuple[int, list[_Scatter]]]":
        """Cut the modules' copies into one chunk per distinct owner class
        present, ascending by class index, as ``(class_idx, scatters)`` pairs.

        A chunk is one packed pull, so every name in it must share a producer —
        which is exactly what an owner class is. The cut is a pure function of
        the bake and the ownership table. Copy order within a chunk is bake
        order, and a module's copies may span chunks (materialize/quant fire on
        its first/last chunk; see ``_build_call_plan``)."""
        assert self._router is not None
        by_class: dict[int, list[_Scatter]] = {}
        for copies in modules:
            for c in copies:
                by_class.setdefault(self._router.class_of(c.src[0]), []).append(c)
        return [(ci, by_class[ci]) for ci in sorted(by_class)]

    def _build_call_plan(self, names: list[str], group_lens: list[int]) -> "_CallPlan":
        """Build the STATIC plan for one whole-sync call.

        Pure — no pulls, no engine state touched — so the result is cached and
        reused every sync. Three passes:
          1. Split ``names`` into gather groups, one chunk per owner class
             present in this worker's baked copies, recording each group's last
             chunk for its ``free_group`` signal (or ``pre_free`` when this worker
             pulls nothing from it). Every group is signaled exactly once by
             construction. The stream has no per-group call boundaries, so group
             L+1's first chunk issues while L's still stream.
          2. Per leaf module, find its FIRST and LAST chunk — materialize on the
             first, quant/kernel/reset on the last, correct by construction
             instead of by runtime counters.
          3. Assemble ``_Chunk``s: dedup keys and precompute the packed layout.
        """
        router = self._router
        assert router is not None, "init_transfer_engine() must run before planning"

        # --- pass 1: gather groups -> per-owner-class scatter chunks -----------
        raw_chunks: list[list[_Scatter]] = []
        free_at: dict[int, list[int]] = {}  # chunk idx -> groups to signal after it
        pre_free: list[int] = []  # groups with no chunk here (signal at start)
        unbaked: list[str] = []
        pos = 0
        for gi, glen in enumerate(group_lens):
            gnames = names[pos : pos + glen]
            pos += glen
            modules: list[list[_Scatter]] = []
            seen: set[int] = set()
            for n in gnames:
                mod = self._name_to_plan.get(n)
                if mod is None:
                    if n in self._live_names:
                        unbaked.append(n)
                elif id(mod) not in seen:
                    seen.add(id(mod))
                    modules.append(mod)
            if not modules:
                # Nothing to pull for this group on this worker; its owners
                # still published it, so signal it done at sync start (they
                # tolerate a signal that arrives before the publish).
                pre_free.append(gi)
                continue
            for _ci, scatters in self._chunk_module_scatters(modules):
                raw_chunks.append(scatters)
            free_at.setdefault(len(raw_chunks) - 1, []).append(gi)
        if unbaked:
            # A live name with no baked plan cannot be loaded: there is no
            # fallback, and its pull would target groups the pipeline has
            # already freed. Fail here, where the names are known.
            raise RuntimeError(
                f"Sharded RDT: {len(unbaked)} live name(s) did not bake, "
                f"first: {unbaked[:3]}. Fix the loader so these names bake "
                "(fully load their module in one pass), or exclude them from "
                "the synced set."
            )

        # --- pass 2: per-module first/last chunk -> materialize/quant ----------
        first_at: dict[int, int] = {}
        last_at: dict[int, int] = {}
        layer_by_id: dict[int, Any] = {}
        for ci, scatters in enumerate(raw_chunks):
            for sc in scatters:
                lid = id(sc.layer)
                layer_by_id[lid] = sc.layer
                first_at.setdefault(lid, ci)
                last_at[lid] = ci
        materialize_at: dict[int, list[Any]] = {}
        quant_at: dict[int, list[Any]] = {}
        for lid, ci in first_at.items():
            materialize_at.setdefault(ci, []).append(layer_by_id[lid])
        for lid, ci in last_at.items():
            quant_at.setdefault(ci, []).append(layer_by_id[lid])

        # --- pass 3: assemble _Chunks (dedup keys + precompute pack layout) ----
        from collections import defaultdict

        chunks: list[_Chunk] = []
        per_owner_seq: dict[int, int] = defaultdict(int)
        for ci, scatters in enumerate(raw_chunks):
            keys: list[FetchKey] = []
            kmeta: dict[FetchKey, tuple[torch.dtype, tuple[int, ...]]] = {}
            for sc in scatters:
                if sc.src not in kmeta:
                    kmeta[sc.src] = (sc.dtype, sc.shape)
                    keys.append(sc.src)
            pack_layout: list[tuple[int, torch.dtype, int, tuple[int, ...]]] = []
            cur = 0
            for k in keys:
                dt, shape = kmeta[k]
                numel = prod(shape) or 1  # type: ignore[arg-type]
                off = (cur + 15) & ~15
                pack_layout.append((off, dt, numel, shape))
                cur = off + numel * dt.itemsize
            owner = router.producer_for(router.consumer_id, scatters[0].src[0])
            chunks.append(
                _Chunk(
                    seq=per_owner_seq[owner],
                    scatters=scatters,
                    keys=keys,
                    pack_layout=pack_layout,
                    pack_bytes=cur,
                    materialize=materialize_at.get(ci, []),
                    quant=quant_at.get(ci, []),
                    free=free_at.get(ci, []),
                    # Every name of a chunk shares an owner class and a group,
                    # so any of them resolves the same producer.
                    owner=owner,
                )
            )
            per_owner_seq[owner] += 1
        return _CallPlan(chunks=chunks, pre_free=pre_free)

    def _signal_group_done(self, group_idx: int) -> None:
        """Fire-and-forget ``free_group`` signal at EVERY owner of the group.

        The per-group barrier: each owner counts one signal per live consumer
        and frees the group (releasing its lookahead credit) on the last one, so
        every owner must hear from every consumer — including owners this worker
        pulls nothing from. Refs are held and drained in ``drain_pending`` so
        every signal has EXECUTED before the sync ends: ``begin_sync`` resets
        the counters, and a straggler landing in the next sync would credit a
        group it does not belong to."""
        assert self._router is not None
        self._pending_frees.extend(self._router.free_group(group_idx))

    def _run_chunk_pipeline(self, plan: "_CallPlan") -> None:
        """Pipelined chunk pulls over the ring of receive slots.

        Issues up to ``ring_depth`` produce RPCs ahead of the blocking gets, so
        while chunk i's RDMA streams the producer serves i+1 into its own ring
        slot and the background thread scatters i-1 out of another. Reads stay
        serialized on the shared NIC — the bandwidth floor, not a loss.

        Slot safety rests on two arguments spelled out in the doc: the producer's
        ring is no shallower than this one and drain-before-issue orders its
        reuse; the consumer's slots are held by ``_issue_pull``'s generation
        handshake.
        """
        from collections import deque

        inflight: deque[tuple[_PendingPull, _Chunk]] = deque()

        # Gather groups with no chunk on this worker (plan.pre_free): signal
        # them done before the pipeline — their owners still published them, and
        # a signal that arrives before its publish is counted (fire-and-forget).
        for gi in plan.pre_free:
            self._signal_group_done(gi)

        def drain_one() -> None:
            pending, chunk = inflight.popleft()
            results = self._complete_pull(pending)
            self._dispatch_item(
                _ProcItem(chunk=chunk, results=results, slot=pending.slot)
            )
            # Each gather group whose LAST chunk this is: its read is done, so
            # this consumer is finished with the group -> signal every owner
            # (fire-and-forget, off the critical path).
            for gi in chunk.free:
                self._signal_group_done(gi)

        for chunk in plan.chunks:
            if not chunk.keys:
                # _chunk_module_scatters never emits empty chunks; keep the
                # signals safe anyway if this ever changes.
                for gi in chunk.free:
                    self._signal_group_done(gi)
                continue
            # Drain BEFORE issue once the ring is full: frees this chunk's slot
            # (generation-wise) and guarantees the producer-slot invariant above.
            # Sync 0 runs SERIAL (max 1 in flight): both sides still grow and
            # register buffers, and a producer registration mid-flight churns its
            # agent-metadata version under an in-flight pull (see __init__).
            depth = 1 if self._completed_syncs == 0 else self._ring_depth
            if len(inflight) >= depth:
                drain_one()
            slot = self._pull_slot
            self._pull_slot = (slot + 1) % self._ring_depth
            inflight.append((self._issue_pull(chunk, slot), chunk))
        while inflight:
            drain_one()

    def _process_item(self, item: "_ProcItem") -> None:
        """Scatter-thread half: materialize this chunk's first-seen modules,
        scatter its slices on the process stream, publish the slot, then hand the
        modules it COMPLETES to the quant thread.

        Mirrors ``_layerwise_process`` minus the loader replay. Once every scatter
        reading ``item.slot`` is enqueued, records the slot's read-done event so
        the RPC thread can block on it before overwriting the slot.
        """
        from vllm.model_executor.model_loader.reload.layerwise import (
            LAYERWISE_INFO,
        )
        from vllm.model_executor.model_loader.reload.meta import materialize_layer

        results = item.results
        chunk = item.chunk
        with (
            torch.cuda.device(self.device),
            torch.cuda.stream(self._proc_stream),
            torch.device(self.device),
        ):
            # PASS 1 — slot readers: materialize the modules whose FIRST scatter
            # is in this chunk (empty HF params, once per module by construction)
            # then scatter this chunk's copies. The scatter copies are the ONLY
            # reads of the receive buffer; quant and kernel-copy operate on the
            # scattered params. Releasing the slot right after the scatters lets
            # the NEXT chunk's RDMA overwrite the buffer while quant still runs.
            try:
                for layer in chunk.materialize:
                    info = LAYERWISE_INFO.get(layer)
                    if info is None or not info.can_load():
                        raise RuntimeError(
                            f"Baked replay: layer {type(layer).__name__} "
                            "was not set up for reload this sync "
                            "(start_weight_update must run before "
                            "update_weights)."
                        )
                    materialize_layer(layer, info)
                for sc in chunk.scatters:
                    param = getattr(sc.layer, sc.param_name)
                    # Through .data, exactly like default_weight_loader: params
                    # that keep requires_grad=True (plain nn.Parameters like
                    # GLM's router bias) reject an in-place copy through an
                    # autograd view of the leaf.
                    dst = param.data.as_strided(sc.shape, sc.stride, sc.offset)
                    with torch._C.DisableTorchFunctionSubclass():
                        dst.copy_(results[sc.src])
                # All reads of this slot's buffer are now enqueued on the process
                # stream; record + publish so the RPC thread can reuse the slot.
                self._slot_read_done[item.slot].record(self._proc_stream)
            finally:
                # Publish even on error so the RPC thread's generation wait
                # unblocks (the failure itself surfaces via _raise_proc_error).
                self._mark_slot_done(item.slot)

            # PASS 2 — param readers only: quant / kernel-copy / reset for the
            # modules whose LAST scatter is in this chunk. Handed to the
            # DEDICATED quant thread (own CUDA stream, event-chained after this
            # chunk's scatters) so it never delays this thread's next pass-1.
            # Running it in order here stalls the RPC thread's slot handshake by
            # ~0.5-0.75s/iter, since every group's quant delays the next publish.
            if chunk.quant:
                ready = torch.cuda.Event()
                ready.record(self._proc_stream)
                if self._quant_queue is None:
                    self._run_quant(chunk.quant, ready)
                else:
                    self._quant_queue.put((chunk.quant, ready))

    def _run_quant(self, layers: "list[Any]", ready: "torch.cuda.Event") -> None:
        """Quant/kernel-copy/reset the given COMPLETED leaf modules, exactly as
        _layerwise_process. Runs on the quant thread's own stream, ordered after
        the modules' scatters via ``ready``; touches only the scattered params
        (never a receive slot), so it can overlap subsequent chunks' RDMA and
        scatters. ``info.reset()`` is what makes finalize skip the layer —
        drain_pending joins the quant queue before finalize runs."""
        from vllm.model_executor.layers.quantization.base_config import (
            QuantizeMethodBase,
        )
        from vllm.model_executor.model_loader.reload.layerwise import (
            LAYERWISE_INFO,
            _copy_and_restore_kernel_tensors,
        )

        stream = self._quant_stream or self._proc_stream
        assert stream is not None  # created in _ensure_proc_worker before use
        with (
            torch.cuda.device(self.device),
            torch.cuda.stream(stream),
            torch.device(self.device),
        ):
            stream.wait_event(ready)
            for layer in layers:
                info = LAYERWISE_INFO.get(layer)
                assert info is not None  # completed leaf module is set up for reload
                quant_method = getattr(layer, "quant_method", None)
                if isinstance(quant_method, QuantizeMethodBase):
                    if hasattr(layer, "_already_called_process_weights_after_loading"):
                        delattr(layer, "_already_called_process_weights_after_loading")
                    quant_method.process_weights_after_loading(layer)
                # Copy into persistent kernel storage (preserves cudagraph refs).
                if info.kernel_tensors is not None:
                    _copy_and_restore_kernel_tensors(layer, info)
                # Reset so finalize_layerwise_reload skips this (loaded) layer.
                info.reset()

    def _quant_worker_loop(self) -> None:
        """Dedicated quant thread: drains (completed_modules, scatter-done event)
        batches. Errors surface via _proc_error like the scatter thread's."""
        torch.cuda.set_device(self.device)
        q = self._quant_queue
        assert q is not None
        while True:
            batch = q.get()
            try:
                if batch is None:
                    return
                self._run_quant(*batch)
            except BaseException as e:  # noqa: BLE001 - surfaced on the RPC thread
                self._proc_error = e
                logger.exception("RDT quant thread failed")
            finally:
                q.task_done()

    def shutdown(self) -> None:
        # Stop the background post-processing thread (drain, then sentinel + join)
        # before dropping the state it touches.
        if self._proc_thread is not None:
            try:
                self.drain_pending()
            except Exception:
                logger.exception("RDT drain during shutdown failed")
            assert self._proc_queue is not None
            self._proc_queue.put(None)  # sentinel
            self._proc_thread.join(timeout=30)
            self._proc_thread = None
            self._proc_queue = None
            self._proc_stream = None
        if self._quant_thread is not None:
            assert self._quant_queue is not None
            self._quant_queue.put(None)  # sentinel
            self._quant_thread.join(timeout=30)
            self._quant_thread = None
            self._quant_queue = None
            self._quant_stream = None
        self._slot_read_done = []
        self._router = None
        # Drop strong references to baked modules so the model can be freed.
        self._name_to_plan.clear()
        self._name_meta.clear()
        self._live_names.clear()
        # Drop the cached plan (holds _Scatter refs to the baked layers).
        self._cached_plan = None
        # Release the receive buffers (their NIXL registration is pinned for the
        # process lifetime; freeing the tensors just drops our strong refs).
        self._dest_buffers = [{} for _ in range(self._ring_depth)]