kvcache-ai / kvcache-ai/Mooncake

[PG] Clarify enqueue-stream ownership across CUDA task slots in the legacy worker

Open
#3,775 3 comments 1 reaction 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.