Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
chen2021673 wants to merge 2 commits into
base: master
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ea5d098
fix(cuda): handle 64-lane warps in broadcast reduction
chen2021673 Sep 1, 2026
93f8603
fix(cuda): include tail lanes in broadcast backward shuffles
chen2021673 Sep 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 22 additions & 14 deletions infini_train/src/kernels/cuda/elementwise.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -487,11 +487,11 @@ __global__ void BinaryBackwardKernel(T *output_a, T *output_b, FuncA fn_a, FuncB
size_t num_elements, const T *grad_output, const T *input_a, const T *input_b) {
extern __shared__ char shared_memory[];
const int tid = threadIdx.x;
const int warp_id = tid / 32;
const int lane_id = tid % 32;
const int lane_id = tid % kWarpSize;
const int logical_warp_id = tid / kWarpSize;

using WarpReduce = cub::WarpReduce<float>;
WarpReduce::TempStorage *temp_storage = reinterpret_cast<WarpReduce::TempStorage *>(shared_memory);
using WarpReduce = cub::WarpReduce<float, kWarpSize>;
auto *temp_storage = reinterpret_cast<typename WarpReduce::TempStorage *>(shared_memory);

size_t idx = blockIdx.x * blockDim.x + tid;
bool in_bounds = (idx < num_elements);
Expand All @@ -509,29 +509,37 @@ __global__ void BinaryBackwardKernel(T *output_a, T *output_b, FuncA fn_a, FuncB
grad_val = common::cuda::Cast<float>(Mul<T>(grad_output[idx], fn_b(a_val, b_val)));
}

unsigned active_mask = __ballot_sync(0xFFFFFFFF, in_bounds);
if (!active_mask) {
using WarpMask = decltype(__ballot_sync(~uint64_t{0}, true));
const WarpMask full_mask = ~WarpMask{0};
const WarpMask physical_active_mask = __ballot_sync(full_mask, in_bounds);
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;
if (active_mask == 0) {
return;
}

int leader = __ffs(active_mask) - 1;
int64_t common_offset = __shfl_sync(active_mask, b_offset, leader);
const unsigned logical_active_mask = static_cast<unsigned>(static_cast<uint64_t>(active_mask) >> logical_base);
const int leader = __ffs(logical_active_mask) - 1;
// All lanes in a nonempty logical warp participate, including out-of-bounds lanes with zero gradients.
// Use the active mask only to select valid offsets so warp_uniform agrees across all lanes before Sum.
const int64_t common_offset = __shfl_sync(logical_lane_mask, b_offset, leader, kWarpSize);

// Check if all active threads share common b_offset
bool warp_uniform = true;
for (int i = 0; i < 32; ++i) {
if (!(active_mask & (1 << i))) {
for (int i = 0; i < kWarpSize; ++i) {
if (!(logical_active_mask & (unsigned{1} << i))) {
continue;
}
int64_t offset_i = __shfl_sync(active_mask, b_offset, i);
const int64_t offset_i = __shfl_sync(logical_lane_mask, b_offset, i, kWarpSize);
if (offset_i != common_offset) {
warp_uniform = false;
break;
}
}

if (warp_uniform) {
float reduced = WarpReduce(temp_storage[warp_id]).Sum(grad_val);
const float reduced = WarpReduce(temp_storage[logical_warp_id]).Sum(grad_val);
if (lane_id == leader) {
// FIXME(lzm): atomicAdd is much slower for bf16 and half compared to float, needs further optimization
atomicAdd(&output_b[common_offset], common::cuda::Cast<T>(reduced));
Expand Down Expand Up @@ -692,7 +700,7 @@ void LaunchBackward(FuncA fun_a, FuncB fun_b, const std::shared_ptr<Tensor> &out
[=](dim3 grid, dim3 block, size_t /*offset*/, auto... ptrs) {
const int block_threads = static_cast<int>(block.x);
const int num_warps = CEIL_DIV(block_threads, kWarpSize);
const size_t smem_size = num_warps * sizeof(cub::WarpReduce<float>::TempStorage);
const size_t smem_size = num_warps * sizeof(cub::WarpReduce<float, kWarpSize>::TempStorage);
BinaryBackwardKernel<<<grid, block, smem_size, stream>>>(output_a_ptr, output_b_ptr, fun_a, fun_b, meta,
num_elements, grad_output_ptr, ptrs...);
},
Expand Down
53 changes: 53 additions & 0 deletions tests/autograd/test_autograd_elementwise_backward.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,59 @@ TEST_P(AutogradElementwiseBackwardTest, MulBackward) {
EXPECT_EQ(grad_inputs.size(), 2);
}

TEST_P(AutogradElementwiseBackwardTest, Float32MulBroadcastBackwardAcrossLogicalWarps) {
auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 64}, DataType::kFLOAT32, GetDevice(), true);
a->Fill(1.0f);
auto b = std::make_shared<Tensor>(std::vector<int64_t>{2, 1}, DataType::kFLOAT32, GetDevice(), true);
b->Fill(2.0f);
auto mul_fn = std::make_shared<autograd::Mul>();
auto result = mul_fn->Apply({a, b});
auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 64}, DataType::kFLOAT32, GetDevice(), true);
grad->Fill(1.0f);

auto grad_inputs = mul_fn->Backward({grad});
ASSERT_EQ(grad_inputs.size(), 2);

test::ExpectTensorFloatEqual(grad_inputs[0], 2.0f);
test::ExpectTensorFloatEqual(grad_inputs[1], std::vector<float>{64.0f, 64.0f});
}

TEST_P(AutogradElementwiseBackwardTest, Float32MulBroadcastBackwardPartialLogicalWarps) {
// A single row covers 31/33-element tails; multiple rows also exercise nonzero B offsets
// and logical warps that straddle rows with different B offsets.
for (int64_t rows : {1, 3}) {
for (int64_t cols : {31, 33}) {
SCOPED_TRACE(::testing::Message() << "rows=" << rows << ", cols=" << cols);
const std::vector<int64_t> a_dims{rows, cols};
const std::vector<int64_t> b_dims{rows, 1};
std::vector<float> a_values(rows * cols), b_values(rows), grad_values(rows * cols);
std::vector<float> expected_grad_a(rows * cols), expected_grad_b(rows, 0.0f);
for (int64_t row = 0; row < rows; ++row) {
b_values[row] = static_cast<float>(row + 2);
for (int64_t col = 0; col < cols; ++col) {
const int64_t idx = row * cols + col;
a_values[idx] = static_cast<float>(idx + 1);
grad_values[idx] = static_cast<float>(col % 3 + 1);
expected_grad_a[idx] = grad_values[idx] * b_values[row];
expected_grad_b[row] += grad_values[idx] * a_values[idx];
}
}

auto a = std::make_shared<Tensor>(a_values.data(), a_dims, DataType::kFLOAT32, GetDevice());
auto b = std::make_shared<Tensor>(b_values.data(), b_dims, DataType::kFLOAT32, GetDevice());
auto mul_fn = std::make_shared<autograd::Mul>();
auto result = mul_fn->Apply({a, b});
auto grad = std::make_shared<Tensor>(grad_values.data(), a_dims, DataType::kFLOAT32, GetDevice());
auto grad_inputs = mul_fn->Backward({grad});
ASSERT_EQ(grad_inputs.size(), 2);
EXPECT_EQ(grad_inputs[0]->Dims(), a_dims);
EXPECT_EQ(grad_inputs[1]->Dims(), b_dims);
test::ExpectTensorFloatEqual(grad_inputs[0], expected_grad_a);
test::ExpectTensorFloatEqual(grad_inputs[1], expected_grad_b);
}
}
}

TEST_P(AutogradElementwiseBackwardTest, BFloat16MulBroadcastBackwardLargeBlock) {
ONLY_CUDA();
auto a = std::make_shared<Tensor>(std::vector<int64_t>{512, 8192}, DataType::kBFLOAT16, GetDevice(), true);
Expand Down
Loading