Skip to content

[Core Refactor] Task-ring publication lacks a valid GPU-to-CPU synchronization edge #99

Description

@XbzOnGit

Summary

The native task ring does not establish a valid CUDA happens-before relationship when publishing GPU-written entries to the CPU.

Its publication sequence is:

write payload and task metadata
-> __threadfence()
-> volatile ready_seq store

The CPU polls ready_seq with an acquire load. However, __threadfence() has device scope and the volatile store is not a system-scope release. The CPU acquire therefore has no matching release operation. Observing ready_seq does not formally guarantee that the associated metadata and payload are ready; PCIe ordering cannot replace the required CUDA memory-model edge.

Why not publish every entry with release/acquire?

Replacing every publication with a GPU system-scope release and CPU acquire would provide the ordering, but a direct RTX 4090 prototype measured:

current publication:       2.393 us
system release/acquire:    3.295 us
difference:               +0.902 us (+37.7%)

A single-counter D2D + publish benchmark found approximately:

Payload Overhead
4-64 KiB ~70%
1 MiB ~51%
16 MiB ~11%
64 MiB ~0.5%

In Qwen3-0.6B vLLM workloads, per-entry system publication increased complete generation time by approximately 0.6-0.8% with hidden-states (28 publications per forward) and 6.7-9.8% with vllm-full (341 publications per forward).

Proposed direction

All currently supported hooks have CPU-known payload sizes and execution order. For this path, keep metadata in a CPU FIFO and use:

  • the existing GPU payload ring;
  • a GPU-local produced-entry counter; and
  • one CPU/GPU-shared monotonic published_head.

Publish published_head with a system-scope release every fixed K hooks and at every forward boundary. The CPU observes it with an acquire load and only then submits the corresponding D2H to the drain stream.

This replaces per-entry managed publication and polling with one release per batch and one polled cache line. It does not reduce payload D2D or D2H traffic, and it delays CPU visibility until a batch boundary. Publishing at every forward boundary ensures that a partial batch cannot remain hidden.

The required ordering chain is:

Stream-A payload writes and writer join
-> GPU system release
-> CPU acquire
-> CPU submits Stream-B D2H
-> Stream-B reads the payload

The same CPU thread must acquire the published value before enqueueing the D2H; the mechanism cannot order a copy that was already enqueued. The counter must be a naturally aligned system-scope atomic in managed memory on a device with concurrentManagedAccess == 1. Publication must join every payload writer, and ring storage cannot be reused until its D2H completes.

GPU-resolved payload sizes

Hooks whose identity and order are CPU-known but whose actual payload size is resolved on the GPU--for example, an EP per-expert output whose row count depends on GPU routing--can use an optional managed ring with one uint64_t size slot per outstanding dynamic entry. The CPU FIFO carries a reservation upper bound and the corresponding slot identity.

In eager execution, a hook that is dynamic only relative to pre-forward planning can instead enqueue the exact size just in time when it is available to the CPU at hook invocation. The managed size ring is only needed when the size remains GPU-resolved at publication time.

The GPU writes the actual size before the next batched release of published_head; the same release/acquire edge then publishes both payload and size, so the size slot needs no separate synchronization. Capacity uses the reservation upper bound, requires actual_size <= reserved_size, and prevents slot reuse until CPU consumption. Dynamic hook identity, count, and execution order remain out of scope.

References

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions