Disaggregated inference across machines: prefill/decode split with a Vx-native KV handoff (extends #319)

Demonstrate that Vx programs a distributed inference system directly: prefill on one machine, decode on another, the KV-cache handoff written in Vx source, and every placement admitted at compile time against the target's fleet file. Extends hiraditya/Vx#319 (single-GPU Llama + FlashAttention over the plugin ABI) with the multi-machine phases, and organizes the campaign into two run structures.

Two run structures

1. Flexibility — generality is the metric, not throughput. Vendor libraries own the kernels (#319 Tracks A/B1); Vx owns placement, admission, staging, and the transport. Ends in a heterogeneous relay across {A100, x86, H100, MI300X}.
2. Performance — both-NVIDIA, Vx emits the device kernels itself (#251, generated from flash_attention_v4.vx source), f16 throughout (#320). The bar is batch-1 latency parity with vLLM on the same dtype and model, reported as cost-of-admission — not a throughput contest against batching machinery this deliberately does not have. Different timeline, same harness; only the kernel provenance flips.

What the repo survey established (2026-08-08)

Facts the plan rests on, verified today:

• The compile-time half of disaggregation exists and is tested: placement, capacity admission (E6009/E6010), routed and costed transfers, and a disaggregated prefill/decode modelling precedent in tests/frontend/pass/rubin_disaggregated.vx.
• The execution half is host-CPU only. vx.launch lowers to vx_plugin_dispatch_async (src/dialect/VxLowering.cpp), whose only provider is runtime/npu_dispatch.mm, built on macOS only. On Linux nothing provides the symbol, so a placed program does not link. This is a prerequisite hiraditya/Vx#319 does not list.
• vx.transfer executes as a host memref.alloc + memref.copy; the topology id is ignored at runtime. Device residency is the performance gate (#319's "real design item").
• std::net has TCP listener/stream + UDP backed by rust_core; tests/backend/pass/ffi_tcp_server.vx executes a bind+drop in the JIT harness, but no test yet exercises an accept/read/write loop across processes.
• Vx models one representative device per declared kind; a runtime device index falls back to device 0 with W1030 (tests/frontend/fail/topology_index_constant_and_runtime.vx). Consequence: disaggregation should be process-per-role, each process seeing one GPU — not one process driving two devices.
• dtypes: f32 executes end-to-end; bf16 executes on the AST path; f16 is declaration-only until hiraditya/Vx#320.
• All three in-tree Llama programs are compile-only or bitrotted; none executes in CI. The most idiomatic base is tests/backend/pass/llama2_v2.vx.

Milestones

M0–M3 — single machine (#319's scope, plus the Linux gap):

• M0 Linux bring-up: portable dispatch shim with the libffi CPU fallback buildable on Linux, done on a trusted EC2 x86 build box; un-bitrot one Llama program into an executing test (base: llama2_v2.vx).
• M1 runtime/cuda_dispatch.cpp + src/plugin/cuda.rs (cuBLAS for GEMM-shaped kernels, cuDNN fused SDPA for attention, REQUIRES: cuda test gating) and the hiraditya/Vx#320 dtype fix.
• M2 device residency: route vx.transfer to vx_device_alloc/vx_device_copy when the target memref carries a device address space (#258). Weights H2D exactly once; KV cache device-resident across the decode loop.
• M3 placed Llama (stories15M) admitted via --machine fleet/a100-80.vx, token-for-token parity with the CPU oracle. The sibling A100-40 file gives an admit/reject boundary without renting a second SKU.

M4 — disaggregation (this issue's core):

• One program text; the role (prefill | decode) picked at runtime; one process per role.
• Wire format kept dumb: header {n_layers, kv_dim, pos, dtype} + raw KV bytes + last token, over std::net TCP — the handoff is Vx source, which is the claim.
• M4a CPU↔CPU across two processes (no GPU dependency; also closes the TCP accept/read/write test gap).
• M4b prefill pod → decode pod (A100 ↔ A100 over the real network). Admission per role per machine file; predicted KV-transfer cost (NIC edges as modelled in fleet/node-8gpu.vx) recorded next to the measured one — this feeds the memory-algebra calibration campaign.

M5 — combo relay {H100 → A100 → MI300X → x86}:

• runtime/rocm_dispatch.cpp (hipBLAS port of the CUDA plugin, same ABI), per-SKU admission, the KV cache relayed machine to machine with a few tokens decoded on each.
• Rent H100/MI300X pods only after A100↔A100 works; container scripts staged in advance (utils/memalg/run_m1.sh discipline).

M6 — performance run: hiraditya/Vx.1#251 kernel emission from flash_attention_v4.vx, fused decode kernels (RMSNorm+GEMV, RoPE), f16, tensor cores as stretch; batch-1 latency vs vLLM. Tracked separately; reuses this harness unchanged.

Acceptance

• M4b: same tokens as the single-machine reference; prefill and decode processes each admitted at compile time; KV handoff size and latency logged, predicted vs measured recorded.
• M5: one unchanged program text; four machine files; four admissions; coherent decode continuing across all four machines.
• No performance number published before M2 lands — naive per-spawn H2D/D2H loses to CPU at stories15M sizes.

Deployment discipline

Source does not ship to rented machines. The execution path is already AOT (src/jit.rs produces a native binary) and admission is compile-time, so pods need no compiler, repo, or .vx source: build vxc, the demo binaries, cuda_dispatch.so (CUDA toolkit stubs — no GPU needed to link) and the admission JSONs on the EC2 x86 box, scp one source-free tarball to the pod, pull results back, nuke the directory, terminate the pod. Never git clone on a pod. The EC2 box doubles as the x86 leg of the M5 relay.

Out of scope

Multi-GPU sharding execution (the hiraditya/Vx.1#284/#290 V1 modelling rule stands), continuous batching / PagedAttention comparisons, and byte-exact OOM claims (the gpu_memory_utilization gap documented in fleet/README.md).

Full plan with risks and open design decisions: docs/discussions/implementation_plans/gpu_disaggregated_inference.md.

Start with docs/discussions/implementation_plans/gpu_disaggregated_inference.md and the M4 milestones. Read src/dialect/VxLowering.cpp, runtime/npu_dispatch.mm, tests/frontend/pass/rubin_disaggregated.vx, tests/backend/pass/ffi_tcp_server.vx, and tests/backend/pass/llama2_v2.vx. The stated M4b acceptance is matching single-machine tokens, compile-time admission for both roles, and logged predicted versus measured KV-handoff size and latency.