kvcache-ai / kvcache-ai/Mooncake
[PG] Clarify enqueue-stream ownership across CUDA task slots in the legacy worker
- Dominant language
- C++
- Stars
- 6.6k
- Forks
- 1.2k
- Avg merge
- 3d 5h
- Merged PRs (30d)
- 312
Description
## Summary
The legacy CUDA worker rotates CUDA work through task slots 2 and 3, but both
slots use a single worker-level `enqueue_stream_`.
`enqueueTaskKernel` is not a short enqueue-only kernel: after publishing the
task, it remains resident until the CPU worker completes the task and clears
`task.active`.
This appears to make enqueue-stream ownership inconsistent with the logical
two-slot task model. I would like to clarify whether serialization through one
shared stream is intentional worker backpressure or whether each CUDA task slot
is intended to own an enqueue stream.
This is not a claim that general `async_op=True` behavior is broken.
## Current legacy-worker behavior
On current main (`b8f294d4`):
- `MooncakeWorker` declares one `enqueue_stream_`:
https://github.com/kvcache-ai/Mooncake/blob/b8f294d4aa6cda6689fa59cd6c37e69e0e27be6e/mooncake-pg/include/mooncake_worker.cuh#L133-L149
- `putTaskCuda` maps `cudaTaskCount` to task IDs 2 and 3, but queues both
slots' copies and `enqueueTaskKernel` on that same stream:
https://github.com/kvcache-ai/Mooncake/blob/b8f294d4aa6cda6689fa59cd6c37e69e0e27be6e/mooncake-pg/src/mooncake_worker_host.cpp#L245-L318
- The enqueue kernel spins while `task.active` remains set:
https://github.com/kvcache-ai/Mooncake/blob/b8f294d4aa6cda6689fa59cd6c37e69e0e27be6e/mooncake-pg/src/mooncake_worker.cu#L18-L42
- The worker clears `task.active` only after task processing completes:
https://github.com/kvcache-ai/Mooncake/blob/b8f294d4aa6cda6689fa59cd6c37e69e0e27be6e/mooncake-pg/src/mooncake_worker_thread.cpp#L414-L432
In abbreviated form:
```text
cudaTaskCount -> slot 2 / slot 3
|
v
shared enqueue_stream_
|
v
enqueueTaskKernel resident for task lifetime
```
The source also says "one enqueue stream per task slot," although the
implementation has one stream for both CUDA slots.
## Why this may matter
With a resident kernel on the shared stream:
```text
issue stream A -> slot 2 -> shared stream -> blocking kernel
issue stream B -> slot 3 -> shared stream -> queued behind slot 2
```
Slot 3 is logically distinct, but its enqueue kernel cannot publish its task
until slot 2's kernel releases the shared stream. This can turn the submission
gate for slot 3 into an indirect wait for slot 2's completion.
That may be intentional capacity backpressure, but it does not appear to match
the "one enqueue stream per task slot" comment.
## Evidence
I ran a bounded two-rank architecture experiment on 2x NVIDIA A40, driver
580.126.09, CUDA 13.0, and PyTorch 2.9.1+cu130.
Rank 1 entered 500 ms late. Rank 0 issued one small async all-reduce on each of
two independent CUDA streams. All results were correct.
| Revision | stream A issue median | stream B issue median |
|---|---:|---:|
| `3b5a594` before #3609 | 0.297 ms | 501.172 ms |
| local per-slot-stream prototype | 0.186 ms | 0.126 ms |
This was a bounded CUDA stream/task-lifetime experiment, not native
RDMA/InfiniBand or production-transport validation. The experiment used forced
TCP transport.
## Prototype observation and limitations
A local prototype gave each CUDA task slot its own enqueue stream and routed
the slot's copies, enqueue kernel, and events through that stream. It removed
the specific serialization between the two independent issue streams above.
The prototype has not been pushed.
The prototype does not solve broader completion/lifetime behavior:
- same-issue-stream consecutive collectives can still serialize;
- a single collective using more than two chunks can still incur
remote-progress-dependent submission backpressure;
- the prototype is not a complete async semantics fix.
## New device collective framework
The newer device collective framework in #3483 / #3484 is opt-in, and the
legacy implementation remains the default worker path during its rollout. This
question concerns enqueue-stream ownership in that legacy/default path.
## Design question
Is one enqueue stream per CUDA task slot the intended legacy worker model, or
is serialization through one shared enqueue stream intentional?
If per-slot ownership is intended, would a narrowly scoped per-slot-stream PR
be useful before addressing broader Work-completion and task-lifetime
semantics?
## Relationship to #3765
This investigation originated while looking at #3765, but this issue does not
depend on #3765's withdrawn async-regression claim.
Contributor guide
Research direction
Start with MooncakeWorker in mooncake-pg/include/mooncake_worker.cuh, then trace putTaskCuda in mooncake_worker_host.cpp, enqueueTaskKernel in mooncake_worker.cu, and task.active clearing in mooncake_worker_thread.cpp. Reproduce the bounded two-rank experiment if the CUDA environment is available; done means documenting the intended ownership model and validating any narrowly scoped change against the stated stream-serialization behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- distributed-systems, performance
- Issue type
- Refactor
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Needs clarification
- Newbie friendliness
- 35/100