Disaggregated inference: prefill on GPU 0, decode on GPU 1, from one program text

The second half of hiraditya/Vx.1#319. hiraditya/Vx.1#319 proves the compiler can admit a placement before the machine is rented and that the admitted build runs. This issue proves the placement can name more than one device, which is the claim that Vx programs a distributed system directly rather than emitting code for a single accelerator.

Disaggregation, not sharding. hiraditya/Vx.1#319 puts multi-GPU sharding out of scope and that stands: no tensor is split across devices, and the V1 modelling rule in hiraditya/Vx.1#284 is untouched. Prefill/decode disaggregation replicates the model on two workers and moves the KV cache between them — the shape DistServe and NVIDIA Dynamo use, and the one that needs no per-instance memory accounting.

The shape

        prompt tokens                          generated tokens
             |                                        |
     +-------v--------+     KV cache      +------------v-----------+
     |  prefill       | ----------------> |  decode                |
     |  Topology::GPU[0] |  peer transfer |  Topology::GPU[1]      |
     |  weights resident |                |  weights resident      |
     +----------------+                   +------------------------+

One function body, instantiated twice:

fn transformer<const D : i32>(...) { ... matmul<D>(...) ... }

fn matmul<const D : i32>(...) -> i32 {
  spawn on(Topology::GPU[D]) { matmul_into(&mut yt, &wt, &xt); }
}

Verified working already: Topology::substitute handles GPU, a const generic forwards through a call chain into a spawn placement, and outer<0>/outer<1> emit topology(500)/topology(501). Only the index differs, which is the same discipline as "only --machine differs" one level up.

What has to be built

• hiraditya/Vx.1#346 — the CUDA plugin's single-device state. The cuBLAS handle is process-wide and bound to the device current at creation, vx_plugin_alloc_and_transfer discards its topology_id, and there is no device-to-device movement. Every one of these turns a two-GPU run into a one-GPU run that reports two.
• vx_plugin_transfer_peer — the KV cache's edge from prefill's device to decode's.
• A two-GPU fleet model, so admission can be asked about the split rather than only about one device.
• The program: weights staged to both devices, prefill and decode holding separate KV arrays so the handoff is load-bearing.

What makes it honest

The handoff must be load-bearing: decode reads a KV cache it can only have obtained through the transfer. Sharing one host array between the phases and copying it device-to-device on the side would be theatre — the run would produce identical output with the transfer deleted. Two separate arrays means deleting the transfer produces garbage, which is the acceptance test.

Stated limits, to be repeated wherever results are:

• Attention, softmax, RoPE and RMSNorm still run on the host (#251). Only the seven projections and lm_head are on device. The KV cache therefore also lives in host memory as the working copy; its device residency exists for the transfer.
• Making the KV cache genuinely device-resident and read there needs strided sub-views of a placed tensor (#344), because a head's slice of the cache has row stride kv_dim, and vx_gemm_plan_decode refuses non-contiguous rows by design.
• Both phases are batch 1 and f32 (#320, hiraditya/Vx.1#333).

Acceptance

• One program text; prefill and decode differ only by a const-generic device index.
• cudaGetDevice at dispatch shows 0 through prefill and 1 through decode, switching exactly at the phase boundary.
• The KV handoff moves a byte count the compiler predicted, recorded against the observed one.
• Token output is identical to the single-device run and to llama2.c.
• Deleting the handoff changes the output — proof it was load-bearing.

Start with hiraditya/Vx.1#346 and the CUDA plugin's single-device state, then trace Topology::substitute, outer<0>/outer<1>, and the vx_plugin_transfer_peer entry point. Verify the two-GPU acceptance checks: device switching at the phase boundary, predicted versus observed transfer bytes, identical output to the single-device run, and changed output when the handoff is removed.