Skip to content

Add devlog: Cross-stream memory management: wait_event instead of record_stream - #43

Merged
weifengpy merged 1 commit into
mainfrom
cross_stream_memory
Sep 30, 2026
Merged

weifengpy merged 1 commit into
mainfrom
cross_stream_memory

Conversation

@weifengpy

Copy link
Copy Markdown
Contributor

New post under content/eager/: Cross-stream memory management: wait_event instead of record_stream.

What it covers

  • Why producer/consumer pipelining across CUDA streams (e.g. FSDP2's all-gather → forward) needs explicit lifetime management: the caching allocator only tracks the stream a block was allocated on.
  • Why Tensor.record_stream is correct but hands block reuse to the allocator's event polling. Peak memory then depends on how far the CPU runs ahead, with links into CUDACachingAllocator.cpp.
  • Two fixes: stall the producer with wait_event, and the recommended one, keep a Python reference and wait_event on the allocation stream before del.
  • How FSDP2 applies this in forward (AllGatherState) and backward (ReduceScatterState, set_reduce_scatter_max_input_buffers).

Files

  • content/eager/2026-09-29-cross-stream-memory-without-record-stream.md
  • static/images/eager/cross-stream-*.svg: 9 swimlane diagrams with light and dark variants via prefers-color-scheme

Verification

  • All PyTorch source links are pinned to 1a0b56b8 on main, and each claim was checked against the code there.
  • The demo script (in the post's appendix) was run on an H100:
    • no lifetime management corrupts 50/50 trials;
    • record_stream, stall and keepalive all pass;
    • they use at most 4, 1 and 2 param blocks per run, respectively.
  • hugo --minify builds cleanly, and the page was checked in headless Chromium in light and dark mode.

…ord_stream

Explains why producer/consumer pipelining across CUDA streams needs explicit
memory lifetime management (the caching allocator only tracks the allocation
stream), why Tensor.record_stream makes peak memory depend on CPU run-ahead,
and the recommended pattern: keep a Python reference, record an event on the
consumer, and wait_event on the allocation stream before dropping the
reference. Includes a single-GPU demo script and walks through how FSDP2
applies the pattern in forward (all-gather) and backward (reduce-scatter).
@weifengpy
weifengpy merged commit 74462e4 into main Sep 30, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants