[vasilevklart]

_capture_thrown_away sizes a capture with torch.cuda.memory_snapshot(pool.handle), and torch drops a reference to the pool id on every such call (torch 2.9.1, 2.12.1, 2.13.0). Once a process has made more than 256 graph pools, the id is no longer a cached int, so the call frees it while the pool still holds it and the next release writes into freed memory: malloc_consolidate(): unaligned fastbin chunk, or a segfault in the caching allocator. In main since #341.

This sums the pool's bytes from the full snapshot by segment_pool_id, which passes torch nothing of ours. The number is the same.

1×H100, 4 GPU test modules in one process	runs	aborted
main	15	15
this PR	2	0

Tests: test_capture_budget.py::test_sizing_a_throwaway_capture_keeps_the_references_to_its_pool_id fails 2/2 on main and passes 2/2 here; test_capture_budget.py 46 passed. CPU Core fails the same 4 tests as main, by name. ruff clean.

The torch bug itself is worth an upstream report, separately.

AI assistance (Claude) was used for this PR.

[NSagan271]

🤖 Automated live test, AI-generated (Claude Code). Not checked line by line by a human. Numbers come from one allocation, so treat small deltas as provisional.

Live test: Fix heap corruption when sizing thrown-away CUDA graph captures

Works. The crash reproduces on main past ~256 graph pools and is gone on the PR. The new size matches the old one on the same pool every time. Every served config builds the same capture plan as main.

H200 · torch 2.9.1 · PR head 9d2e8b1 vs its merge base 5522394 · Rust runtime · FlashInfer 0.6.18

The crash

400 calls of the installed _capture_thrown_away, one process per run	runs	crashed
main	3	3 (SIGSEGV after ~250 pools, in capture_begin / memory_snapshot)
PR	3	0 (all 400 pools sized)

The number

Check	Result
100 pools below the 256 limit, each sized both ways on the same pool: main's memory_snapshot(pool) sum vs the PR's filter of the whole snapshot by segment_pool_id	✅ equal 100/100 (22 and 54 MiB pools)
Served configs, 2 restarts per arm: Qwen3.5-0.8B, Whisper large-v3, Orpheus colocated, Orpheus P/D, BAGEL P/D, Qwen3-Omni P/D	✅ the sorted Captured accelerator graph … / executable-memory lines hash identically across main, this PR and #975 for every config. Every request completed, with no segfault or heap error at exit under MALLOC_PERTURB_=165
Startup (sizing now reads the whole snapshot)	✅ ready times no slower than main's restarts on any config: Qwen3.5 194, 194 vs 199–211 s, Whisper 124–136 vs 125–154 s, Qwen3-Omni P/D 260–309 vs 261–342 s

Tests

• test_sizing_a_throwaway_capture_keeps_the_references_to_its_pool_id passes 3/3. The GPU capture/KV modules (8 files, 105 tests) pass in 2 of 3 runs; the third run fails one test, test_kv_ipc_transfer.py::test_a_read_from_another_process…. It is an exit segfault in an IPC reader subprocess, and main and Fix exit segfault from KV caches shared over CUDA IPC #975 fail it 1 in 3 as well (see the Fix exit segfault from KV caches shared over CUDA IPC #975 report), so it isn't caused by this PR.

Commands

Harness: _capture_thrown_away(capture, …) around a small two-GEMM graph, 400 times, plus a fresh pool per iteration kept alive · same-pool check: capture into graph_pool_handle(), then sum memory_snapshot(pool) vs memory_snapshot() filtered by segment_pool_id

[NSagan271]

LGTM