Skip to content

fix(cuda): handle 64-lane warps in broadcast backward reduction - #216

Open
chen2021673 wants to merge 3 commits into
masterfrom
fix/cuda-64-lane-warp-reduction
Open

chen2021673 wants to merge 3 commits into
masterfrom
fix/cuda-64-lane-warp-reduction

Conversation

@chen2021673

@chen2021673 chen2021673 commented Sep 1, 2026

Copy link
Copy Markdown
Contributor

背景

BroadCast 反向传播 kernel 之前默认物理 WarpSize 为 32,在 32-lane 设备上运行正常,但在部分 WarpSize 为 64 的设备上会出现规约错误。

修改内容

  • 将物理 Warp 划分为独立的 32-lane 逻辑 Warp
  • 保留完整的 ballot mask,避免 64-lane mask 被截断
  • 统一 shuffle、CUB 规约和共享内存的逻辑 WarpSize
  • 新增跨多个逻辑 Warp 的广播乘法反向传播测试

测试

新增 {2, 64} * {2, 1} Broadcast Mul Backward 用例,验证:

  • grad_a 均为 2
  • grad_b{64, 64}

maca C550 WarpSize=64 环境中测试:
image

CUDA 测试:
image
image

const int physical_lane = tid % warpSize;
const int logical_base = (physical_lane / kWarpSize) * kWarpSize;
const WarpMask logical_lane_mask = static_cast<WarpMask>(uint64_t{0xffffffff} << logical_base);
const WarpMask active_mask = physical_active_mask & logical_lane_mask;

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这个 pr 的主要解决的问题是,当 physical warp size != 32 时,能够正确得到当前 logical warp 内真正 in_bounds 的线程。

但这里有两个容易混淆的概念,建议区分一下 warpSize 和 kWarpSize 的语义:

  • warpSize 是 CUDA built-in variable,表示 hardware physical warp size;
  • kWarpSize 实际表示这个 kernel 使用的 logical warp size。

当前实现中的

uint64_t{0xffffffff} << logical_base

unsigned logical_active_mask

都隐含了 logical warp 固定为 32 lane 的假设,因此建议显式增加:

static_assert(
    kWarpSize == 32,
    "BinaryBackwardKernel assumes a 32-lane logical warp");

同时建议将 kWarpSize 重命名为 kLogicalWarpSize,避免与 built-in warpSize 混淆。

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

已修改,显式增加 static_assert 并将 kWarpSize 重命名为 kLogicalWarpSize。

@chen2021673
chen2021673 force-pushed the fix/cuda-64-lane-warp-reduction branch from 93f8603 to dea656b Compare September 18, 2026 05:04
Rename kWarpSize to kLogicalWarpSize to distinguish it from the CUDA
built-in warpSize, which represents the physical hardware warp size.

Also add a static assertion documenting that BinaryBackwardKernel assumes
a 32-lane logical warp, since its mask construction and active-lane
handling currently rely on that layout.
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.

3 participants