执行摘要

• 一句话：升级 flash-linear-attention 到 0.5.2，修复 KDA 训练崩溃
• 推荐动作：该 PR 值得精读，尤其是对于维护或使用 fla 的工程师。原因：
  • 它展示了一个典型依赖升级的处理方式：清晰的动机说明、明确的 A/B 验证、对 API 兼容性的确认。
  • 它揭示了 triton 版本变化对 jit 代码的破坏性影响，对编写 triton 内核的开发者有警示意义。
  • 对 KDA 训练相关功能而言，这是一个必要的修复，建议合并后尽快使用新镜像验证训练流程。

功能与动机

PR body 说明：fla 0.4.2 中的 chunk_intra_token_parallel 在 @triton.jit 内调用 triton.next_power_of_2，被新版 triton 拒绝（"Unsupported function referenced"），导致 KDA 训练（glm5_next）前向失败。0.4.2 是最后一个存在此 bug 的版本，当前每个 KDA 训练部署都需要手动修补 site-packages。升级到 0.5.2 可消除手工修补，并恢复训练功能。

实现拆解

该 PR 的实现非常直接，仅涉及两个 Dockerfile 中的依赖版本号变更：

1. docker/Dockerfile：将 flash-linear-attention==0.4.2 改为 flash-linear-attention==0.5.2，位于 pip 安装部分，紧邻 mbridge 安装和 tilelang 安装之间。
2. docker/Dockerfile.rocm：同样将版本号从 0.4.2 改为 0.5.2，位于 mooncake 安装之后、tile_kernels 相关注释之前。

这两处变更分别影响 CUDA/通用镜像和 ROCm 镜像，确保两种硬件平台都能获取修复后的版本。无需修改任何代码，因为 0.5.2 的 API 与 0.4.2 兼容（PR body 中明确说明 fla.ops.cp.build_cp_context、cp_context kwargs、FusedRMSNormGated/ShortConvolution、lower_bound 均未改变），因此没有源码或测试配套改动。

关键文件：

• docker/Dockerfile（模块 部署脚本；类别 infra；类型 infrastructure）: 这是主要变更文件，版本号升级直接影响 CUDA/通用镜像的构建，是 KDA 训练功能恢复的关键。
• docker/Dockerfile.rocm（模块 部署脚本；类别 infra；类型 infrastructure）: 该文件用于构建 ROCm 镜像，版本升级确保 AMD 平台也能获得修复。

关键符号：未识别

评论区精华

该 PR 的 review 评论较少，主要讨论集中在 PR body 中。

• 作者在 PR body 中详细说明了升级动机（triton 兼容性 bug）和验证结果（A/B 对比，0.5.2 与 0.4.2 在 KDA 前向逐位一致，GDN 前向仅 bf16 舍入级差异）。
• yueming-yuan 给予了 APPROVED，无额外评论，表明审核者认可该变更。
• claude[bot] 的评论是自动提示，非实质性讨论。

没有发现未解决的疑虑或争议点。

• 暂无高价值评论线程

风险与影响

• 风险：该 PR 的风险主要在于：
  1. GDN 前向差异：PR body 明确说明 GDN（gate delta rule）前向与 0.4.2 相比有 bf16 舍入级差异（最大 9.8e-4），这可能影响 qwen3-next 系列模型的 logprob 指标，但仍在噪声范围内。
  2. 依赖变化：虽然 API 未变，但 0.5.x 在 PyPI 上是一个封装包，实际代码在 fla-core 依赖中，升级后需确保 fla-core 正确安装，否则可能出现运行时报错。
  3. ROCm 兼容性：Dockerfile.rocm 的变更未提供针对 ROCm 的专门验证，可能存在 ROCm 下的差异，但 PR body 中的 A/B 测试未明确说明是否覆盖 ROCm。

这些风险由于版本升级验证充分，总体可控，但建议关注 CI 结果。

• 影响：该 PR 的影响范围主要为使用这两个 Dockerfile 构建镜像的部署环境，特别是涉及 KDA 训练的 glm5_next 和 qwen3-next 系列模型。升级后：
  • 消除了手动修补 site-packages 的必要，降低部署复杂度。
  • KDA 训练前向恢复可用，梯度和前向输出与旧版基本一致。
  • GDN 相关模型的 logprob 指标可能出现微小漂移，但预计在噪声范围内。

对于团队而言，需要重新构建镜像并验证相关模型的训练和推理流程；对于依赖该镜像的下游用户，可能需要同步更新镜像版本。

• 风险标记：依赖升级

关联脉络

• PR #2714 Bump sglang to v0.5.18: 同为依赖版本升级的 PR，涉及 Dockerfile 和 CI 调整，可参考其升级流程和验证方式。