vx-lang / vx-lang/Vx

Dispatch does not say which device to run on, which disaggregation requires

Open
#331 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

codegen enhancement runtime
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):

  1. I32Attr:$topology on LaunchOp, set from the spawn's topology in SpawnOpLowering where the launch is created.
  2. A topo=<id> entry in the NUL-separated payload blob in LaunchOpLowering. Passing it as a separate argument would change the ABI signature; the blob exists so facts can be added without one.
  3. 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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.