Prhub

#51464 [ROCm] update triton in base docker for gluon compatibility

原始 PR 作者 hongxiayang 合并时间 2026-08-13 00:26 文件变更 1 提交数 5 评论 12 代码增减 +1 / -1

执行摘要

升级 ROCm 基础镜像 Triton 分支,修复 Gluon MLA 编译

PR body 明确说明升级目的:选取 ROCm/triton PR#960 修复 gluon MLA kernel 编译问题——旧 Triton commit 拒绝 DistributedLinearLayout,导致 vllm-openai-rocm:nightly 首次 MLA decode 时在 aiter/ops/triton/gluon/mla_gluon.pytriton.compiler.errors.CompilationError,提示期望 offsets type layout 为 BlockedLayoutSliceLayout。升级后 gluon kernel 可以编译并服务 K3,且保持准确率。

该 PR 属于基础设施依赖更新,值得快速浏览以了解 vLLM 如何通过 pin Triton commit 解除内核编译阻塞,但无需精读。值得关注的两个设计点:一是外部依赖固定到具体 commit 而非 tag 的做法,能精确复现修复但也带来漂移风险;二是验证策略采用 aiter op 测试 + K3 GSM8K 端到端组合,可作为后续 ROCm 内核兼容性修复的参考模板。若团队维护 ROCm 镜像,建议跟进补充自动化的 Triton 内核冒烟测试以覆盖镜像构建路径。

讨论亮点

讨论集中在 CI 覆盖与合并节奏,没有技术争议。tjtanaa 请求 AndreasKaratzas 与 Rohan138 审查;Rohan138 打上 ready 标签以触发 AMD CI,并明确表示 Added ready to trigger amd-ci as well, but LGTM;AndreasKaratzas 在 AMD CI 通过后选择 rebase 触发 upstream CI 并合并,原话是 we know that amd ci is good after the above run (so happy), so i'm rebasing to just trigger upstream ci and get this merged;AMD 侧 reviewer dllehr-amd 批准并评价 Thanks! Looks good!

实现拆解

  1. 根因定位:旧 Triton commit 0263a6agl.amd.cdna4.async_copy.buffer_load_to_shared 的 layout 校验只接受 BlockedLayoutSliceLayout,而 Gluon MLA 内核传入的是 DistributedLinearLayout,因此首次 MLA decode 即编译失败。
  2. 单行配置更新:修改 docker/Dockerfile.rocm_base,将 ARG TRITON_BRANCH"0263a6a" 更新为 "532137f"(release/internal/3.7.x 分支 08/07 快照),该版本包含 ROCm/triton PR#960 的修复,这是 PR 唯一的代码变更。
  3. 验证:作者先运行 aiter 的 op 测试确认内核可编译,再在 vLLM nightly docker 中重建 Triton,用 Kimi-K3 无 spec 服务验证 gluon-MLA 路径启动成功、无编译错误与内存故障,GSM8K 全量(1319 问)准确率 94.5%。
  4. CI 与合并流程:打上 ready 标签触发 AMD CI 验证 ROCm 路径,通过后 rebase 到最新 main 触发 upstream CI(Buildkite #83431/#83436/#83450),最终合并。没有新增自动化测试或文档改动。
文件 模块 状态 重要度
docker/Dockerfile.rocm_base 基础镜像 modified 2.78

关键源码片段

docker/Dockerfile.rocm_base configuration

唯一变更文件,通过修改 Triton 分支固定版本修复 Gluon MLA 内核在 ROCm 上的编译失败,是 vLLM ROCm 基础镜像的构建入口配置。

# ROCm 基础镜像关键依赖分支固定(构建入口)
ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.3-complete
# 将 Triton 分支从 0263a6a 升级到 532137f(release/internal/3.7.x, 08/07)
# 目的是携带 ROCm/triton PR#960 的修复:解决 Gluon MLA kernel 中
# DistributedLinearLayout 被 async_copy 拒绝导致的 CompilationError
# 此前 vllm-openai-rocm:nightly 在 Kimi-K3 首次 MLA decode 时编译失败
ARG TRITON_BRANCH="532137f" # release/internal/3.7.x as of 08/07
ARG TRITON_REPO="https://github.com/ROCm/triton.git"
ARG PYTORCH_BRANCH="6bbd260" # release/2.12 as of 08/02
ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git"

评论区精华

请求维护者审查(PTAL) question

tjtanaa 在 PR 打开后第一时间 Tag 两位维护者 `@AndreasKaratzas @Rohan138` 请求审查,并触发 `/ci run`。

结论:维护者随后介入:Rohan138 检查后表示 LGTM,AndreasKaratzas 负责触发 CI 并最终合并。 · 已解决

触发 AMD CI 以覆盖 ROCm 验证 测试

Rohan138 补充说明 `Added ready to trigger amd-ci as well, but LGTM`,即打上 `ready` 标签以触发 AMD CI,同时表达认可。

结论:AMD CI 通过,验证了 Triton 升级在 ROCm 上无回归。 · 已解决

Rebase 策略与 upstream CI 触发 other

AndreasKaratzas 评论:`we know that amd ci is good after the above run (so happy), so i'm rebasing to just trigger upstream ci and get this merged`,即先确认 AMD CI 通过,再 rebase 触发 upstream CI,避免重复跑完整流水线。

结论:连续 rebase 触发 Buildkite #83431/#83436/#83450 后合并。 · 已解决

风险与影响

主要风险是外部依赖版本升级:532137f 是 ROCm Triton 内部 release/internal/3.7.x 分支的快照而非稳定 tag,升级同时影响镜像中所有依赖 Triton 的内核路径(Gluon MLA、各类 quantize/attention kernel),存在编译行为变化的可能。其次,PR 未为镜像构建新增自动化测试,回归验证依赖人工的 aiter op 测试和 K3 GSM8K 端到端结果,若上游 CI 未覆盖 ROCm Triton 内核编译,类似问题可能在合入后暴露。风险面限定在 docker/Dockerfile.rocm_base,非 ROCm 平台不受影响。

对用户:使用 vllm-openai-rocm:nightly 部署 Kimi-K3 的用户从启动即失败变为可正常服务,GSM8K 94.5% 准确率说明精度无回退;其它 ROCm 模型理论上受 Triton 版本变化影响,但验证中未发现异常。对系统:变更仅影响 Docker 构建阶段,不触及运行时 Python 代码,镜像的依赖快照会随新分支更新而更新,可能影响构建缓存复用。对团队:为后续 ROCm K3/Gluon 内核支持建立了可工作的 Triton 基线,但需要关注该内部分支与上游的同步节奏。

外部依赖升级 影响 ROCm 镜像构建 缺少自动化测试覆盖

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论