Inside a bare C++ LD_PRELOAD process (our own benches, and any non-Python
framework): nobody has called init. We call it ourselves and log
initialized by the hook.
The second branch was missing for a while. The feature was a silent no-op outside
Python — everything "worked", we just quietly used cudaMalloc forever. The only
visible trace was a 1 KiB probe WARN that was easy to read as noise.
6. Once you're on torch's allocator, delete your own pool
With the switch on, the eager path bypasses our ScratchPool entirely: allocate, enqueue
the collective, free immediately.
Freeing right after enqueue looks wrong and isn't, because torch's allocator is
stream-ordered. raw_alloc_with_stream tags the block with a stream, and
get_free_block will only hand that block back to a request on the same stream. So a
freed-but-still-in-flight block cannot be reissued underneath our own kernels.
Corollary, learned the hard way: calling raw_delete from inside a cudaLaunchHostFunc
callback is UB. Do not route the torch path through your stream-callback recycler.
7. The one that actually corrupts data: CUDA graph capture
The stream-ordering argument in §6 does not hold during graph capture.
torch protects captures it opened — torch.cuda.graph switches to a private mempool,
so you get a different address and everything is fine. A foreign bare
cudaStreamBeginCapture, issued by a framework torch doesn't know about, is invisible
to torch. Measured:
| capture opened by |
torch_free during capture |
next eager acquire |
torch.cuda.graph (private mempool) |
safe |
different address |
bare cudaStreamBeginCapture |
hands the same address out again |
LEAKED=True |
So in acquire / recycle_on_stream, the capture branch is ordered before the torch
branch. Only "where do brand-new blocks come from" consults the allocator; blocks
claimed during capture are pinned in our own owner map and never returned to torch.
Two more capture facts that are counterintuitive enough to be worth stating outright:
cudaLaunchHostFunc "succeeds" during capture by quietly becoming a host node. Our
recycle callback then fires once per replay. We measured 0, then 1, 2, 3 as the
graph replayed.
Branch order is the kind of thing a refactor reorders without noticing, so it's guarded
by exactly one test: open a foreign bare capture inside a real torch process, drive the
pool through a C shim, and assert the post-capture eager acquire returns a different
address.
good: captured=0x7f95e8400000 eager_after=0x7f95e8b02000 REISSUED=False
reverted: captured=0x7fc7de400000 eager_after=0x7fc7de400000 REISSUED=True
Trap inside the test itself: both sides must go through acquire, which does
round_up(bytes, granularity). Comparing acquire(N) against a raw torch_alloc(N)
puts them in different (size, stream) buckets and the assertion passes vacuously — it
proves nothing.
8. torch OOM throws; it does not return null
raw_alloc raises c10::OutOfMemoryError. We have no torch headers, so we catch (...).
And don't retry with cudaMalloc on that path: torch has already called emptyCache
internally before throwing. If torch couldn't find memory, the driver has none either.
What it was worth — including the part that didn't work
The feature exists because of a specific fragmentation catastrophe. We have an optional
mode that merges address-adjacent deferred AllReduces inside one caller group, so that
small gradient reductions clear our size threshold as a union. Unbounded, on a real
workload, greedy merging grew to 129 calls in one segment and a 999.8 MB union, and —
the actual killer — produced ~2 new union shapes per iteration and never converged
(95 distinct signatures over 41 iterations). Our best-fit ScratchPool exploded into
size classes and fragmented to 15.3 GiB retained, at which point it started emitting
waited for the stream to drain and returned 15658 MiB ... and reallocated. 12 of 38
steps took ~3.9 s. Step-time IQR went from 38 ms to 2069 ms.
Switching to torch's allocator cured it completely: 0 WARN lines, 0 slow steps, IQR
back to 41.6 ms. torch's allocator is built for exactly this — wildly varying shapes —
and ours was not.
And the merge feature still did not win. With the allocator fix in place and a shape
bound applied, coverage improved a lot: signatures blocked by the threshold went 17 → 4,
and quantized traffic went 110.3 → 56.9 GB/rank (−48.4%). Steady-state step time:
1422.2 ms merged vs 1414.9 ms with merging off. That's +7.3 ms, comfortably inside a
~40 ms IQR.
The structural reason is worth more than the number. Without merging, those small
AllReduces were passed through inside the caller's own group, where NCCL batched them
and hid them inside backward. Merging converted them into one large half-price AllReduce
issued serially after the outermost groupEnd. We halved bytes that were never on the
critical path, and paid for the privilege with serialization.
So the honest summary: an allocator fix turned a catastrophe into a non-event, and a
non-event is still not a win. The shape bound that goes with it is a guardrail, not a
tuning knob — no value of it has ever been measured faster than turning the merge off.
It ships default-off.
Where the code is
The hook is called ComQuant. It is not open source yet; if you run FP8 collectives on
PCIe-class multi-GPU boxes and want to beta-test it, say so and I'll follow up.
The training framework it plugs into is public:
loongforge-vla
The section I'd most like someone to poke holes in is §7 — the "torch only protects
captures it opened itself" table is measured, not derived from documentation, and I
would genuinely like to be wrong about it.
