执行摘要

本次 PR 修复了 AMD ROCm QuickReduce 在 CUDA Graph 模式下由于 flag 状态被 graph 固定导致的累计精度错误。通过将 flag_color 从 launch 标量改为设备端指针，在 kernel 内部动态递增，确保了多次 replay 的正确隔离。同时新增手工回归测试。

功能与动机

QuickReduce 使用 flag color 来同步多个 GPU 的规约轮次。在 CUDA Graph 中，如果 flag_color 作为 kernel launch 参数，它会被固化在 graph 中，每次 replay 时值不变，导致后续轮次无法区分新旧 flag，从而读取错误数据。PR 正文给出详细复现代码和结果，显示了错误的累加值。

实现拆解

1. 问题定位：发现 kernel allreduce_prototype_twoshot 接收的 flag_color 标量在 graph capture 时被固定。
2. 参数改造：将参数类型从 uint32_t flag_color 改为 uint32_t* d_flag_counters，在 kernel 内部从设备内存读取并递增，写回。
3. 宏同步更新：修改 TWOSHOT_DISPATCH 宏，将 d_flag_counters 传入 kernel launch 函数。
4. 生命周期管理：在 DeviceComms 结构体中新增成员 d_flag_counters，在 init() 中分配并初始化为 1，在 destroy() 中独立释放（绕过 initialized 标志）。
5. 测试覆盖：在 test/manual/test_quick_allreduce.py 中新增 qr_graph_replay 函数和 TestQuickreduceGraphReplay 测试类，构造单个 graph 并多次 replay 的场景，验证每次 replay 结果正确。

sgl-kernel/csrc/allreduce/quick_all_reduce.h

核心修改：将 flag_color 由 launch 标量改为设备端指针，在 kernel 内部动态递增

#include <hip/hip_runtime.h>
#include <vector>
#include "quick_all_reduce.cuh"
​
namespace quickreduce {
​
template <typename AllReduceKernel, typename T>
__global__ __quickreduce_launch_bounds_two_shot__ static void allreduce_prototype_twoshot(
    T const* A,
    T* B,
    uint32_t N,
    uint32_t num_blocks,
    int rank,
    uint8_t** dbuffer_list,
    uint32_t data_offset,
    uint32_t* d_flag_counters, // 每个 block 的 color 计数器（设备端）
    int64_t data_size_per_phase) {
  int block = blockIdx.x;
  int grid = gridDim.x;
​
  // 从设备内存读取此 block 的 flag_color，使每次 graph replay 得到新值
  uint32_t flag_color = d_flag_counters[blockIdx.x];
​
  while (block < num_blocks) {
    AllReduceKernel::run(A, B, N, block, rank, dbuffer_list, data_offset, flag_color, data_size_per_phase);
    block += grid;
    flag_color++;
  }
  // 整个 block 最终值一致，由 thread (0,0) 写回
  if (threadIdx.x == 0 && threadIdx.y == 0) {
    d_flag_counters[blockIdx.x] = flag_color;
  }
}
​
#define TWOSHOT_DISPATCH(__codec) \
  if (world_size == 2) { \
    using LineCodec = __codec<T, 2>; \
    using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
    hipLaunchKernelGGL( \
        (allreduce_prototype_twoshot<AllReduceKernel, T>), \
        dim3(grid), \
        dim3(kBlockTwoShot), \
        0, \
        stream, \
        A, B, N, num_blocks, rank, dbuffer_list, data_offset, \
        d_flag_counters, \
        this->kMaxProblemSize); \
  } else if (world_size == 4) { \
    // ... 类似 ( 省略 ) \
  } else if (world_size == 8) { \
    // ... 类似 ( 省略 ) \
  }
​
struct DeviceComms {
  uint32_t* d_flag_counters = nullptr; // 新增：per-block color 指针
​
  void init() {
    // ... 其他初始化
    // 为每个 block 分配 counter，初始值为 1（避免与零初始化的 flags 冲突）
    HIP_CHECK(hipMalloc(&d_flag_counters, kMaxNumBlocks * sizeof(uint32_t)));
    std::vector<uint32_t> init_color(kMaxNumBlocks, 1u);
    HIP_CHECK(hipMemcpy(d_flag_counters, init_color.data(),
                        kMaxNumBlocks * sizeof(uint32_t),
                        hipMemcpyHostToDevice));
  }
​
  void destroy() {
    // 在 initialized 标志前释放 d_flag_counters，避免部分初始化泄漏
    if (d_flag_counters) {
      HIP_CHECK(hipFree(d_flag_counters));
      d_flag_counters = nullptr;
    }
    if (initialized) {
      // 释放其他资源 ...
    }
  }
};
​
} // namespace quickreduce

评论区精华

• 资源泄漏风险：机器人建议扩展 destroy() 中的资源释放范围，PR 仅处理了 d_flag_counters。（未解决）
• 测试覆盖要求：HaiShaw 要求添加测试，作者随后添加了手动测试。（已解决）

风险与影响

• 风险：核心路径（通信 kernel）的修改可能引入潜在 GPU 级错误；新测试未集成 CI，需手动运行；资源泄漏防护不完整。
• 影响：影响所有 AMD MI300 等使用 QuickReduce 且启用 CUDA Graph 的用户，修复了严重的正确性问题，无性能回退。不影响其他架构。

关联脉络

本 PR 与 AMD QuickReduce 模块的其他历史 PR 共同完善了 AMD 平台的通信可靠性。目前无直接跨 PR 的依赖。未来应考虑将手动测试集成到 CI 中，并跟进资源泄漏建议。