Skip to content

issue/1553 - fix: let the default runtime honour INFINICORE_DEFAULT_DEVICE_INDEX - #1554

Open
JoeZhang-0x000 wants to merge 1 commit into
InfiniTensor:mainfrom
JoeZhang-0x000:fix/default-runtime-device-index
Open

issue/1553 - fix: let the default runtime honour INFINICORE_DEFAULT_DEVICE_INDEX#1554
JoeZhang-0x000 wants to merge 1 commit into
InfiniTensor:mainfrom
JoeZhang-0x000:fix/default-runtime-device-index

Conversation

@JoeZhang-0x000

@JoeZhang-0x000 JoeZhang-0x000 commented Sep 7, 2026

Copy link
Copy Markdown

关联 issue: #1553

问题

ContextImpl 的构造函数把默认 runtime 的设备号硬编码为 0,所以任何使用非 0 号卡的进程,都会在 0 号卡上额外建一整套 runtime(primary context + 一条流 + 一个 infiniop handle + 两个分配器),尽管它从不在那张卡上计算。PyTorch 没有这个行为。

复现方式和实测数据见 #1553

改动

只动 src/infinicore/context/context_impl.cc,三处:

  1. 新增 default_device_index(),读环境变量 INFINICORE_DEFAULT_DEVICE_INDEX 决定默认 runtime 建在哪张卡上,构造函数用它取代硬编码的 0变量未设置、取值非法或越界时一律回落到 0,与现行行为逐字一致,因此对现有用法完全向后兼容。

  2. getCurrentRuntime() 的惰性初始化改为扫描整行找第一个非空 runtime,不再假定下标 0。这是第 1 处的必要配套:默认 runtime 不再必然位于 0,否则这里会读到 nullptr 并一路回落到 CPU runtime。

  3. setDevice() 不再为了做一次比较而调用 getCurrentRuntime(),改为直接判断 current_runtime_ 是否为空。这一处独立成立:在冷线程上,setDevice(cuda:N) 的第一行原本就会先把默认卡的 runtime 惰性建出来,仅仅是为了回答一个马上会被覆盖的比较 —— 这也是调用方「提前 set_device()」无法规避该问题的原因。

嵌入方只需在第一次调用 InfiniCore 之前,把本进程要用的卡号写进 INFINICORE_DEFAULT_DEVICE_INDEX

不新增除 <cstdlib> 外的依赖,不改任何公开 API,不动构建配置。

验证

环境:MetaX C550(8 卡)/ MACA 3.8.0.23 / PyTorch 2.10.0 / vLLM 0.22.0。

单进程隔离,绑定到 3 号卡:

场景 进程持有 context 的卡
修复前 0 和 3
修复后,设 INFINICORE_DEFAULT_DEVICE_INDEX=3 仅 3
修复后,不设该变量 0 和 3(与原行为一致)

端到端,通过 vLLM 张量并行,运行中用 mx-smi 统计 0 号卡上的进程数:

后端 TP 0 号卡上的进程 结果
vllm-metax(对照) 4 2 正常
InfiniCore 修复前 4 4 卡死
InfiniCore 修复后 4 2 正常

修复后 16 组张量并行用例(3 个模型 × TP 1~8 × 两种后端)全部有效;此前 TP>=4 的用例无一能跑完。

格式化:clang-format 21.1.8 对本文件 --dry-run --Werror 通过。

说明

MACA 驱动为何会因这些多余的 runtime 而超时,我们没有定论,本 PR 也不做任何机制上的声称。已排除的有:不是显存(多余的 runtime 占 0 MiB,把 gpu_memory_utilization 降到 0.5 照样挂)、不是每卡 context 数量(8 个进程同时在一张卡上持有 context 正常)、不是每卡流的容量(单进程建 4096 条流正常)。合成负载也复现不出该失败。

本 PR 依据的是因果证据:去掉这个多余的 runtime,TP>=4 就从必然卡死变成全部跑通。 而且无论驱动侧的上限是多少,一个进程都不应该在它从不使用的卡上建 runtime。

我们也考虑过在构造函数里用 infinirtGetDevice() 自动探测当前设备,这样就不需要环境变量。但它返回的是 InfiniRT 自己维护的 CURRENT_DEVICE_ID,而非驱动层的当前设备,在 InfiniRT 被显式 setDevice 之前拿不到宿主框架已经设好的卡号。要做到自动探测需要新增直达驱动的接口(各后端已有 hcGetDevice / cudaGetDevice 之类的调用),改动面较大。如果维护者更倾向那个方向,我们乐意按该思路重做。

ContextImpl built its default Runtime on device index 0 whatever card the
process actually uses, so any process working on another card also opened a
whole extra runtime on card 0 -- primary context, a stream, an infiniop handle
and two allocators -- on a card it never computes on. In a tensor-parallel job
that accumulates: measured with vLLM at TP=4, card 0 carried four processes
against two for the reference backend, and on MACA that made TP>=4 hang in an
unbounded mxkwCreateQueueBlock retry.

default_device_index() now reads INFINICORE_DEFAULT_DEVICE_INDEX so the
embedder can name the card this process will use. Unset, invalid or
out-of-range values all fall back to 0, matching the previous behaviour
exactly.

Two supporting changes: getCurrentRuntime() scans the row for the first
non-null runtime instead of assuming index 0, since the default one no longer
has to live there; and setDevice() no longer calls getCurrentRuntime() merely
to compare, because on a cold thread that eagerly builds the default card's
runtime just to answer a comparison it is about to overwrite -- which is why
calling setDevice() early could not avoid the stray runtime.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@JoeZhang-0x000
JoeZhang-0x000 requested a review from a team September 7, 2026 11:39
@JoeZhang-0x000

Copy link
Copy Markdown
Author

我手动检验过,详细见#1553 里面描述的问题,在沐曦机器上会出现错误。怀疑和沐曦驱动有关系。
最开始以为是primary context,stream,handler 这些实例在一张卡有上限,后来做了实验发现不是这么简单。
具体为什么导致的TP=4 出错不太清楚,但是 原本的做法中确实存在了隐患,在 card 0 中存的完全是一份副本。
而且改掉后原本报错的地方能跑通了。

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant