RFC: Async overlap of CPU and GPU compute during dynamic inference step

[tdene]

This RFC the following plan for how to best optimize the dynamic inference step function.

There are several interconnected issues at play:

• Dynamic sampling code is currently very unoptimized. There is a PR draft that reimplements it.
• async_generate_output_tokens_dynamic_batch mixes CPU and GPU operations indiscriminately.
• async_generate_output_tokens_dynamic_batch may be declared async, but it has no good way of yielding the event loop. A lot of CPU time is wasted waiting for the GPU, and can be reclaimed.

The ideal solution appears to be:

• Fix dynamic sampling code.
• Clearly separate CPU and GPU operations.
• Provide a place to yield the event loop.

The PR series suggested by this RFC are:

1. Break async_generate_output_tokens_dynamic_batch apart into multiple sub-methods, which are clearly labeled as "CPU compute" vs "GPU compute".
   • Achieved by Clean up dynamic inference step #1992.
2. Implement barebones unoptimized dynamic sampling code.
   • Achieved by Allow mixed-batch sampling in dynamic inference #1927.
3. Tensorize the dynamic sampling bookkeeping.
   • Achieved by Tensorize dynamic inference mixed sampling #2105.
4. Reorganize the step function to allow for async step
   • Achieved by Break apart dynamic inference step into 2 methods #2192. No new functionality, just moving code around.
5. Reorder the sub-methods from point 1) so that CPU/GPU compute forms separate continuous blocks of code, and yield the event loop after the CPU compute via torch polling.
   • Due to all the prep work, this will be a tiny, extremely readable, PR.
   • Achieved by Make the dynamic engine step async #2193.
6. Optimize dynamic sampling code via graphed FlashInfer sampling.
   • A draft has been written by @kanz-nv; @tdene will finish it.
7. Refactor dynamic logprobs computation to follow the same style as the new sampling code.
   • A draft has been written by @tdene.
8. Reconcile with main's implementation of top_n_logprobs.
9. Wait for a torch update, or brainstorm a way to yield the event loop without polling in the current version of pytorch.
   • Maybe by sampling on a single rank, instead of the current sampling on every rank?
   • Will discuss further in comments.

[tdene]

My starting thoughts on points 6 and 7.

For point 6, it is trivial to yield the event loop in pytorch in the following manner:

test_evt = torch.cuda.Event()

do_GPU_work()

test_evt.record()
while not test_evt.query():
    await asyncio.sleep(10 ** -5)

This is a very bad solution. However, to the extent of my current knowledge, this is the only way to achieve this in non-distributed pytorch. I consider this a flaw of pytorch.
I would like to do some community brainstorming and see if there is a way to yield the event loop in a non-polling manner here.

For point 7, it is possible to yield the event loop in a non-polling manner in distributed pytorch. This type of method has seen a lot of use in my code:

def work_to_awaitable(work: dist.Work, loop: asyncio.AbstractEventLoop) -> Awaitable:
    """Convert a dist work object to an awaitable.

    torch.dist.work has a get_future method, but that torch.futures.Future unblocks as soon
    as the GPU operation has been enqueued, not when the GPU operation completes.

    torch.dist.work has a get_future_result method, but that only returns the status of the
    result, not the result itself. However, it unblocks only when the result is available.
    But get_future_result runs in a thread and asyncio.future.set_result() is not thread-safe.

    The options are the following:
     - Poll the Work's is_completed method
     - Use the Work's wait method to block
     - Use a Future's wait method which is the same as the Work's wait
     - Use a Future's done-based methods, which unblock too early
     - Use the done-based methods of the result Future, which needs loop.call_soon_threadsafe
    """
    ret_fut = loop.create_future()
    # Handle case where pytorch version is too old for get_future_result.
    try:
        torch_fut = work.get_future_result()
    except AttributeError:
        torch_fut = work.get_future()

    torch_fut.add_done_callback(
        lambda _: loop.call_soon_threadsafe(lambda: ret_fut.set_result(None))
    )
    return ret_fut

I consider this a good solution.

Unfortunately, this requires a torch.distributed.Work handle, which is only generated by distributed work. One way to take advantage of this would be by doing dummy distributed work, if we can find a dummy workload with minimal overhead.

Another way to take advantage of this would be by performing actual distributed work: we only sample on 1 rank per group, and broadcast the sampled tokens. That broadcast allows us to yield the event loop cleanly.

[tdene]

As previously discussed, performing sampling on a single rank has the following pros and cons.

Pros:

• Less total compute is performed.
  • However, since the ranks need to synchronize anyway, the amount of total compute does not practically matter. Only the critical path matters, and that remains unchanged.
• There is no need to worry about synchronizing sampling RNG or related factors across different ranks.
• It greatly simplifies the process of doing advanced sampling, where the sampling function can evolve beyond some simple CUDA ops and become a complicated algorithm (e.g. speculative decoding).
• This also simplifies integration with advanced coordinator features that we may eventually have demand for.

Cons:

• The critical path lengthens due to the additional broadcast.
  • However, to my understanding, yielding the event loop as desired by this RFC can make this issue inconsequential. We do place an extra GPU op on the critical path, but we're already overlapping GPU and CPU computation, thus the GPU op's latency ideally gets hidden behind our CPU ops. Would need to test and prove.