diff --git a/infini_train/src/kernels/cuda/elementwise.cu b/infini_train/src/kernels/cuda/elementwise.cu
index fc423b35..8c57c18b 100644
--- a/infini_train/src/kernels/cuda/elementwise.cu
+++ b/infini_train/src/kernels/cuda/elementwise.cu
@@ -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);
@@ -509,21 +509,29 @@ __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;
@@ -531,7 +539,7 @@ __global__ void BinaryBackwardKernel(T *output_a, T *output_b, FuncA fn_a, FuncB
     }
 
     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));
@@ -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...);
             },
diff --git a/tests/autograd/test_autograd_elementwise_backward.cc b/tests/autograd/test_autograd_elementwise_backward.cc
index f7eb0d5f..aca8e4af 100644
--- a/tests/autograd/test_autograd_elementwise_backward.cc
+++ b/tests/autograd/test_autograd_elementwise_backward.cc
@@ -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);