Concepts

Tracing Python to CUDA: A Practical Guide to Full-Stack GPU Call Attribution

Sofia Moran Sofia Moran
9 min read
Tracing Python to CUDA: A Practical Guide to Full-Stack GPU Call Attribution

Most profiling tools give you either a Python call stack or a GPU kernel timeline. Neither one alone is useful for finding bottlenecks in an AI inference service. The Python view shows you which Python functions ran the longest, but the expensive work happens below the Python layer. The GPU timeline shows you which kernels ran and for how long, but without the Python context you cannot trace a slow kernel back to the user code responsible for it.

Stitching these two views together is the core challenge of whole-stack GPU profiling. This post explains how zymtrace approaches it using eBPF, what the call path looks like from a Python operator call down to the CUDA kernel, and how to read the combined output when something is slow.

The Dispatch Chain From Python to Device

When you call torch.matmul(a, b) in Python, the execution path crosses several distinct layers before any GPU work starts. Understanding these layers helps you interpret where in the chain a stall or unexpected cost appears.

The first layer is Python itself. The function call goes through CPython's object system, dispatches through PyTorch's Python bindings, and reaches the ATen C++ operator dispatch table. This costs between 1 and 5 microseconds for a typical operator call, depending on the operation and whether autograd is enabled.

The second layer is ATen dispatch. The dispatcher selects the right backend implementation based on tensor device type (CPU, CUDA, XLA, etc.) and data type. For CUDA tensors it selects the CUDA kernel implementation, which lives in the CUDA extension shared library. This is also where autograd records operations if you are in a torch.no_grad() context or not.

The third layer is the CUDA runtime API. The selected kernel implementation calls cudaLaunchKernel (or the newer cuLaunchKernelEx) to queue the kernel on a CUDA stream. This is an asynchronous call by default: it returns immediately on the CPU side while the GPU executes the kernel independently.

The fourth layer is the driver and device. The CUDA driver transfers the kernel launch to the hardware command buffer, the GPU's hardware scheduler issues it to available SMs, and the kernel executes. The CPU and GPU are now running in parallel until there is an explicit synchronization point.

Where eBPF Intercepts the Chain

eBPF programs run inside the Linux kernel with verified safety guarantees. They can attach to kernel tracepoints, kprobes, and uprobes. For userspace library functions like the CUDA runtime, eBPF uprobes attach directly to the shared library symbols without modifying the library or the application binary.

zymtrace's collection agent attaches uprobes to key CUDA runtime symbols: cudaLaunchKernel, cudaMemcpy, cudaMemcpyAsync, cudaStreamSynchronize, cudaDeviceSynchronize, and several others. When any of these functions are called, the eBPF program runs in kernel context, reads the current CPU stack, and records the event with a timestamp and the thread's call stack.

The CPU stack at the moment of cudaLaunchKernel includes the full path from the Python frame (via the CPython interpreter's frame objects, which are readable from eBPF user-stack walkers) through the C extension, through ATen, to the CUDA driver call. That stack is the bridge: it links the Python-level context to the GPU-level kernel event.

On the GPU side, zymtrace correlates the kernel launch event with the kernel's actual execution time using CUPTI (CUDA Profiling Tools Interface) event callbacks where available, or CUDA event timestamps inserted around kernel launches. The result is a record that says: Python function X called PyTorch operator Y which dispatched CUDA kernel Z, and that kernel ran for W microseconds.

What Shows Up in the Flamegraph

The combined call path appears in the flamegraph as a continuous stack from the Python root frame down to the CUDA kernel symbol. There is no gap where the profiler loses track of execution. The transition from Python to C extension to ATen to CUDA driver to kernel is visible as adjacent layers in the same stack.

This matters most when you have multiple code paths that call the same kernel. Consider a transformer inference service where aten::mm (matrix multiply) appears in three places: the attention projection, the feed-forward projection, and an embedding lookup path. In a flat profiler view, all three aggregate into a single aten::mm entry. In the whole-stack flamegraph, each occurrence appears under its own call path. If the attention projection's matrix multiply is slower than expected but the others are fine, the flamegraph shows you exactly which path to investigate.

Synchronization Stalls and How They Look

One of the most useful things the combined view reveals is implicit synchronization. Certain PyTorch operations force the CPU to wait for GPU execution to complete before continuing. The most common examples are calls to .item() on a GPU tensor (which requires the value to be copied to host), .numpy(), printing a tensor, or calling Python's bool() on a GPU tensor.

When these happen inside an inference loop, they create a fence: the CPU cannot proceed with the next batch dispatch until the GPU finishes and returns the value. The effect is that the CPU and GPU stop running in parallel and you serialize what was previously overlapping work.

In the zymtrace flamegraph, these appear as wide frames attributed to cudaDeviceSynchronize or cudaMemcpy (device to host direction), sitting below the Python call that triggered them. The width represents the GPU idle-on-CPU-wait time plus the transfer time. The Python frame above it tells you exactly which line caused the synchronization.

Practical Tracing: An Example Scenario

Consider an inference endpoint that processes variable-length inputs. The service receives a batch of requests, tokenizes them, pads to the maximum sequence length in the batch, runs the model forward pass, and then strips padding tokens from the output before returning results.

The output stripping step determines the actual sequence lengths of each output sequence by calling something like output_lengths = (output_ids != pad_token_id).sum(dim=1).tolist(). The .tolist() call on a GPU tensor triggers a synchronization and a device-to-host copy. In the flamegraph, this shows as a cudaMemcpy frame of around 20 to 40 microseconds for a batch of 32, sitting under the post-processing Python function.

Without the Python context, you would see a device-to-host copy in the GPU timeline but have no idea what triggered it. With the combined stack, you see exactly which line of Python is responsible, and you can evaluate alternatives: could the output lengths be computed differently to avoid the synchronization, or could the stripping be done in a GPU kernel that does not require copying lengths to CPU?

Limitations of the eBPF Approach

The eBPF user-stack walker reads Python frames by traversing CPython's frame pointer chain. This works reliably for CPython 3.9+ with debug symbols or with the frame pointer preserved. Code compiled with -fomit-frame-pointer loses the frame chain for C extension frames, which means those layers may appear as unknown frames in the stack. zymtrace's agent handles this by using DWARF debug information for symbol resolution where available.

CUDA kernels launched on non-default streams with significant queue depth can show timing offsets between when the launch was recorded on the CPU and when the kernel actually executed on the GPU. We track this with CUDA event timestamps, but the correlation adds latency to the trace data rather than impacting the application itself.

The combined view is most reliable on CUDA 11.8 and later with Python 3.9+. Older CUDA versions have limited CUPTI support for concurrent kernel timing. The core CPU-side stack attribution works on older versions, but GPU execution time precision is reduced.

Getting Started

The practical starting point is to run the zymtrace agent against a running inference process, collect 60 seconds of call-path data under production-realistic load, and look at the flamegraph. The first question to answer is whether memory copies appear prominently in the combined stacks. If they do, the Python context above them will tell you which code path is responsible. That usually surfaces 2 or 3 concrete changes that reduce synchronization, improve batching, or eliminate unnecessary device-host round trips. Those changes are typically more impactful than any kernel-level optimization because they address structural inefficiencies in how the Python orchestration layer interacts with the GPU.

try it yourself

See your GPU call paths in a live flamegraph

Free tier covers 1 GPU node and 90 days of call-path history. No credit card required.