Slowdown in CUDA graph execution using cudaMemsetAsync (vs. a fill kernel)

cbcase: <p>We have a PyTorch-based inference system that heavily uses graph capture+replay. Looking at profiles, every time there is a cudaMemset, there are gaps in the execution that look like this:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/1/d12c28ee63753f52af5d9dbe092bd94b0521a1d8.png" data-download-href="https://discuss.pytorch.org/uploads/default/d12c28ee63753f52af5d9dbe092bd94b0521a1d8" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/d/1/d12c28ee63753f52af5d9dbe092bd94b0521a1d8.png" alt="image" data-base62-sha1="tQqp0QBiigD80hBKtbjHaQJvgdq" width="690" height="53" data-dominant-color="C7CEDA"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">806×62 5.91 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>In this case, the (4byte) cudaMemsetAsync takes 1.7us, but the gap between kernels is almost 13us. PyTorch itself makes very little use of memset (it seems to prefer fill kernels), but plenty of libraries use memset as part of various kernel prep (here, the TransformerEngine fp8 layernorm kernel).</p>
<p>I was able to reproduce this effect with a small microbenchmark. It looks roughly like this:</p>
<pre data-code-wrap="python"><code class="lang-python">z = torch.zeros(1, device="cuda")
x = torch.randn(512, 1024, device="cuda")
w = torch.randn(1024, 1024, device="cuda")

# [warmup, setup, etc], then:
# capture
with torch.cuda.graph(g):
  for _ in range(1000):
    torch.ops.memset_cuda.memset_cuda(z)  # Case 1: use memset
    # z.zero_()  # Case 2: use vectorized_elementwise_kernel
    torch.matmul(x, w.T)

# Measure graph time
torch.cuda.synchronize()
tic = time.time()
g.replay()
torch.cuda.synchronize()
toc = time.time()
print(f"{(toc - tic) * 1000.}ms")
</code></pre>
<p>On a single H100-SXM5, I get these measurements for the graph time:</p>
<pre><code class="lang-auto">using memset: 34.8ms
using zero_: 33.4ms
</code></pre>
<p>This is almost a 5% E2E slowdown due to memset vs. a fill kernel! Looking at the profile, I see similar gaps as above, though their relative size is smaller. Any ideas why this might be? Or any workarounds? In our own code we can prefer the fill kernels, but it is impossible to workaround (eg) cuBLAS’ use of memset. Tagging <a class="mention" href="/u/ptrblck">@ptrblck</a> for NV awareness.</p>
<p><em>Aside: the call to <code>memset_cuda</code> is a 5-line cpp extension I added that simply calls:</em></p>
<pre><code class="lang-auto">cudaMemsetAsync(x.data_ptr&lt;float&gt;(), 0, x.numel() * sizeof(float),
    at::cuda::getCurrentCUDAStream().stream());</code></pre>

ptrblck: <p>This is an interesting observation as <a href="https://discuss.pytorch.org/t/replacing-torch-zeros-internals-with-cudamemset-instead-of-fill-kernel/208272/3">this topic</a> claims the opposite at least for eager execution.</p>
<p>Could you post your full extension code to reproduce the issue?</p>

cbcase: <p>Thanks for taking a look <a class="mention" href="/u/ptrblck">@ptrblck</a>!</p>
<p>Re: eager mode differences, it’s true that the reported operation time for the set-to-zero is lower with memset. There’s a little variation, but staring at some <code>nsys nvprof --print-gpu-trace</code> output, it looks like for this 4-byte memset I see:</p>
<ul>
<li><code>cudaMemsetAsync</code>: 0.8us</li>
<li><code>vectorized_elementwise_kernel</code>: 1.1-1.3us</li>
</ul>
<p>But this few-hundred-nanosecond difference is swamped in graph execution by the gaps with memset. (Do CE operations not have fast-dependent-launch?)</p>
<p>I don’t see a way to attach files, so here are the three files to reproduce my original numbers:</p>
<p>1: <code>setup.py</code></p>
<pre data-code-wrap="python"><code class="lang-python">from setuptools import setup
from torch.utils import cpp_extension

setup(
    name="memset_cuda",
    ext_modules=[
        cpp_extension.CUDAExtension(
            "memset_cuda",
            ["memset_cuda.cu"],
        )
    ],
    cmdclass={"build_ext": cpp_extension.BuildExtension},
)
</code></pre>
<p>2: <code>memset_cuda.cu</code></p>
<pre data-code-wrap="cpp"><code class="lang-cpp">#include &lt;c10/cuda/CUDAStream.h&gt;
#include &lt;cuda.h&gt;
#include &lt;torch/extension.h&gt;

namespace memset_cuda {

void memset_cuda(torch::Tensor&amp; x) {
    AT_DISPATCH_FLOATING_TYPES_AND2(
        at::ScalarType::Half, at::ScalarType::BFloat16, x.scalar_type(), "memset_cuda", ([&amp;] {
            cudaMemsetAsync(x.data_ptr&lt;scalar_t&gt;(), 0, x.numel() * sizeof(scalar_t),
                            at::cuda::getCurrentCUDAStream().stream());
        }));
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {}

TORCH_LIBRARY(memset_cuda, m) { m.def("memset_cuda(Tensor x) -&gt; ()"); }

TORCH_LIBRARY_IMPL(memset_cuda, CUDA, m) { m.impl("memset_cuda", &amp;memset_cuda); }

}  // namespace memset_cuda
</code></pre>
<p>3: <code>test.py</code></p>
<pre data-code-wrap="python"><code class="lang-python">import time
import torch
import memset_cuda


def main():
    z = torch.zeros(1, device="cuda")
    x = torch.randn(512, 1024, device="cuda")
    w = torch.randn(1024, 1024, device="cuda")
    use_memset = True

    g = torch.cuda.CUDAGraph()
    s = torch.cuda.Stream()
    s.wait_stream(torch.cuda.current_stream())

    # Warmup
    with torch.cuda.stream(s):
        for _ in range(30):
            if use_memset:
                torch.ops.memset_cuda.memset_cuda(z)
            else:
                z.zero_()
            torch.matmul(x, w.T)
    torch.cuda.current_stream().wait_stream(s)

    # Capture
    with torch.cuda.graph(g):
        for _ in range(1000):
            if use_memset:
                torch.ops.memset_cuda.memset_cuda(z)
            else:
                z.zero_()
            torch.matmul(x, w.T)

    # Measure
    torch.cuda.synchronize()
    tic = time.time()
    g.replay()
    torch.cuda.synchronize()
    toc = time.time()
    print(f"{(toc - tic) * 1000.}ms")


if __name__ == "__main__":
    main()
</code></pre>
<p>Run / measure:</p>
<pre><code class="lang-auto"># Build the extension
python setup.py build
# Run with use_memset=True
# Path will be different if you have a different python version
PYTHONPATH=build/lib.linux-x86_64-cpython-311 python test.py
# Then modify to have use_memset=False and re-run
</code></pre>
<p>Fwiw, we are on CUDA 12.1.</p>

ptrblck: <p>Thank you Carl! Let me take a look at it and profile the usage.</p>

cbcase: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a> – gently pinging this one to see if you’ve learned anything – thanks!</p>

ptrblck: <p>Thanks for the ping!<br>
I’ve talked to <a class="mention" href="/u/eqy">@eqy</a> and we were speculating if the usage of custom extensions by itself is adding enough CPU overhead (we’ve seen it in the past), but would need to profile it properly and maybe even compare it to a “native” implementation.</p>

cbcase: <p>Thanks for taking a look. Fwiw, I don’t think this is PyTorch-specific (so maybe this is the wrong forum?). Here is a sample cuda file that does nothing other than use either a kernel or cudaMemsetAsync to zero-out a 4-byte flag:</p>
<pre><code class="lang-auto">#include &lt;cuda.h&gt;
#include &lt;stdio.h&gt;

#include &lt;cassert&gt;
#include &lt;chrono&gt;
#include &lt;iostream&gt;

#define CHECK_CUDA(val) check((val), #val, __FILE__, __LINE__)
void check(cudaError_t err, const char* const func, const char* const file, const int line) {
    if (err != cudaSuccess) {
        std::cerr &lt;&lt; "CUDA Runtime Error at: " &lt;&lt; file &lt;&lt; ":" &lt;&lt; line &lt;&lt; std::endl;
        std::cerr &lt;&lt; cudaGetErrorString(err) &lt;&lt; " " &lt;&lt; func &lt;&lt; std::endl;
        std::exit(EXIT_FAILURE);
    }
}

__global__ void zero_out(float* a) { *a = 0; }

int main(int argc, char** argv) {
    if (argc != 2) {
        std::cerr &lt;&lt; "Usage: " &lt;&lt; argv[0] &lt;&lt; " [kernel|memset]" &lt;&lt; std::endl;
        std::exit(EXIT_FAILURE);
    }
    bool use_kernel = false;
    if (std::string(argv[1]) == "kernel") {
        use_kernel = true;
    } else if (std::string(argv[1]) != "memset") {
        std::cerr &lt;&lt; "Usage: " &lt;&lt; argv[0] &lt;&lt; " [kernel|memset]" &lt;&lt; std::endl;
        std::exit(EXIT_FAILURE);
    }
    CHECK_CUDA(cudaSetDevice(0));
    float* d_z;
    CHECK_CUDA(cudaMalloc(&amp;d_z, sizeof *d_z));
    CHECK_CUDA(cudaMemset(d_z, 0, sizeof *d_z));

    cudaStream_t stream, stream2;
    cudaGraph_t graph;
    CHECK_CUDA(cudaStreamCreate(&amp;stream));
    CHECK_CUDA(cudaStreamCreate(&amp;stream2));
    CHECK_CUDA(cudaDeviceSynchronize());

    // 1) Capture
    CHECK_CUDA(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal));
    for (int i = 0; i &lt; 1000; ++i) {
        if (use_kernel) {
            zero_out&lt;&lt;&lt;1, 1, 0, stream&gt;&gt;&gt;(d_z);
        } else {
            CHECK_CUDA(cudaMemsetAsync(d_z, 0, sizeof *d_z, stream));
        }
    }
    CHECK_CUDA(cudaStreamEndCapture(stream, &amp;graph));
    // 2) Instantiate and warmup
    cudaGraphExec_t graphExec;
    CHECK_CUDA(cudaGraphInstantiate(&amp;graphExec, graph));
    for (int i = 0; i &lt; 5; ++i) {
        CHECK_CUDA(cudaGraphLaunch(graphExec, stream2));
    }
    // 3) Measure
    CHECK_CUDA(cudaDeviceSynchronize());
    auto tic = std::chrono::high_resolution_clock::now();
    CHECK_CUDA(cudaGraphLaunch(graphExec, stream2));
    CHECK_CUDA(cudaDeviceSynchronize());
    auto toc = std::chrono::high_resolution_clock::now();

    auto duration = std::chrono::duration_cast&lt;std::chrono::microseconds&gt;(toc - tic);
    std::cout &lt;&lt; "Execution time: " &lt;&lt; duration.count() &lt;&lt; "us" &lt;&lt; std::endl;
    return 0;
}
</code></pre>
<p>Supposing this file is <code>test.cu</code>, here are the results running on my own machine (with an H100-SXM5):</p>
<pre><code class="lang-auto">&gt; nvcc -O3 -arch=sm_90 -o test test.cu
&gt; ./test memset
Execution time: 1490us
&gt; ./test kernel
Execution time: 1045us
</code></pre>
<p>In other words, the average time per-memset is ~1us for the kernel and ~1.5us for memset. If you look at traces, though, you’ll see that an individual memset takes closer to 800ns, so half of the time is gaps in execution. This matches what I saw all the way at the beginning of the thread, though the gaps we see in real applications are much larger (can be as large as 10us) – certainly if they were never more than 1us, that would be great!</p>