a tale of two runtimes

problem: the GPU wasn't saturated

The current execution model used one CUDA stream and one DataFusion partition (analogous to a Tokio task), meaning only one host thread was feeding the GPU. The GPU could execute only one kernel or host-to-device copy at a time (with the sole exception of the cuDF read_parquet kernel, which uses streams internally).

The reason we had this simple model was that it was easier to reason about correctness and lifetimes. The safeguards that Rust has around shared ownership and lifetimes don't really apply to data on the GPU. For example, the borrow checker and type system cannot detect potential cross-stream dependencies or ensure that a necessary stream sync happens before reading data on the GPU.

However, as the project matured, we started looking for the next performance gain, namely in the form of concurrency and parallelism. We wanted to:

1. efficiently schedule cuDF/GPU operations from the host runtime; and
2. concurrently run operations on the GPU when it has available resources.

Doing so requires migrating to a multi-task, multi-CUDA-stream model, which I'll dive into below.

first: stop blocking host threads

Every DataFusion operator produces a futures::Stream that polls its input, performs the operation, and yields a RecordBatch. For example, a DataFusion AggregateExec implements something like this:

fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Result<Batch>>> {
    let batch = ready!(self.input.poll_next_unpin(cx));
    let output = aggregate_on_cpu(batch)?;
    Poll::Ready(Some(Ok(output)))
}

For our GPU operators, we implemented a simple translation of the above. A CuDFAggregateExec is almost the same:

fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Result<Batch>>> {
    let batch = ready!(self.input.poll_next_unpin(cx));
    let output = aggregate_on_gpu(batch)?; // new: launches a kernel
    Poll::Ready(Some(Ok(output)))
}

The problem is that aggregate_on_gpu calls into cuDF through FFI. cuDF enqueues kernels asynchronously, but its APIs are synchronous from the caller's perspective and may sync streams internally for correctness. This means poll_next cannot yield while executing cuDF code, resulting in long polls that block the calling Tokio executor thread.

I inspected a query using the Tokio Console and saw the issue firsthand. In TPC-H Q1, the root task executing the query was marked busy for 9 of 10 seconds, with an average poll around 27 ms.

         state   total   busy   sched   idle   polls   avg poll
before    ▶       10s     9s    628ms   36µs     335     ~27ms

This meant that for 9 seconds, the executor could not schedule another future on that worker thread until the call returned.

In PR #89, I implemented a simple fix. Because cuDF calls can block their calling thread, I moved the GPU aggregate kernel work to Tokio's blocking pool. spawn_blocking returns a future, allowing poll_next to return Pending while the work continues on another thread.

match &mut self.state {
    State::ReadingInput => {
        // Move the blocking GPU call to another thread.
        self.state = State::Aggregating(
            tokio::task::spawn_blocking(move || aggregate_on_gpu(batch)),
        );
        continue;
    }
    // Later calls poll the GPU task's JoinHandle.
    State::Aggregating(task) => match Pin::new(task).poll(cx) {
        Poll::Pending => Poll::Pending,
        Poll::Ready(result) => finish(result),
    },
}

In the Tokio Console, our poll duration improved to 1 ms and the task was busy for far less time.

         state   total   busy   sched   idle   polls   avg poll
after     ⏸       17s     2s     15s     3ms    1993      ~1ms

second: add tasks and CUDA streams

In CPU land, DataFusion parallelizes a plan in two directions. Vertical parallelism splits the plan into pipelines at operators such as RepartitionExec. When first polled, RepartitionExec spawns producer tasks to drive the pipeline below it, allowing the pipelines above and below the boundary to run concurrently. Horizontal parallelism runs multiple partitions of each pipeline, with each partition processing a different slice of the input.

For example, this plan:

SortExec
  AggregateExec
    RepartitionExec
      ProjectionExec
        FilterExec
          DataSourceExec

can run as four concurrent tasks: two pipelines/segments, s0 and s1, each with two partitions, p0 and p1.

s0,p0: SortExec <- AggregateExec <- RepartitionExec
s0,p1: SortExec <- AggregateExec <- RepartitionExec

s1,p0: ProjectionExec <- FilterExec <- DataSourceExec
s1,p1: ProjectionExec <- FilterExec <- DataSourceExec

PR #92 applies a similar idea to GPU aggregates. It adds an optimizer rule that recognizes an aggregate over a Parquet scan with pipelineable operators in between (note that aggregates are pipeline breakers because they must consume their entire input before emitting a batch). The rule splits these pipelines into parallel partitions and assigns a non-blocking CUDA stream to each.

before

Aggregate · single       [1]
  Projection             [1]
    Filter               [1]
      Coalesce       [8 → 1] // everything above this runs in one stream
        Parquet scan     [8]

after

Aggregate · final        [1]
  Coalesce           [8 → 1] // everything below is parallelized on 8 streams
    Aggregate · partial  [8] // per-partition aggregation to reduce the data size
      Projection         [8]
        Filter           [8]
          Parquet scan   [8]

The GPU-backed record batches reference the stream they are on, so filters, projections, and partial aggregates enqueue kernels on the same stream, avoiding any cross-stream dependencies. Only the final aggregate synchronizes all the streams and merges their partial results. The final aggregate would have had to wait for its entire input anyway because it's a natural pipeline breaker.

results

TPC-H SF100 on a g7.4xlarge with an NVIDIA RTX PRO 4500 Blackwell Server Edition (32 GB, 165 W). These are mean end-to-end query times from the PR. They compare the original plan with the complete multi-partition, multi-stream rewrite.

• 2.53× — Q15 speedup at 8 streams
• 2.07× — Q1 speedup at 8 streams
• 1.77× — Q20 speedup at 2 streams

the profile tells the same story

An NVIDIA Nsight Systems profile of Q1 tells the same story. Extra streams turn empty GPU capacity into useful overlap: the runtime drops by half while all three utilization measures rise sharply.

• 8% → 47% — SM warp occupancy
• 20% → 85% — SM activity
• 37% → 95% — GPU activity

before · one stream

after · eight streams