CUDA Streams and Kernel Concurrency: The Overlap Engine Behind ZeRO-3 and FSDP

1. What a CUDA Stream Actually Is

Strip away the marketing language and a CUDA stream is disarmingly simple: a FIFO queue of GPU operations that execute in submission order, asynchronous to the CPU host. That is the complete definition. Every cudaMemcpyAsync, every kernel launch, every NCCL collective you issue against a given stream is serialized within that stream — but the host thread that submitted those commands returns immediately and moves on.

The GPU has a hardware work distributor — the GigaThread Engine on NVIDIA architectures — that dequeues commands from multiple streams and dispatches them to the SM pool concurrently, subject to resource constraints. This is what people mean when they say “streams enable concurrency”: the serialization guarantee is per-stream, not global.

Default stream vs. non-default streams is where things get subtle. The default stream (stream 0) in CUDA’s legacy mode is synchronizing: any operation on the default stream acts as a barrier against all other streams on the same device. This behaviour was chosen for correctness in the early days of CUDA when multi-stream code was rare; it is now almost universally a footgun. Modern code should use either non-default streams created with cudaStreamCreate or the per-thread default stream (--default-stream per-thread compile flag or cudaStreamCreateWithFlags(..., cudaStreamNonBlocking)), which does not synchronize against other streams. PyTorch’s CUDA backend creates a pool of non-blocking streams for exactly this reason, and torch.cuda.stream(s) context managers re-schedule CUDA work into s for the duration of the block.

What concurrency looks like physically: two non-default streams issuing independent kernels share the same 132 SMs on an H100. The SM scheduler time-slices resident warps across streams. If kernel A from stream 0 is bandwidth-bound and leaves compute units idle, kernel B from stream 1 can fill those slots. This is the mechanism — not magic, not a second GPU — behind every “overlap” claim in distributed training.

2. Why Naive Distributed Training Is Sequential

Before overlap, the standard data-parallel training loop looks like this in wall-clock time:

[Forward pass, all layers]
        ↓
[AllReduce gradients across all ranks]
        ↓
[Optimizer step]
        ↓  (next iteration)

In plain DDP (DistributedDataParallel with find_unused_parameters=False), PyTorch fires a single AllReduce after the entire backward pass completes. This is trivially sequential: the GPU sits idle during the AllReduce, and NCCL sits idle during the backward. The machine has provisioned two resources — compute SMs and InfiniBand NICs — but uses only one at a time.

For small models on fast interconnects this is tolerable. For a 70B model on InfiniBand NDR at 100 GB/s effective bandwidth per node, an AllReduce over 140 GB of BF16 gradients costs roughly 2.6 seconds while the compute step on 128 H100s runs in ~100–200 ms. The ratio is 13–26×. The cluster is compute-idle for 92–96% of each step. This is the canonical failure mode.

ZeRO-3 (Rajbhandari et al., arXiv:1910.02054) and PyTorch FSDP (Zhao et al., arXiv:2304.11277) both exist to escape this regime. But escaping it requires understanding streams.

3. The Overlap Trick: Anatomy of a Pipelined Step

The key insight is that gradient communication and forward computation for the next mini-batch (or the next layer in the backward pass) are data-independent. The gradient of layer N does not depend on anything produced by layer N+1’s forward. You can therefore issue both concurrently.

The timeline, in words:

Time →
Stream A (compute): | Layer N+1 fwd | Layer N+2 fwd | Layer N+3 fwd | ... |
Stream B (comm.):       | AllReduce(grad_N) |  AllReduce(grad_N-1) | ... |
                         └── overlapped ──┘

Concretely, in DDP with gradient bucketing (the default):

The backward hook for layer N fires when grad_N is complete.
PyTorch enqueues ncclAllReduce(grad_bucket_N, ..., comm_stream) on the communication stream.
The backward continues computing grad_N+1 on the default compute stream.
The SM scheduler runs both simultaneously if SMs are available.

In FSDP (arXiv:2304.11277, §3), the pattern is more aggressive. Before the forward pass, FSDP issues AllGather on the next layer’s parameter shard on comm_stream while the current layer’s forward computation runs on compute_stream. The prefetch lookahead (forward_prefetch=True, backward_prefetch=BackwardPrefetch.BACKWARD_PRE) ensures that by the time computation arrives at layer k, the full weights are already materialized in GPU memory. The communication has been amortized into the shadow of the computation.

Image 1 (the timeline diagram in the project) shows exactly this structure: Stream A carries Layer N forward while Stream B carries AllReduce(Layer N-1) offset in time, with the overlapped region visually occupying ~38% of total step time (serial: 100ms → overlapped: 62ms). The reduction is precisely the fraction of AllReduce time that fits inside the compute shadow.

The resource sharing reality: both streams compete for the same SM pool. An AllReduce that saturates InfiniBand typically does so from a small number of NCCL kernel threads — the SMs are not fully occupied by the communication kernel. This is why overlap works: ncclAllReduce on the H100 at 100 GB/s outbound uses perhaps 4–8 SMs for the reduction tree and DMA engine; the remaining 124+ SMs remain available for the matmul on the compute stream.

4. Where It Breaks: The InfiniBand Wall

Overlap is not a free lunch. It fails precisely when the AllReduce duration exceeds the compute window — when there is no compute shadow large enough to hide the communication.

Quantified: Llama-7B forward pass on 16 H100 nodes

Llama-7B has ~7B parameters → ~14 GB in BF16. With ZeRO-3, each of 128 GPUs holds ~109 MB of parameters. Before each layer’s forward, FSDP issues an AllGather to reconstruct the full layer weights. Per-layer AllGather volume ≈ 14 GB / 32 layers ≈ 437 MB.

At 100 GB/s inter-node bandwidth per node (InfiniBand NDR dual-NIC), AllGather time per layer:

t_comm = 437 MB / 100 GB/s ≈ 4.4 ms per layer

Now compute. A Llama-7B layer’s dominant operation is two MLP matmuls: [batch × seq, 4096] × [4096, 11008] and its transpose. At batch=4, seq=2048, that is [8192, 4096] × [4096, 11008]. FLOPs = 2 × 8192 × 4096 × 11008 ≈ 739 GFLOPs per matmul. On 8 H100s with TP=8, each GPU handles 1/8 of this: ~92 GFLOPs per GPU. At 330 TFLOPS effective (33% MFU):

t_compute ≈ 92 × 10⁹ / 330 × 10¹² ≈ 0.28 ms per matmul

Two matmuls + attention + norms ≈ ~0.8 ms total compute per layer.

The ratio is 4.4 ms communication : 0.8 ms compute = 5.5×. There is no shadow. Overlap recovers at most 0.8 ms of the 4.4 ms communication cost; the remaining 3.6 ms is pure blocking stall per layer. Across 32 layers: 115 ms of unrecoverable communication overhead per forward pass, against ~26 ms of compute. The cluster is 82% network-bound even with perfect overlap implementation.

This is why Llama-7B at small batch is a poor fit for large-scale ZeRO-3 without additional techniques (gradient checkpointing to reduce activation memory, sequence parallelism to increase per-layer compute, or simply more batch).

5. NCCL Implementation Specifics: The Knobs That Actually Matter

NCCL (NVIDIA Collective Communications Library) is not a black box that magically performs AllReduce. It has an internal architecture that directly determines how well the overlap trick works.

Separate communication SMs: on Hopper, NCCL reserves a dedicated SM subset from the device for communication kernels when using the NCCL_MIN_NCHANNELS / NCCL_MAX_NCHANNELS environment variables. The default channel count on H100 is 16–32 channels (each channel maps to a ring segment and uses roughly one SM for the reduction tree). Setting NCCL_MAX_NCHANNELS=2 starves NCCL of SMs but leaves more for compute overlap; setting it to 32 maximizes bandwidth at the cost of SM contention with the compute stream. The optimal value is workload-specific and should be swept at cluster commissioning time.

NCCL_PROTO: NCCL implements three protocols — Simple, LL (low-latency, 8-byte atomic), and LL128 (128-byte atomic for Volta+). For InfiniBand at large message sizes (>1 MB), Simple is almost always fastest because it minimizes software overhead per byte. For small tensors or latency-critical operations, LL or LL128 reduce the per-message setup cost. Setting NCCL_PROTO=Simple explicitly prevents NCCL from auto-selecting LL for large AllReduces, which can save 5–15% on bandwidth-bound operations.

NCCL_NSOCKS_PERTHREAD and NCCL_SOCKET_NTHREADS: these control the socket-level parallelism in the proxy threads that manage InfiniBand QPs. At NDR speeds (400 Gb/s), the bottleneck can shift to the proxy thread CPU cycles. Setting NCCL_NSOCKS_PERTHREAD=8 and NCCL_SOCKET_NTHREADS=2 (4 threads × 8 sockets = 32 parallel socket operations) is a common starting point for NDR clusters; the NVIDIA NCCL documentation recommends tuning these empirically using nccl-tests/all_reduce_perf against your specific topology before declaring a cluster production-ready.

NCCL_BUFFSIZE: the internal pipeline buffer per channel. Default 4 MB. On InfiniBand, increasing to 8–16 MB reduces the number of pipeline stages for large tensors and improves effective bandwidth by ~10% on transfers >100 MB. The tradeoff is increased GPU memory usage (NCCL_BUFFSIZE × NCCL_MAX_NCHANNELS bytes reserved per AllReduce).

The cudaLaunchKernel + ncclGroupStart/End pattern: when issuing multiple NCCL operations in a step (FSDP issues AllGather + ReduceScatter per layer), wrapping them in ncclGroupStart() / ncclGroupEnd() coalesces the operations into a single NCCL kernel launch. Without grouping, each AllGather is a separate CUDA kernel launch with associated setup overhead (~5–15 µs per launch at scale). With grouping, the launches are fused. FSDP’s PyTorch implementation uses this internally; custom kernels should replicate it.

6. PyTorch / Triton Implications: When `torch.cuda.stream()` Helps and When It Hurts

When it helps: anywhere you have work that is logically independent and should not serialize. The canonical pattern:

compute_stream = torch.cuda.Stream()
comm_stream    = torch.cuda.Stream()

with torch.cuda.stream(compute_stream):
    output = layer_forward(input)          # matmul, etc.

with torch.cuda.stream(comm_stream):
    dist.all_reduce(grad_tensor, async_op=True)   # non-blocking NCCL

# Explicit synchronization only at dependency boundary
compute_stream.wait_stream(comm_stream)

This works because dist.all_reduce(async_op=True) returns a Work handle without blocking, and the NCCL kernel runs on comm_stream while the matmul runs on compute_stream. The CUDA SM scheduler interleaves them. FSDP’s implementation (arXiv:2304.11277, Algorithm 1) does exactly this — the AllGather for layer k+1 is issued on a separate stream before layer k’s computation completes, with a stream dependency (comm_stream.wait_stream(compute_stream)) ensuring the AllGather doesn’t read uninitialized memory.

When it hurts:

First, spurious synchronization. If you call .item(), .numpy(), any Python-side tensor inspection, or print a CUDA tensor inside a streamed block, PyTorch forces a cudaDeviceSynchronize(). This collapses the entire stream concurrency into a sequential barrier. In practice, debugging code (print(loss.item())) silently destroys the overlap that the surrounding infrastructure carefully constructed.

Second, Triton kernel launches. Triton kernels launched via @triton.jit functions currently do not respect torch.cuda.stream() context managers unless you explicitly pass the stream via triton.runtime.driver.active.get_current_stream() and set it on the kernel launch. As of PyTorch 2.x, the torch.cuda.stream() manager does set the current CUDA stream for CUDA API calls, but Triton’s internal launcher path can bypass this on some builds. The safe pattern when mixing Triton and NCCL is to explicitly synchronize between streams at Triton kernel boundaries.

Third, stream proliferation. Each torch.cuda.Stream() allocates GPU memory for the stream’s command buffer. Creating streams per-layer (one AllGather stream per transformer layer) in FSDP exhausts the command buffer pool and introduces scheduling overhead. PyTorch FSDP uses a fixed pool of two streams — one for computation, one for communication — and reuses them across layers precisely to avoid this.

7. The Synthesis: What the Overlap Pattern Actually Buys You

The compute-communication overlap in ZeRO-3 / FSDP is a temporal reuse trick, not a bandwidth trick. It does not move fewer bytes. FSDP transfers more bytes than DDP (AllGather + ReduceScatter instead of AllReduce = ~3× the total volume). What it achieves is spreading that transfer over the full step duration rather than concentrating it at a synchronization barrier.

The practical ceiling on the gain is min(t_comm, t_compute). When compute exceeds communication — large batch, large model, fast network — overlap is nearly perfect and communication overhead approaches zero. When communication dominates — small model, small batch, slow network — overlap recovers at most t_compute worth of communication cost; the rest is irreducible. The Llama-7B example above sits firmly in the second regime at small batch.

The engineers who extract real MFU from InfiniBand clusters are the ones who understand this not as a library feature but as a scheduling problem: your job is to maximize the compute shadow. That means: large global batch, aggressive sequence parallelism to deepen the per-step compute, and per-layer AllGather granularity fine-tuned so the AllGather of layer k+1 completes before the compute of layer k finishes — not a millisecond before (memory waste), not a millisecond after (stall).

The hardware does not forgive misalignment. The SM scheduler is not a forgiving abstraction. But when you get it right, two orders of magnitude of bandwidth gap become nearly invisible — and that is one of the more remarkable engineering sleights of hand in modern ML infrastructure.

References

Zhao, Y. et al. (2023). PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel. arXiv:2304.11277.
Rajbhandari, S., Rasley, J., Ruwase, O., & He, Y. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. arXiv:1910.02054.
NVIDIA. NCCL Developer Guide. https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/index.html
NVIDIA. H100 Tensor Core GPU Architecture Whitepaper. 2022.
Shoeybi, M. et al. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv:1909.08053.
Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.

BibTeX

@article{fp4-2606011,
  title   = {CUDA Streams and Kernel Concurrency: The Overlap Engine Behind ZeRO-3 and FSDP},
  author  = {fp4 editorial desk},
  year    = {2026},
  url     = {https://fp4.dev/silicon/cuda-streams-overlap/},
  journal = {fp4}
}