3 ms·

Does anyone have an intuition on why this offers significant gains over CUDA Graphs?. The CPU launch cost of a graph is tiny which implies most of the work has

by skavi 1y ago

Does anyone have an intuition on why this offers significant gains over CUDA Graphs?. The CPU launch cost of a graph is tiny which implies most of the work has been offloaded to the GPU's own scheduler. I'd expect that some I/O marshalling at kernel boundaries could be avoided with megakernels. Maybe some loop fusion? Are there any more interesting optimizations they enable?

refulgentis 1y ago
You've hit the nail on the head. The CPU launch cost of a pre-compiled CUDA graph is tiny. CUDA Graphs are a huge step up from manually launching kernels, but they still treat kernels as monolithic, black-box operations. A megakernel erases the boundaries between those operations. With CUDA Graphs, as in the example in the article, if you have Matmul -> AllReduce, the AllReduce kernel cannot start until the entire Matmul kernel has finished. The dependency is at the kernel level. With a megakernel, they break these ops into fine-grained "tasks" scheduled across SMs. An AllReduce task that needs data from the first slice of the Matmul can begin as soon as that slice is computed by a few SMs, while other SMs are still working on the rest of the Matmul. This fine-grained software pipelining and compute/communication overlap is simply not possible when the dependency unit is the entire kernel.
- skavi 1y ago
  Ah, that makes a lot of sense. Is this fine grained task scheduling related to CUDA Dynamic Parallelism at all? If not would you have a pointer on where to look? I suppose I could look through the code of this project, but I’d hate to have detangle that from the compiler infrastructure.
  - touisteur 1y ago
    Think more of it as 'tasking by hand' where you have one kernel driving the 18k+ cores and you manually (or using device libraries) fine-grained synchronize them, handle memory traffic asynchoronously and pipeline as much as you can. You might have a look at cooperative groups, also things cuda::pipeline in libcudacxx to handle asynchronous and pipelined memory traffic, and also most of block/warp CUB primitives, and move on up to cuFFTDx, cuBLASDx and now cuSolverDx as the starting toolbox for your fused kernel journey.
saagarjha 1y ago
> The CPU launch cost of a graph is tiny Absolutely not; it’s comparable to the launch overhead of a kernel.
- skavi 1y ago
  Fair enough. I should have clarified that “approximately the cost of a single kernel launch” is pretty much what I meant by “tiny”. What I was getting at was that a “megakernel” and a captured graph should have similar launch costs.
  - zhihaojia 1y ago
    You are right that CUDA graph can help reduce launch overhead but does not support overlapping computation/communication across layers, since data dependencies are described at the kernel level.
  - touisteur 1y ago
    It's not so much kernel overhead than memory traffic between global memory and L2 cache / shared memory that you (or at least I) target with fused kernel approach. Kernel launch overhead can be drastically reduced with cuda-graph indeed. I'm not sure it applies so well in LLMs though (should read the paper...).

by skavi 1y ago

refulgentis 1y ago

You've hit the nail on the head. The CPU launch cost of a pre-compiled CUDA graph is tiny. CUDA Graphs are a huge step up from manually launching kernels, but they still treat kernels as monolithic, black-box operations. A megakernel erases the boundaries between those operations. With CUDA Graphs, as in the example in the article, if you have Matmul -> AllReduce, the AllReduce kernel cannot start until the entire Matmul kernel has finished. The dependency is at the kernel level. With a megakernel, they break these ops into fine-grained "tasks" scheduled across SMs. An AllReduce task that needs data from the first slice of the Matmul can begin as soon as that slice is computed by a few SMs, while other SMs are still working on the rest of the Matmul. This fine-grained software pipelining and compute/communication overlap is simply not possible when the dependency unit is the entire kernel.

skavi 1y ago
Ah, that makes a lot of sense. Is this fine grained task scheduling related to CUDA Dynamic Parallelism at all? If not would you have a pointer on where to look? I suppose I could look through the code of this project, but I’d hate to have detangle that from the compiler infrastructure.
- touisteur 1y ago
  Think more of it as 'tasking by hand' where you have one kernel driving the 18k+ cores and you manually (or using device libraries) fine-grained synchronize them, handle memory traffic asynchoronously and pipeline as much as you can. You might have a look at cooperative groups, also things cuda::pipeline in libcudacxx to handle asynchronous and pipelined memory traffic, and also most of block/warp CUB primitives, and move on up to cuFFTDx, cuBLASDx and now cuSolverDx as the starting toolbox for your fused kernel journey.

saagarjha 1y ago

> The CPU launch cost of a graph is tiny Absolutely not; it’s comparable to the launch overhead of a kernel.

skavi 1y ago
Fair enough. I should have clarified that “approximately the cost of a single kernel launch” is pretty much what I meant by “tiny”. What I was getting at was that a “megakernel” and a captured graph should have similar launch costs.
- zhihaojia 1y ago
  You are right that CUDA graph can help reduce launch overhead but does not support overlapping computation/communication across layers, since data dependencies are described at the kernel level.
- touisteur 1y ago
  It's not so much kernel overhead than memory traffic between global memory and L2 cache / shared memory that you (or at least I) target with fused kernel approach. Kernel launch overhead can be drastically reduced with cuda-graph indeed. I'm not sure it applies so well in LLMs though (should read the paper...).