[Bugfix][Distributed] Make custom allreduce graph capture VMM-safe

[Leonccaa]

Purpose

Backport and harden the capture-scoped allocator guard from vllm-project/vllm#43923 so 1Cat custom all-reduce can capture CUDA Graph buffers when either supported PyTorch allocator variable enables expandable segments.

Custom all-reduce exports graph buffers with legacy cudaIpcGetMemHandle. CUDA VMM allocations created by expandable segments cannot be exported through that API, so graph-buffer registration otherwise fails with invalid argument. The guard temporarily disables expandable segments for the complete capture and registration window, then restores the full effective allocator configuration.

Scope

• one runtime implementation file;
• one CPU-runnable focused test file;
• no SM70-only dispatch, model, quantization, attention, or native-kernel change;
• disabled custom all-reduce and configurations without an effective expandable_segments:True remain no-ops.

The implementation follows PyTorch 2.10 precedence: a present PYTORCH_CUDA_ALLOC_CONF wins, including an empty value; otherwise it falls back to PYTORCH_ALLOC_CONF. It changes only the expandable_segments token in the selected full configuration, preserving and restoring options such as max_split_size_mb, roundup_power2_divisions, and garbage_collection_threshold.

Duplicate-work and provenance

This started as an explicit 1Cat backport of the still-open vLLM PR #43923, not a competing implementation. Review on this PR identified two correctness gaps in the original patch: lossy partial allocator updates and missing unified-variable support. The follow-up commit fixes both locally. No open 1Cat PR covers this path. The original author remains co-author of the backport commit.

1Cat #248 concerns VMM-safe native CPU KV offload and does not cover custom all-reduce CUDA IPC graph buffers.

Tests

• focused pytest: 12 passed;
• changed-file pre-commit hooks: all passed, including Ruff, mypy, SPDX, forbidden-import, and CUDA-API checks;
• git diff --check: passed.

The tests cover full mixed-config state preservation, legacy/unified variable selection and precedence, bracketed allocator syntax, normal and exceptional restoration, graph-buffer registration ordering, and inactive/no-option no-op behavior. One stateful test models PyTorch parseArgs resetting omitted allocator fields, rather than only checking setter call strings.

PyTorch 2.10 / V100 validation

Fresh isolated validation used 2 x Tesla V100 PCIe 32 GB and torch==2.10.0+cu128 without changing a service or production route.

• Mixed legacy config snapshot: max_split_size_mb:512, roundup_power2_divisions:4, and garbage_collection_threshold:0.8 remained unchanged while expandable segments toggled True -> False -> True; the complete post-context allocator snapshot equaled the pre-context snapshot.
• Unified-only PYTORCH_ALLOC_CONF snapshot: the same True -> False -> True transition and exact post-context restoration passed.
• TP2 custom all-reduce CUDA Graph register/replay passed on both ranks with the mixed legacy config. CUSTOM was selected, each rank registered its graph address, output matched the expected all-reduce result, and allocator state was exactly restored.
• The same TP2 graph gate also passed with unified-only PYTORCH_ALLOC_CONF.

Historical full-service evidence for the same capture mechanism remains:

• Before the guard, both the fork runtime and a base runtime failed during graph-buffer registration with custom_all_reduce.cuh invalid argument.
• With the guard, both ranks registered 848 graph addresses, CUDA Graph capture completed, and the service returned the expected deterministic output.
• AWQ and FP8, each with MTP off and K4, completed the TP2 portion of a 40-cell 8K-to-256K matrix.
• In one fixed AWQ/MTP2/concurrency-4 hot A/B, CUSTOM reached 211.07 aggregate tok/s versus 188.15 for PYNCCL-only, +12.2%. This is configuration-specific evidence, not a universal throughput claim.

The PR is ready after the requested changes and latest-main source audit. The fresh gate above is a focused allocator/custom-all-reduce test, not a new current-main full-model serving matrix.

AI assistance

AI assistance was used for isolation, focused tests, validation review, and drafting. Leon reviewed the design and authorized the follow-up implementation; he remains responsible for understanding and defending the change before it is marked ready again.

Latest-main audit

• Replayed unchanged onto main 7bee39776; both commits retain the original author/co-author credits and DCO sign-offs.
• Confirmed the still-open upstream vLLM #43923 root cause and compared its patch; this version additionally preserves the complete selected allocator configuration and supports the PyTorch 2.10 unified variable.
• Confirmed PyTorch 2.10 CUDA allocator environment precedence from its installed/source header implementation.
• CPU-focused allocator/capture suite: 12 passed. All changed-file pre-commit hooks and git diff --check passed.
• No model, checkpoint, architecture, quantization, attention, or kernel selector is introduced.

[yangzhuxinyzx]

源码与轻量验证结论：当前版本不能合并，存在两个 allocator 状态正确性阻断项。

1. torch.cuda.memory._set_allocator_settings(...) 会重解析整套配置，并把未显式给出的 max_split_size_mb、roundup_power2_divisions、garbage_collection_threshold 重置为默认值（PyTorch 2.10 的 AcceleratorAllocatorConfig::parseArgs 明确如此）。当前 guard 先传入仅含 expandable_segments:False 的字符串，退出时又仅传入 expandable_segments:True，所以例如 max_split_size_mb:512,garbage_collection_threshold:0.8,expandable_segments:True 在一次 capture 后会永久丢失前两项。我在本机 PyTorch 2.9.1/CUDA 上通过 torch.cuda.memory._snapshot()["allocator_settings"] 复现：max_split_size 从 512 MiB 变为 -1，GC threshold 从 0.8 变为 0.0，恢复 True 后仍未恢复。新增测试虽然使用了混合配置，但只检查调用字符串，因而没有覆盖这个回归。请用保留其余选项的完整配置切换，并在退出时恢复完整原配置，同时补充对应测试。

2. 仓库 CUDA 依赖已锁定 PyTorch 2.10.0；该版本也支持新的统一变量 PYTORCH_ALLOC_CONF。当用户只通过该变量启用 expandable_segments:True 时，当前代码只读取 PYTORCH_CUDA_ALLOC_CONF，guard 会直接 no-op，原始 CUDA IPC 失败仍然存在。请按 PyTorch 2.10 的实际环境变量选择/优先级兼容统一变量并补测试。

现有新增测试：6 passed；Ruff 0.14.0 check/format 均通过。上面两项是测试未覆盖的运行时状态问题。PR 仍为 Draft，未做合并。

[Leonccaa]

@yangzhuxinyzx 谢谢仔细复现，两个阻断项都已在 2fc729478 修复。

• 现在传给 allocator setter 的是完整的生效配置，只在 capture/register 窗口将 expandable_segments:True 替换为 False，退出时传回完整原字符串。max_split_size_mb、roundup_power2_divisions 和 garbage_collection_threshold 不再丢失。
• 环境变量选择已按 PyTorch 2.10 实际优先级实现：只要 PYTORCH_CUDA_ALLOC_CONF 存在（包括空值）就优先，否则回落 PYTORCH_ALLOC_CONF。
• focused pytest 现为 12 passed；changed-file hooks 和 git diff --check 全部通过。新测试包含一个模拟 PyTorch parseArgs 重置行为的状态测试，不再只检查 setter 调用字符串。
• 另用 torch==2.10.0+cu128 和 2×V100 完成了真实验证：legacy mixed config 与 unified-only config 的 allocator snapshot 均为前后完全相等；两种变量配置下 TP2 CUSTOM CUDA Graph 注册/replay 均在两个 rank 通过，输出正确且 allocator 完整恢复。

PR 保持 Draft，请重新审阅。当前 main 中 cumem.py 和 _scoped_allocator_max_split 的同类 partial-setting 模式不混入本 PR，会单独登记和处理。

[yangzhuxinyzx]

复核当前 head 2fc729478f2b0d6a5ec7a21efa42ada560d20a4f：上一轮两个 allocator correctness 阻断项均已修复。

• guard 现在基于完整的生效配置仅替换 expandable_segments:True，退出时恢复完整原配置；新增状态型测试能够捕获 PyTorch parseArgs 对缺失字段的重置。
• legacy/unified 变量选择与仓库锁定的 PyTorch 2.10 实际行为一致：PYTORCH_CUDA_ALLOC_CONF 只要存在（包括空值）即优先，否则使用 PYTORCH_ALLOC_CONF。
• 我在本机 torch==2.10.0+cu128 的真实 CUDA allocator snapshot 上独立验证了 mixed config：expandable_segments 为 True -> False -> True，max_split_size_mb、roundup_power2_divisions、garbage_collection_threshold 在上下文前后逐项完全一致。
• 聚焦测试 12 passed；Ruff check/format、SPDX、forbidden-import 和 git diff --check 均通过。

当前剩余的合并门禁是 DCO：PR 的两个提交 3fce76cb8、2fc729478 都没有 Signed-off-by: trailer，而本仓库明确要求 git commit -s。请由贡献者补齐相应 sign-off 后再标记 Ready；届时只需按新 head 复核提交元数据和小型测试。

非阻断建议：当前调用 torch.cuda.memory._set_allocator_settings 会在 PyTorch 2.10 每个 worker 打出两条 FutureWarning；可改用该版本已有的 torch._C._accelerator_setAllocatorSettings（必要时保留兼容 fallback）以避免启动日志噪声。

PR 仍为 Draft，本轮不合并。

[Leonccaa]

@yangzhuxinyzx 已按仓库 DCO 要求补齐两个提交的 Signed-off-by：

• ef0cad3：初始 backport 提交
• 08c7ec5：allocator correctness 修复提交

这次仅重写 commit metadata；新旧 head 的 tree hash 均为 8354d694987b216fb8f94bcbcc80ea261fe9f565，代码树完全一致。新 head 上聚焦测试重新运行：12 passed。

PR 继续保持 Draft，等待维护者复核并决定何时标记 Ready。FutureWarning 建议本轮不扩入，避免在 correctness 与 DCO 已收口后增加新的代码差异。

[yangzhuxinyzx]

Re-review completed on exact head 08c7ec5bfa4ecf7e0c70801421ea263be6f0a58f against current main@6c7a8617f.

• Both rewritten commits now carry a Signed-off-by trailer matching the commit author.
• The source tree is byte-for-byte identical to the previously audited and allocator-validated 2fc729478 tree.
• No-commit merge onto current main is clean; git diff --check passes.
• Focused allocator tests: 12 passed.
• Ruff check and format: passed.

The prior allocator-correctness and DCO change requests are resolved, so I approve this exact source head. The PR remains Draft; this approval does not authorize changing it to Ready or merging it before the author completes the Draft workflow.

[yangzhuxinyzx]

Current-main revalidation on exact head 08c7ec5bfa4ecf7e0c70801421ea263be6f0a58f: the prior APPROVED review remains valid.

The two commits are DCO-signed; git diff --check, Ruff 0.14 check/format, SPDX, forbidden-import, and torch-accelerator checks pass; all 12 allocator tests pass; and the merge-tree is clean against main@66becb6423b6. The hosted pre-run-check failed only because the PR lacks the repository verified/ready workflow label and the author count threshold, so pre-commit was skipped.

No source blocker is open. The remaining blockers are workflow state: this PR is still Draft and required CI has not run on the approved head. A maintainer should move it through the normal Ready/verified workflow; once it is non-Draft with green required checks, it is the current merge candidate.