CUDA plugin keeps one device's worth of state: cuBLAS handle, allocation and D2H all ignore the topology

select_device is correct and its doc comment already says what it is for -- "the whole of running prefill on one device and decode on another". But it is called from exactly one place (vx_plugin_dispatch_async), and the three other entry points that touch device memory do not honour a topology at all. Each is a silent wrong-device bug, and together they mean a two-GPU program would run entirely on GPU 0 while its traces claimed otherwise.

1. The cuBLAS handle is process-wide and bound to whichever device was current at first use

cublasHandle_t cublas_handle() {
  static cublasHandle_t handle = [] { ... cublasCreate(&h); ... }();
  return handle;
}

A cuBLAS handle is tied to the device current at cublasCreate. select_device then calls cudaSetDevice(1) and run_gemm uses a handle belonging to device 0. This is the worst of the three because it is the one that produces numbers: the dispatch appears to target GPU 1, prints [Vx CUDA] device 1, and computes against device 0's context.

Fix: a small per-device handle cache, created lazily after the device is selected.

2. vx_plugin_alloc_and_transfer discards its topology_id

void *vx_plugin_alloc_and_transfer(size_t bytes, void *host_ptr, uint32_t topology_id) {
  (void)topology_id;

The parameter exists, callers pass 500, and llama2.vx's comment says the weights are staged "into the memory of the topology that will read them". They are staged into the memory of whichever device happens to be current. With one GPU that is accidentally right; with two it puts the decode replica's weights on the prefill GPU.

3. vx_plugin_transfer_device_to_host has no topology parameter

int32_t vx_plugin_transfer_device_to_host(void *device_ptr, void *host_ptr, size_t bytes);

Unified addressing makes this mostly work by accident, since the runtime can infer the device from the pointer. It should still name the device it is reading from, both for symmetry with the alloc side and because the ABI is the contract a non-CUDA backend implements, where inference may not be available.

4. Missing: a device-to-device movement

There is no way to express "this buffer is on GPU 0 and must end up on GPU 1". That is the movement a disaggregated placement is made of -- the KV cache produced by prefill has to reach the decode device -- and it is the one edge the plugin ABI cannot currently name.

Proposed:

void *vx_plugin_transfer_peer(void *src_device_ptr, uint32_t src_topology_id,
                              uint32_t dst_topology_id, size_t bytes);

CUDA answers it with cudaMemcpyPeer, which stages through the host by itself when no P2P path exists, so one implementation covers both an NVLink pod and a PCIe-only one. The CPU backends answer it with a memcpy, so the same program stays correct on a machine with no GPUs -- the property the rest of the plugin ABI already has.

Blocks the disaggregated half of hiraditya/Vx.1#319.

Start by tracing select_device from vx_plugin_dispatch_async, then inspect cublas_handle, vx_plugin_alloc_and_transfer, and vx_plugin_transfer_device_to_host across the CUDA and CPU backends. Map how topology IDs reach each entry point and how device-to-device movement should fit the ABI; done means handles, allocation, host transfers, and peer transfers honor the selected topology on both backend types.