Prhub

#56015 [XPU] Honor device_ids for worker device placement

原始 PR 作者 chaojun-zhang 合并时间 2026-09-16 11:53 文件变更 3 提交数 4 评论 12 代码增减 +44 / -19

执行摘要

修复 XPU worker 忽略 device_ids 的设备放置

PR body 给出了非常直接的动机:「XPUWorker.init_device directly binds torch.device(f"xpu:{local_rank}"), ignoring the engine's resolved device_ids (parallel_config.assigned_physical_gpu_ids). As a result, --device-ids (or LLM(device_ids=...)) had no effect on XPU: every single-GPU worker landed on visible ordinal 0, regardless of the requested device.」也就是说,这不是「设备选错」的偶发 bug,而是整条 device_ids 契约在 XPU 上从未接线:engine 侧不解析,worker 侧不消费。修复策略也写明了要对齐 CUDA:「This mirrors the CUDA worker changes from #45026」,把 XPU 从特例分支收回到与 CUDA 相同的代码路径,避免平台间行为漂移。

值得精读,尤其适合关注平台差异收敛、多后端设备绑定与 CI 依赖治理的工程师。

阅读建议:

  • 先看 vllm/v1/worker/xpu_worker.py 的 init_device,重点是「映射存在 / 不存在」两条分支与 logical_device_id_to_visible_device_id() 的引入,这是理解 vLLM 里 logical rank 与 visible ordinal 两套编号体系的好样本。
  • 再看 vllm/v1/engine/core.py 的 _set_visible_devices 与 _set_assigned_physical_gpu_ids,理解 device_ids 是如何从 CLI 一路解析并写回 VllmConfig 的(这条链路也直接决定 Ray 部署下 DP rank 的设备分配)。
  • 特别关注 review 中 jikunshang 对 dpctl 的质疑与作者的复现数据,这是「依赖一个非直觉的外部库副作用来修正确性」的典型案例。

可跟进的改进(若团队在 XPU + Ray 生产部署):补一个覆盖 assigned_physical_gpu_ids 非 None 的单元测试(mock torch.xpu.get_device_properties 与 logical_device_id_to_visible_device_id 即可低成覆盖断言分支);把 dpctl 依赖写入文档或环境检查;评估是否把映射发布挪到断言之后以消除全局状态污染。

讨论亮点

核心交锋有两处,且都直接决定合并版本长什么样:

  1. 断言与平台全局状态副作用(hlin99,vllm/v1/worker/xpu_worker.py:132):hlin99 提出「is it better to move set_assigned_physical_gpu_ids after assert? so the mapping is validated before processing」,即先校验映射再发布,避免校验失败时平台全局状态已被改写。作者以一致性为由回绝:「the CUDA gpu_worker.py follows the same order today: it calls set_assigned_physical_gpu_ids(...) first and then runs the bounds asserts right after. I kept this PR's XPU path consistent with that existing pattern rather than diverging from it.」结论:保持原顺序,评论被定性为 nit,不阻塞合并。

  2. CI 安装 dpctl 的争议(jikunshang,.buildkite/intel_jobs/test-intel.yaml:100):jikunshang 明确反对:「I don't want to install this.. I know ray is using this. cc @faaany I remember there are some compatible issue right?」作者用最小复现回应(Ray 2.58、8x Intel Arc Pro B60、ZE_AFFINITY_MASK=2):不装 dpctl 得到 {'accelerator_ids': ['0'], 'ZE_AFFINITY_MASK': '2'}(MISMATCH,复现 bug),装上后得到 {'accelerator_ids': ['2'], 'ZE_AFFINITY_MASK': '2'}(一致)。结论:以数据说服评审,该行进入合并版本。

  3. 与「移除 ray」计划的对齐(jikunshang):「#56013 plan to remove ray. @yma11 can you algin first?」,yma11 回复「I will put ray in test requirements, should not conflict with this PR」,确认无冲突后 jikunshang 执行 /ci run。

另外,claude[bot] 的自动审查对 fork PR 被禁用(「This pull request is from a fork — automated review is disabled」),因此本次改动没有经历自动化代码审查,仅靠两位人类 reviewer 的 APPROVED 与 Intel CI 覆盖。

实现拆解

变更入口:问题横跨「engine 侧解析映射」与「worker 侧消费映射」两端,两端都需要改,因此 PR 同时改了 engine core 与 XPU worker,并在 CI 里补一个 Ray 侧的前提依赖。

  1. 源头短路修复(vllm/v1/engine/core.py):_set_visible_devices 中原本是 if current_platform.is_xpu(): pass else: ...,XPU 的 DP engine-core actor 因此从不调用 _set_assigned_physical_gpu_ids,get_physical_gpu_ids_for_local_dp_rank(device_control_env_var, local_dp_rank, world_size, user_assigned_gpu_ids=...) 的解析结果也就不会被写回 vllm_config.parallel_config.assigned_physical_gpu_ids。删除该分支后 XPU 与 CUDA 共用一个实现,后续 worker 才有非 None 的映射可读;这也是本 PR 唯一影响共享代码路径的改动,对非 XPU 平台语义等价。

  2. worker 侧消费映射并重做边界校验(vllm/v1/worker/xpu_worker.py):在 XPUWorker.init_device 中新增分支——若 assigned_physical_gpu_ids 非 None,先调用 set_assigned_physical_gpu_ids() 把逻辑→物理映射发布到平台层,再把原来的 self.local_rank < visible_device_count 断言换成 self.local_rank < len(assigned_physical_gpu_ids);否则回退到原来的 torch.accelerator.device_count() 断言。这样「校验对象」与「实际使用的映射表」保持一致,同时保留非 DP / 无映射场景的旧语义。

  3. 跨节点后端的校验豁免:新增的 local_world_size <= len(assigned_physical_gpu_ids) 断言只在 distributed_executor_backend 不属于 ray/external_launcher 时执行。代码注释解释了原因:local_world_size 由 parallel_config.nnodes 推导,而 nnodes 只在 mp 多节点后端下被设置,ray/external_launcher 下恒为 1,会让 local_world_size 退化成整个 world_size,跨节点部署时误报。取舍是这两类后端只信任 local_rank 上界。

  4. 逻辑序号到可见序号的翻译:设备选择从 torch.device(f"xpu:{self.local_rank}") 改为先取 current_platform.logical_device_id_to_visible_device_id(self.local_rank),再构造 torch.device、调用 torch.accelerator.set_device_index(),并且 torch.xpu.get_device_properties() 的入参与 init_gpu_memory 统计同样改用该可见序号;而写回环境变量的 LOCAL_RANK、init_worker_distributed_environment() 与 oneCCL warm-up 仍沿用逻辑 local_rank,保持进程组与集合通信的既有语义不变。

  5. CI 配套与验证(.buildkite/intel_jobs/test-intel.yaml):xpu-v1-test 步骤的 shell 从 cd tests && pytest ... 变成 cd tests && pip install dpctl && pytest ...。原因是作者的复现显示:在 ZE_AFFINITY_MASK=2 的受限容器里,不装 dpctl 时 Ray 报 accelerator_ids=['0'],装上后才报 ['2'],即 Ray 在 Intel GPU 上的 accelerator id 枚举依赖 dpctl,这正是 --device-ids 端到端成立的前提。

文件 角色 关键动作
vllm/v1/engine/core.py 源头 删除 XPU 空分支,DP actor 统一写 assigned_physical_gpu_ids
vllm/v1/worker/xpu_worker.py 消费端 发布映射、改断言、逻辑→可见序号翻译
.buildkite/intel_jobs/test-intel.yaml 配套 xpu-v1-test 安装 dpctl

测试配套:本 PR 未新增任何单测;材料显示源码文件没有对应测试文件变更,验证依赖既有 v1/core、v1/engine、v1/worker 等 Intel CI 套件,以及作者在 review 中贴出的独立 Ray 复现脚本。需要说明的是,提交 ac546a9 的提交信息还提到为 ray_executor_v2.py/interface.py/xpu.py 增加 Platform.ray_accelerator_id_to_physical_device_id() 覆写,但最终变更文件只有 3 个,说明那条平台钩子路径在合并前被放弃或回退,改用 CI 安装 dpctl 的方式规避(此处以最终 diff 为准,属于推断)。

文件 模块 状态 重要度
vllm/v1/worker/xpu_worker.py 设备放置 modified 7.19
vllm/v1/engine/core.py 引擎核心 modified 6.15
.buildkite/intel_jobs/test-intel.yaml 测试流水线 modified 2.9

关键符号

XPUWorker.init_device _set_visible_devices _set_assigned_physical_gpu_ids set_assigned_physical_gpu_ids logical_device_id_to_visible_device_id get_physical_gpu_ids_for_local_dp_rank

关键源码片段

vllm/v1/worker/xpu_worker.py dependency-wiring

修复主体:XPUWorker.init_device 从「直接用 local_rank 当设备序号」改为「消费 assigned_physical_gpu_ids + 逻辑→可见序号翻译」,并重做边界断言与跨节点后端豁免,是 --device-ids 在 XPU 生效的关键。

def init_device(self):
    # 前段:依据 data_parallel_index / replica_world_size 推导本 engine 在节点内的 DP 槽位,
    # 叠加到 self.local_rank(本次改动未触及)
    self.local_rank += dp_local_rank * replica_world_size
​
    device = self.device_config.device
    if (
        isinstance(device, torch.device)
        and device.type == "xpu"
        and current_platform.is_xpu()
    ):
        # 关键改动:engine 侧解析出的物理设备 ID 列表若存在,就发布到平台层,
        # 后续 logical_device_id_to_visible_device_id() 才能把逻辑 rank 换成可见序号;
        # 改动前这里直接 self.device = torch.device(f"xpu:{self.local_rank}"),
        # --device-ids 因此完全失效
        assigned_physical_gpu_ids = self.parallel_config.assigned_physical_gpu_ids
        if assigned_physical_gpu_ids is not None:
            from vllm.platforms.interface import set_assigned_physical_gpu_ids
​
            set_assigned_physical_gpu_ids(assigned_physical_gpu_ids)
            # 校验对象从「可见设备数」换成「映射表长度」:
            # 这是 device_ids 生效后必须满足的新不变量
            assert self.local_rank < len(assigned_physical_gpu_ids), (
                f"local_rank {self.local_rank} is out of bounds for "
                f"assigned_physical_gpu_ids {assigned_physical_gpu_ids}"
            )
            # ray / external_launcher 下 parallel_config.nnodes 恒为 1,
            # local_world_size 会退化成整个 world_size,跨节点部署时误报,
            # 因此这两类后端只信任上面的 local_rank 上界
            if parallel_config.distributed_executor_backend not in (
                "ray",
                "external_launcher",
            ):
                assert parallel_config.local_world_size <= len(
                    assigned_physical_gpu_ids
                ), (
                    f"local_world_size ({parallel_config.local_world_size})"
                    " exceeds assigned_physical_gpu_ids count "
                    f"({len(assigned_physical_gpu_ids)})"
                )
        else:
            # 无映射场景(非 DP、非 Ray 等)保持旧语义:按可见设备数校验
            assert self.local_rank < torch.accelerator.device_count(), (
                f"DP adjusted local rank {self.local_rank} is out of "
                f"bounds for {torch.accelerator.device_count()} devices."
            )
​
        # 逻辑 rank -> 可见序号;设备选择、set_device_index 与显存统计必须全部用同一序号,
        # 否则统计到的是另一张卡
        visible_device_index = (
            current_platform.logical_device_id_to_visible_device_id(self.local_rank)
        )
        self.device = torch.device(f"xpu:{visible_device_index}")
        torch.accelerator.set_device_index(self.device)
        current_platform.check_if_supports_dtype(self.model_config.dtype)
        torch.accelerator.empty_cache()
        self.init_gpu_memory = torch.xpu.get_device_properties(
            visible_device_index
        ).total_memory
    else:
        raise RuntimeError(f"Unsupported device type: {self.device_config.device}")
​
    # 分布式相关仍沿用逻辑 local_rank:进程组与 oneCCL warm-up 语义不变
    os.environ["LOCAL_RANK"] = str(self.local_rank)
    init_worker_distributed_environment(
        self.vllm_config,
        self.rank,
        self.distributed_init_method,
        self.local_rank,
        current_platform.dist_backend,
    )
vllm/v1/engine/core.py core-logic

链路源头:删除 `_set_visible_devices` 中 XPU 专属的 `pass` 空分支,使 XPU 的 DP engine-core actor 与 CUDA 一样把解析结果写回 parallel_config.assigned_physical_gpu_ids,是 worker 侧新逻辑能被触发的前提。

def _set_visible_devices(self, vllm_config: VllmConfig, local_dp_rank: int):
    from vllm.platforms import current_platform
​
    # 改动前这里是 `if current_platform.is_xpu(): pass` 的短路分支:
    # XPU 的 DP engine-core actor 从不写 assigned_physical_gpu_ids,
    # worker 侧只能读到 None,--device-ids / LLM(device_ids=...) 全部失效。
    # 现在 XPU 与 CUDA 走同一条路径
    device_control_env_var = current_platform.device_control_env_var
    self._set_assigned_physical_gpu_ids(
        vllm_config, local_dp_rank, device_control_env_var
    )
​
​
def _set_assigned_physical_gpu_ids(
    self,
    vllm_config: VllmConfig,
    local_dp_rank: int,
    device_control_env_var: str,
):
    # 本方法未在本次改动中修改,列出以说明写回路径:
    # 按 local_dp_rank 切出本 engine 负责的那一段物理设备 ID
    world_size = vllm_config.parallel_config.world_size
    try:
        physical_gpu_ids = get_physical_gpu_ids_for_local_dp_rank(
            device_control_env_var,
            local_dp_rank,
            world_size,
            user_assigned_gpu_ids=(
                vllm_config.parallel_config.assigned_physical_gpu_ids
            ),
        )
        # 写回同一个 VllmConfig 对象,随后被 worker 读取
        vllm_config.parallel_config.assigned_physical_gpu_ids = physical_gpu_ids
    except IndexError as e:
        raise Exception(
            f"Error computing assigned_physical_gpu_ids: "
            f"local range: [{local_dp_rank * world_size}, "
            f"{(local_dp_rank + 1) * world_size}) "
            f'base value: "{os.getenv(device_control_env_var)}"'
        ) from e

评论区精华

set_assigned_physical_gpu_ids 是否应放在断言之后 设计

hlin99 在 vllm/v1/worker/xpu_worker.py:132 提问「is it better to move set_assigned_physical_gpu_ids after assert? so the mapping is validated before processing」;作者回复 CUDA 的 gpu_worker.py 当前也是先 publish 再 assert,为保持两条路径一致不改顺序。

结论:维持原顺序,以与 CUDA 路径一致为优先;被定性为 nit,不阻塞合并。副作用是断言失败时平台全局映射已被改写。 · resolved-as-nit

Intel CI 是否应该安装 dpctl 测试

jikunshang:「I don't want to install this.. I know ray is using this. cc @faaany I remember there are some compatible issue right?」作者给出最小复现:Ray 2.58、8x Intel Arc Pro B60、ZE_AFFINITY_MASK=2 时,不装 dpctl 得到 accelerator_ids ['0'](MISMATCH),装上得到 ['2'](一致)。

结论:以复现数据说服评审,该行进入合并版本;但生产环境是否安装 dpctl 成为隐性前提,Ray 部署的设备 ID 正确性依然依赖外部库。 · resolved-with-residual-concern

跨节点后端跳过 local_world_size 校验 设计

代码注释说明 local_world_size 由 parallel_config.nnodes 推导,nnodes 只在 mp 多节点后端设置,ray/external_launcher 下恒为 1,会使 local_world_size 退化成整个 world_size 并在跨节点部署误报,因此这两类后端只保留 local_rank 上界校验。

结论:有意弱化 local_world_size 的交叉校验,改由 local_rank < len(assigned_physical_gpu_ids) 单独保证映射有效,误配置更晚暴露。 · documented-in-code

与「移除 ray」计划的对齐 other

jikunshang 提示「#56013 plan to remove ray. @yma11 can you algin first?」;yma11 回复「I will put ray in test requirements, should not conflict with this PR」,随后 jikunshang 执行 /ci run。

结论:确认不冲突,本 PR 继续合并;但 ray 在 XPU 设备绑定中的角色长期仍需重新评估。 · 已解决

风险与影响

  1. 平台全局状态在断言前被改写(设计取舍,未修复):set_assigned_physical_gpu_ids() 先于 assert self.local_rank < len(assigned_physical_gpu_ids) 执行,一旦断言失败,进程内的逻辑→物理映射已被污染,异常路径上的后续逻辑(或同进程内的重试)可能读到不一致状态。这与 hlin99 的 nit 是同一问题,最终以「与 CUDA 一致」保留,属于有意接受的隐患。

  2. Ray 路径的正确性依赖外部 dpctl:作者自己的复现证明,在没有 dpctl 的环境里 Ray 会把受限容器(ZE_AFFINITY_MASK=2)报成 accelerator_ids=['0']。CI 里装了 dpctl,但生产用户的 XPU + Ray 部署不一定会装;而提交信息中提到过的平台钩子 Platform.ray_accelerator_id_to_physical_device_id() 并不在最终 3 个变更文件里,意味着 Ray 部署下的物理 ID 恢复仍可能错位。这是本 PR 最需要在真实集群验证的残余风险。

  3. 无新增单测,分支覆盖薄弱:vllm/v1/worker/xpu_worker.py 新增的两条分支(有映射 / 无映射)、跨节点后端的豁免逻辑、以及 logical_device_id_to_visible_device_id 的翻译结果都没有针对性测试;而 CI 里 v1/worker 套件本身还 --ignore=v1/worker/test_worker_memory_snapshot.py,覆盖进一步打折。

  4. 行为变更从「静默」转为「启动失败」:以前所有单卡 XPU worker 都落在 0 号设备,配置错误只是被静默吸收;新逻辑一旦映射越界会直接 assert 失败。对依赖旧行为的部署(例如容器里只暴露一张卡但声明的 device_ids 与实际可见设备不一致)会表现为启动报错,属于预期内的收紧,但需要发布说明。

  5. 跨模块影响有限但真实:_set_visible_devices 是 CUDA/XPU 共用的方法,删除 XPU 分支后 XPU 会走 current_platform.device_control_env_var + get_physical_gpu_ids_for_local_dp_rank(),若 XPU 平台上的环境变量语义(ZE_AFFINITY_MASK)与该解析函数的假设不完全匹配,会抛 IndexError 并被包装成通用 Exception;非 XPU 平台的代码路径逐行等价,无回归风险。

用户侧:XPU 用户终于可以真正使用 --device-ids 或 LLM(device_ids=...) 指定物理设备,此前这些参数在 Intel GPU 上完全无效;DP + Ray 的多卡部署能够落到请求的设备上,避免了多 worker 争抢 0 号卡的隐性串行化。代价是错误配置会以启动断言失败的形式暴露,而不是被静默吸收。

系统侧:改动集中在 XPU 设备绑定链路(engine core 的 visible-device 设置与 worker 的 init_device),不触及调度、KV cache、attention 等核心路径;CUDA/ROCm 路径语义不变。真正的系统级影响在 CI:Intel CI 的 XPU V1 测试环境新增 dpctl 依赖,改变了 Ray 的 accelerator 枚举行为,后续 Intel job 的镜像与依赖管理需要同步。

团队侧:这是一个典型的「平台特例分支收敛」改动,把 XPU 从 pass 短路拉回与 CUDA 同构的实现,降低了后续维护两套设备绑定逻辑的成本;同时也暴露了 vLLM 在 Ray + 非 NVIDIA 加速器上依赖第三方库来获得正确设备 ID 的结构性问题——这与 jikunshang 提到的 PR#56013「移除 ray」方向一致,团队需要决定 Ray 在 XPU 场景的长期定位。

影响程度:单平台、单链路的中等影响变更,对绝大多数只跑 CUDA 的用户不可见。

缺少单测覆盖 共享 v1 引擎路径变更 CI 新增 dpctl 依赖 Ray/XPU 设备枚举依赖外部库 XPU 行为由静默转为启动失败

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论