Dispatch does not say which device to run on, which disaggregation requires
Nobody has claimed this yet.
- Dominant language
- Rust
- Stars
- 14
- Forks
- 2
- Avg merge
- 12h 42m
- Merged PRs (30d)
- 61
Description
vx.kernel carries the topology the programmer wrote (vx.kernel @vx_npu_kernel_0(500) for spawn on(Topology::GPU), from I32Attr:$topology in include/VxDialect.td). vx.launch -- the op that becomes the call into a plugin -- does not:
let arguments = (ins FlatSymbolRefAttr:$callee, Variadic<AnyType>:$operands);
So vx_plugin_dispatch_async(payload, size, args, tags, num_args) has no way to learn which device the program asked for. Every dispatch reaches a plugin identically, and a CUDA backend runs everything on whatever device is current.
That is the whole of disaggregated inference. Prefill on Topology::GPU[0] and decode on Topology::GPU[1] is a program that says exactly which device each stage runs on, and the runtime today cannot tell the two apart. The allocation entry points already take a topology_id, so the asymmetry is only in dispatch.
Shape of the fix
Add the topology to vx.launch and carry it into the payload the way the kernel kind and operand roles already travel (#325):
I32Attr:$topologyonLaunchOp, set from the spawn's topology inSpawnOpLoweringwhere the launch is created.- A
topo=<id>entry in the NUL-separated payload blob inLaunchOpLowering. Passing it as a separate argument would change the ABI signature; the blob exists so facts can be added without one. cudaSetDevice()from it in runtime/cuda_dispatch.cpp, and the same treatment in whatever a second backend does. A plugin that ignores the entry behaves exactly as it does today, which is what makes this safe to land ahead of the backends that use it.
Indexed topologies (Topology::GPU[1]) need the index to survive too, not just the kind -- worth checking what the i32 encoding does with the index, since 500 is the bare GPU.
Why file it now rather than fix it now
Nothing single-GPU needs it, so it is not blocking M1. It blocks the multi-device milestone, and it is better identified while the surrounding code is being written than discovered on two rented GPUs.
Related: hiraditya/Vx.1#321 (the GPU campaign), hiraditya/Vx.1#325 (the payload blob and the facts it carries).
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with I32Attr:$topology in include/VxDialect.td and follow SpawnOpLowering into LaunchOpLowering, comparing the payload facts described in #325. Then inspect runtime/cuda_dispatch.cpp and the i32 encoding for indexed topologies. Done means the launch topology survives into the payload and can select the requested device without changing the plugin ABI.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, rust
- Domain
- backend, compilers
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 52/100