fix(dist): make raw IPC input pools usable — remove init_dist_env's vestigial signal/buffer block, add explicit raw-pool override

[ThomasNing]

Fixes #4921.

Summary

Two commits:

1. init_dist_env no longer breaks raw IPC input pools — the vestigial signal/buffer block is removed.
2. AITER_CUSTOM_AR_RAW_INPUT_POOL — an explicit trigger for the raw (plain-hipMalloc) input pool that works, because the existing trigger (expandable segments) turns out to be broken further down regardless.

Root cause — three stacked failures, not one

#4921 reports the crash at communication.py:78 (ca_comm.buffer = ca_comm._pool["input"].tensor). Measured on 4×/8× MI350X (gfx950, both VF and bare metal), the full chain under PYTORCH_HIP_ALLOC_CONF=expandable_segments:True is:

1. The first failure is earlier than the issue's traceback: register_input_buffer(signal) exports the torch.zeros signal tensor's pointer via hipIpcGetMemHandle, and under expandable segments that pointer is VMM-backed — the export dies at custom_all_reduce.cu:417 ("invalid argument"). The issue's line-78 traceback is the signature of the explicit-override trigger, where the allocator stays default and registration survives.
2. Line 78 — the raw pool has no backing torch.Tensor, and .tensor raises by design.
3. With 1 and 2 fixed, a third failure appears (previously unreachable): CUDA-graph capture completes, then flush_graph_buffers dies at custom_all_reduce.cuh:3448 exporting capture-time output tensors — every out = torch.empty_like(inp) allocated during capture is VMM-backed, get_output_buffer_RD records it, and the post-capture flush tries to IPC-export it. 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 today (failure 3 needs output staging — a separate, kernel-level change). The explicit override gives the raw pool a trigger that works: under the default allocator everything else (meta pool, capture-time outputs, graph flush) stays exportable, and only the input pool moves to hipMalloc.

Why removal instead of .tensor → .data_ptr

The issue asked which contract ca_comm.buffer serves. Answer: none —

• ca_comm.buffer / ca_comm.signal are only ever written, never read (swept aiter and a downstream engine, including hasattr gates);
• every internal consumer of the input pool goes through _pool["input"].data_ptr;
• C++ register_input_buffer only inserts a pointer-translation entry keyed by the registered tensor's own address, consulted only when an allreduce is invoked with that exact tensor as input — which never happens for the signal tensor (open_ipc_handle's cache fills on demand, so no pre-warming is lost);
• gfx1250 already skips the entire block in-tree and works.

The old comment said the attributes are "never read anywhere" while the code still ran the registration — removal makes the code match the comment. Numerics were validated, not just startup (see below), since a wrong guess here would corrupt allreduce silently.

Measured validation (GLM-5.2-MXFP4, TP4, --level 3 graphs, gpu_memory_utilization 0.90)

Two environments, stated per row — provenance matters here because bring-up
behaviour in this problem space has proven state-sensitive (see motivation):

• [A] bare-metal gfx950, image nightly_202608041536 + a source overlay of the two files this PR touches plus their csrc (JIT-rebuilt in-container; ≡ aiter e404860 + this PR).
• [B] gfx950 VF, stock image nightly_202608191459 (aiter 4fa508e) with this PR's files bind-mounted, next to another tenant's live engine.

config	env	result
unpatched, expandable segments	A	❌ dies at cu:417 (signal export — failure 1)
unpatched, expandable segments	B	❌ same signature
patched, expandable segments	A	❌ dies at cuh:3448 post-capture (failure 3 — expected, out of scope)
patched + AITER_CUSTOM_AR_RAW_INPUT_POOL=1	A	✅ READY at full envelope, 0 IPC errors
patched, default alloc	A	✅ READY (no-regression)
patched + override, second co-resident TP4 engine	A	✅ both READY at 0.90; TPOT identical to solo (0.0196 s) — see IPC-matrix caveat below
patched + override	B	✅ READY, 0 IPC errors; TPOT 0.01561 vs 0.01664 unpatched-default (no regression)

Numerics: inference through the raw pool produces correct output (spot-checked deterministic arithmetic prompts; allreduce path exercised end-to-end at TP4 with HIP graphs).

Operational notes for the co-resident-engines use case that motivated #4921:

• With this patch, on the older engine-stack image used for the dual-engine run, second-engine bring-up was nondeterministic across IPC-namespace configurations (measured matrix, one run per cell: each engine solo READY under host or private IPC; host+host → e2 dies at NCCL init; host+private sequential → both READY, benched below; private+private → e2 dies at NCCL init whether brought up sequentially or concurrently). We do not claim an IPC-namespace mechanism — and, importantly, no configuration axis has survived testing: an independent replication ran the identical unpatched two-engine command that had succeeded 3/3 in one window and had it fail hours later on a verified-clean node (no orphaned IPC segments, no stale KFD holders, 0% VRAM). Unpatched second-engine bring-up appears state- or timing-dependent under identical configuration; image, IpcMode, envelope, GC settings, capabilities and shm size have all failed to separate outcomes. A success-rate measurement is in progress; treat every single-run cell in the matrix above (including the successes) accordingly. The dual-engine numbers in the table were measured on the one configuration of that matrix that came up (engine 1 --ipc host, engine 2 --ipc private, verified via docker inspect).
• Without this patch, the unpatched two-engine workaround is nondeterministic under identical configuration: an independent VRAM-instrumented effort (clean boxes, IpcMode captured via docker inspect) had one successful bring-up of the 0.62/128 two-engine configuration — whose engine pair then stayed durable through six consecutive benchmark passes — and the identical command failed at hipIpcGetMemHandle hours later on a node verified clean beforehand. Every envelope above 0.62/128 has failed in all unpatched attempts, and no configuration axis tried so far (image, IpcMode, GC settings, capabilities, shm size) separates success from failure. The failure is concentrated entirely in bring-up: it fails fast (~50 s), leaves the first engine healthy, and pairs that do come up stay durable — so the operational shape of the problem is a bring-up lottery, not a serving instability. That is the practical motivation for this override — with one scope limit found by the same investigation: under co-residency on at least one stack, NCCL_DEBUG=INFO shows RCCL's own P2P transport failing hipIpcGetMemHandle first (p2p.cc:256), with AITER's cu:417 failure downstream of it. So this override removes the AITER-side export dependency (an exportable-by-construction hipMalloc pool instead of allocator-state-dependent pointers); it does not claim to make two co-resident engines reliable by itself, because RCCL's own export sits upstream and is outside this PR. (NCCL_DMABUF_ENABLE=1 was evaluated for that layer and did not help: 0/3 with it and 0/3 without, interleaved, p2p.cc:256 present in all six.) Side by side: 0.62/128 unpatched with an unquantified success rate — vs 0.90 util / 512 seqs patched (measured above, single-engine).

Co-residency bring-up rates (appended after measurement — the promised control)

Measured on the matched image (nightly_202608191459), sequential bring-up, second engine on disjoint GPUs, NCCL_DEBUG enabled, teardown + clean-node checks between attempts:

arm	host	attempts	successes	first failure
patched + override (pool armed, logline-proven)	bare-metal, iommu=pt	5	0	RCCL p2p.cc:256 every time
unpatched control, same envelope, same host	bare-metal, iommu=pt	3	0	RCCL p2p.cc:256 every time
patched + override	VF, no iommu flag	9+	1	RCCL p2p.cc:256 every time

Conclusions, stated at their real strength:

• This patch is orthogonal to co-residency bring-up — patched 0/5 vs unpatched 0/3 on the same host and envelope, identical failure signature. It neither causes nor cures the bring-up lottery, which is what "no co-residency claim" means concretely.
• The first failure on both hosts, both stacks, both IpcMode shapes is RCCL's own P2P transport (hipIpcGetMemHandle → invalid argument at p2p.cc:256), before any AITER pool exists. A separate issue against RCCL is being prepared with full NCCL_DEBUG=INFO evidence from both hosts.
• The one successful bring-up (pooled ~1/30 across hosts) was durable: six benchmark passes without incident, and across ~25 consecutive second-engine failures the incumbent engine never dropped or needed restarting — the fault is scoped to the joining process's export path.
• What this PR fixes is unaffected by all of the above: the [Bug] raw_cached IPC pool is unusable: init_dist_env reads _pool['input'].tensor, which raises for plain-hipMalloc buffers #4921 contract break dies at init deterministically without it, and the override arms the raw pool (logline-proven on two hosts).
• The dual-engine numbers above are not "co-residency is free": at c=32 each engine ran ~21–27% below solo in a short window (aggregate +51%), with identical TPOT — decode is unaffected; the loss is prefill/queueing contention. Note the dual aggregate curve turns over past the knee (c=128+128 measured below c=32+32 in this short-window harness); don't extrapolate aggregate upward with concurrency.

Test

op_tests/multigpu_tests/test_init_dist_env.py (new): brings up init_dist_env per rank and checks one allreduce under three modes — default, expandable, raw_override — asserting the override actually selects the raw pool. The existing test_custom_allreduce.py performs its own init and never executes init_dist_env, which is how the regression shipped.

Not covered by CI: failure 3 (needs graph capture + expandable segments); documented here and reproducible with the table above.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4924 --add-label <label>

PR title tags:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf] and op tags like [MLA] are left untouched. Add the no-auto-title label to opt this PR out of title tagging.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: e2a3a827ae

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Set allocator configuration before importing aiter

Under the normal import path, this top-level import loads aiter, which imports custom_all_reduce.py; its module initialization calls _detect_gfx1250() and queries the CUDA device. Because spawned workers import this module before entering _worker, CUDA can be initialized before PYTORCH_HIP_ALLOC_CONF is set on lines 38–39. The allocator snapshot then remains non-expandable, and the warning-only branch on lines 96–103 lets the expandable run pass without exercising the regression. Move the aiter import after the environment assignment or launch workers with the allocator variable already present in their environment.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Assert the result of the allreduce comparison

checkAllclose returns a nonzero mismatch ratio for ordinary incorrect results and does not raise with its default catastrophic_check=False. Since this return value is discarded, even an allreduce whose entire output is wrong can leave every mode passing, so the advertised correctness regression check is ineffective. Assert an acceptable returned error ratio.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Skip raw-pool introspection for the VMM transport

On gfx1250 systems using the supported VMM path—ROCm older than the IPC threshold, an undetectable ROCm version, or AITER_CUSTOM_AR_FORCE_VMM=1—ca_comm._pool["input"] is a _GFX1250BufferProxy._Entry, which has data_ptr and max_size but no _raw_cached. This unconditional access therefore raises AttributeError in every run mode before the allreduce executes. Gate this IPC-specific assertion on not ca_comm._use_vmm or represent VMM as a separate pool mode.

Useful? React with 👍 / 👎.