[Bug] raw_cached IPC pool is unusable: init_dist_env reads _pool['input'].tensor, which raises for plain-hipMalloc buffers

[ThomasNing]

Summary

aiter/ops/communication.py:78 (init_dist_env) unconditionally does:

ca_comm.buffer = ca_comm._pool["input"].tensor

But when CustomAllreduce builds its input pool with raw_cached=True, that buffer is a plain hipMalloc region with no backing torch.Tensor, so the tensor property raises by design (custom_all_reduce.py:447):

RuntimeError: Uncached IPCBuffer has no backing tensor; use .data_ptr

Net effect: the raw_cached (plain-hipMalloc, IPC-exportable) input pool cannot be used at all. Any configuration that selects it fails during distributed init.

That matters because raw_cached is the path that keeps the pool exportable via hipIpcGetMemHandle, and it is currently the only mitigation for co-resident engines — see Impact.

Reproduction

8× MI350X (gfx950, SR-IOV VF), ROCm 7.2, rocm/atom-dev:nightly_202608191459 (bundled AITER 4fa508ef), GLM-5.2-MXFP4, one TP4 engine:

RuntimeError: Uncached IPCBuffer has no backing tensor; use .data_ptr
  File "/app/aiter-test/aiter/ops/communication.py", line 78, in init_dist_env
    ca_comm.buffer = ca_comm._pool["input"].tensor
  File "/app/aiter-test/aiter/dist/device_communicators/custom_all_reduce.py", line 447, in tensor
  File "/app/ATOM/atom/model_engine/model_runner.py", line 941, in _setup_device_and_distributed

raw_cached is gated at custom_all_reduce.py:1062:

raw_cached = _expandable_segments_enabled()

so it is reached two ways, and both fail identically:

1. PYTORCH_HIP_ALLOC_CONF=expandable_segments:True — the supported trigger.
2. A one-line local patch adding an explicit override (... or _env_flag("AITER_CAR_RAW_CACHED")).

Controls run to exclude our own environment:

configuration	result
patched file mounted, override off (raw_cached=False)	server READY — so the patch itself is harmless
override on, with our unrelated local AITER patches mounted	dies
override on, pristine container (no other patches)	dies — same traceback

So this is not caused by anything we mount, and it is not a version skew: the only custom_all_reduce* change between fc2e5d57 (v0.1.20) and 4fa508ef is @compile_ops(..., develop=True) on three fused_allreduce_mhc_* stubs.

Impact

On this VF hardware, running two TP4 engines on one 8-GPU node fails at the default envelope:

custom_all_reduce.cu:417 fail to call hipIpcGetMemHandle(...) ---> [HIP error](invalid argument)

The first engine starts; the second cannot export its handle. Consistent with the torch.empty pool returning a caching-allocator pointer that is not IPC-exportable — exactly what raw_cached exists to avoid, and what #4174 / #4621 address for expandable segments.

Because raw_cached is unusable, the only working mitigation we found is to buy allocator headroom:

--gpu-memory-utilization / --max-num-seqs	two engines?
0.78 / 256	no
0.70 / 128	no
0.62 / 256	no
0.62 / 128	yes

Both dimensions must drop together. That costs real KV cache and half the sequence slots. Measured node throughput at 0.62/128 is 8465 tok/s aggregate; the alternative (guarding the second engine with NCCL_P2P_DISABLE=1 + AITER_DISABLE_IPC_COLLECTIVES=1) drops that engine's own end-to-end throughput by ~50% and yields 6440 aggregate.

Suggested fix

init_dist_env should use the pointer rather than the tensor, since IPCBuffer already exposes data_ptr for exactly this case — e.g. keep .tensor only where a tensor is genuinely required, or have init_dist_env branch on uncached/raw_cached. Happy to send a PR if you confirm the intended contract for ca_comm.buffer (tensor vs raw pointer) — it is consumed downstream and I did not want to guess.

Separately, it would help operationally if raw_cached were selectable on its own (e.g. AITER_CAR_RAW_CACHED=1) rather than only as a side effect of expandable_segments: co-resident deployments need the exportable pool but do not necessarily want expandable segments, which independently prevents this engine from starting.

Environment

• 8× MI350X (gfx950) SR-IOV VF, ROCm 7.2, TP4 per engine
• rocm/atom-dev:nightly_202608191459, bundled AITER 4fa508ef2935110ff99adf2743ea93807dbd9c67
• Model GLM-5.2-MXFP4, --level 3 --method mtp --num-speculative-tokens 3
• PYTORCH_HIP_ALLOC_CONF unset unless stated

[ThomasNing]

Correction to the Impact section — the mitigation needs one more condition

The envelope table in the body is incomplete in a way that would block reproduction, so stating it explicitly:

Two co-resident TP4 engines also require that the containers do NOT share the host IPC namespace (i.e. no docker run --ipc=host). With host IPC shared, the second engine fails on hipIpcGetMemHandle → invalid argument even at 0.62 / 128. Controlled A/B, identical envelope:

envelope	--ipc=host	second engine
0.62 / 128	yes	fails
0.62 / 128	no	starts, serves

So the working configuration is both conditions together: own IPC namespace and --gpu-memory-utilization 0.62 --max-num-seqs 128. Neither alone is sufficient — without host IPC, 0.78/256 and 0.70/128 still fail; with host IPC, even 0.62/128 fails.

I do not have a mechanism for the --ipc=host dependency and am not going to guess at one. HIP IPC goes through KFD rather than SysV IPC, so a docker namespace flag should not matter, and the obvious hypothesis — ROCr keeping IPC bookkeeping in host POSIX shared memory — was tested and falsified: host /dev/shm entry count, psm_* segment count and SysV segment count are all unchanged across baseline / first-engine-up / second-engine-failed. It may or may not be related to the raw_cached issue this report is about; I mention it only so the mitigation is reproducible.

For completeness on the other numbers in the body: the guarded-engine figure is measured at 4252 (unguarded) vs 2188 (guarded) end-to-end output tok/s, i.e. −48.5%, and the aggregate figures are 8465 (both unguarded, 0.62/128) vs 6440 (second engine guarded, 0.78/256).

[ThomasNing]

Fix up as #4924, with the full failure chain root-caused and measured. Three findings that update this issue's text:

1. Trigger 1's first failure is earlier than the traceback above. Under expandable_segments:True, register_input_buffer(signal) dies first — the torch.zeros signal tensor is VMM-backed and hipIpcGetMemHandle rejects it at cu:417. The line-78 .tensor crash is trigger 2's signature. Both are fixed by removing the block (it is vestigial — details in the PR).

2. A third failure was hiding behind these two: with both fixed, capture completes and flush_graph_buffers dies at cuh:3448 exporting capture-time output tensors (torch.empty_like → VMM; get_output_buffer_RD records them). fix: car graph mode err when pytorch set expandable_segments:True #4621's copy-in guard covers inputs only. So expandable segments cannot arm the raw pool usefully; fix(dist): make raw IPC input pools usable — remove init_dist_env's vestigial signal/buffer block, add explicit raw-pool override #4924 adds AITER_CUSTOM_AR_RAW_INPUT_POOL=1 as a working trigger instead (default allocator, only the input pool moves to hipMalloc).

3. On the Impact section (co-resident engines): measured on both bare-metal and VF gfx950 — the second engine's failure is the shared IPC namespace, not the allocator. Two TP4 engines at the full 0.90 envelope both come up clean when the second uses --ipc private (with --ipc host on both: NCCL init + cu:417 failures). With the fix + override + private IPC: dual-engine TPOT is identical to solo (0.0196s), no 0.62/128 squeeze, no NCCL_P2P_DISABLE guards. The 0.62/128 mitigation table above may therefore be obsolete — a VRAM-instrumented re-measure is in progress separately.

[ThomasNing]

Correction to my point 3 above, after an independent VRAM-instrumented replication (clean box, no other containers, IpcMode logged):

"The wall is the shared IPC namespace" holds for PATCHED engines only. For unpatched AITER, private IPC is not sufficient — the published 0.62/128 workaround succeeds on one gfx950 VF host and fails with hipIpcGetMemHandle → invalid argument on another, same config, private IPC verified on both. So the unpatched working envelope is host/image-specific rather than a portable recipe, which is the stronger motivation for #4924's explicit override: the mitigation this issue documented is not reliably reproducible, independent of its KV cost.

[ThomasNing]

Additional data: the workaround has an image floor, and the mechanism is not the one you'd guess

Following up on the envelope workaround in the body, with two results that bear on how much this issue is worth fixing.

1. The 0.62 / 128 workaround is version-fragile. Same host, same envelope, same private IPC namespace, clean baseline (8 GPUs at 0% VRAM, no other containers, per-GPU VRAM logged, IpcMode=private verified by inspect rather than assumed):

ATOM image	bundled AITER	second TP4 engine
rocm/atom-dev:nightly_202608181633	fc2e5d57 (v0.1.20)	fails — hipIpcGetMemHandle(...) → [HIP error](invalid argument) at custom_all_reduce.cu:417
rocm/atom-dev:nightly_202608191459	4fa508ef	both engines run, unguarded

So the allocator-headroom workaround does not merely cost KV cache — whether it works at all depends on the image. That seems worth weighing when prioritising the fix: an exportable pool would remove the need to find a working envelope per image.

2. I am not attributing that flip to a commit, and I'd caution against assuming one. The AITER delta between fc2e5d57 and 4fa508ef contains no change to the IPC/buffer path — the only Python change is @compile_ops(FUSED_AR_MHC_MD_NAME, develop=True) on three fused_allreduce_mhc_* stubs (#4832), plus C++ edits to fused_ar_mhc_post. The two images also differ in ATOM's own source, so the responsible change may not be in AITER at all. Treat the floor as empirical.

3. For completeness, the working configuration once it does work (replicated on a second host, 3 reps, clean baseline): two TP4 GLM-5.2-MXFP4 engines on one 8-GPU MI350X-VF node, disjoint ROCR_VISIBLE_DEVICES, private IPC namespaces on both, --gpu-memory-utilization 0.62 --max-num-seqs 128, no collective guards → 8570 tok/s aggregate at node concurrency 256, engines balanced within 1–2%. Under host IPC the second engine fails at the same envelope, so private namespaces are necessary but not sufficient — the envelope is required too, which is the part this issue would eliminate.

I have no mechanism for the host-IPC dependency and am not going to speculate: the obvious hypothesis (ROCr IPC bookkeeping contending over host POSIX shared memory) was tested and falsified — host /dev/shm entry count, psm_* segment count and SysV segment count are unchanged across baseline / first-engine-up / second-engine-failed.

[ThomasNing]

Final measurements are in; #4924 carries the full table. Summary for this issue:

What #4924 fixes (deterministic, validated): the contract break reported here — init_dist_env dies at init on any raw input pool (and, first, on exporting the VMM-backed signal tensor under expandable segments). With the fix, raw-pool engines initialize, serve, and produce correct output on both a bare-metal and a VF gfx950 host; the new per-rank logline proves which pool armed in every run.

What it deliberately does not fix — measured, both hosts: co-resident second-engine bring-up. Rates at the full 0.90/512 envelope on the matched image: patched 0/5 and unpatched 0/3 on the same bare-metal host (patch is orthogonal — neither causes nor cures it), patched 1/9+ on the VF host. First failure in every attempt on both hosts is RCCL's own P2P transport (p2p.cc:256 hipIpcGetMemHandle → invalid argument), before any AITER pool exists. The one successful bring-up was durable (six benchmark passes), and across ~25+ consecutive second-engine failures the incumbent engine never dropped — the fault is scoped to the joining process's export path. A separate issue against RCCL is being prepared with full NCCL_DEBUG=INFO evidence from both hosts; nothing about it is AITER's cu:417, which is downstream.

Not predictive of the outcome, all tested: image, envelope (0.62–0.90), IpcMode shape, NCCL_DMABUF_ENABLE (0/3 vs 0/3 interleaved), collective guards, iommu=pt (present on the 0/5 host, absent on the 1/9 host), AITER pool mode. No mechanism is claimed.