perf: reduce redundant CUDA initialization and synchronization

example/gpt2/main.cc

@@ -479,10 +479,10 @@ void Train(const nn::parallel::Rank &rank) {
 
                 LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish loss forward";
 
-                auto loss_cpu = loss->To(Device());
-                lossf += static_cast<const float *>(loss_cpu.DataPtr())[0];
                 LOG(INFO) << "Rank " << rank.GlobalRank() << ": start backward";
                 loss->Backward();
+                auto loss_cpu = loss->To(Device());
+                lossf += static_cast<const float *>(loss_cpu.DataPtr())[0];
                 LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward";
             }
 

example/llama3/main.cc

@@ -456,10 +456,10 @@ void Train(const nn::parallel::Rank &rank) {
 
                 LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish loss forward";
 
-                auto loss_cpu = loss->To(Device());
-                lossf += static_cast<const float *>(loss_cpu.DataPtr())[0];
                 LOG(INFO) << "Rank " << rank.GlobalRank() << ": start backward";
                 loss->Backward();
+                auto loss_cpu = loss->To(Device());
+                lossf += static_cast<const float *>(loss_cpu.DataPtr())[0];
                 LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward";
             }
 

example/mnist/main.cc

@@ -70,6 +70,8 @@ int main(int argc, char *argv[]) {
             optimizer.ZeroGrad();
 
             auto loss = loss_fn.Forward({outputs[0], new_label});
+            loss[0]->Backward();
+
             auto loss_cpu = loss[0]->To(cpu_device);
             float current_loss = static_cast<float *>(loss_cpu.DataPtr())[0];
             total_loss += current_loss;
@@ -79,7 +81,6 @@ int main(int argc, char *argv[]) {
                            << " loss: " << current_loss;
             }
 
-            loss[0]->Backward();
             optimizer.Step();
             train_idx += 1;
         }

infini_train/src/kernels/cuda/concat.cu

@@ -186,7 +186,7 @@ std::vector<std::shared_ptr<Tensor>> ConcatBackward(const std::shared_ptr<Tensor
     grads.reserve(input_dims_list.size());
     for (const auto &dvec : input_dims_list) {
         auto t = std::make_shared<Tensor>(dvec, dtype, device);
-        t->Fill(0.0);
+        // ConcatBackwardKernel maps every grad_output element to exactly one grad tensor element; no Fill is needed.
         grads.push_back(t);
     }
 

infini_train/src/kernels/cuda/cross_entropy.cu

@@ -204,7 +204,7 @@ std::shared_ptr<Tensor> CrossEntropyBackward(const std::shared_ptr<Tensor> &inpu
                                  DataTypeList<INFINI_ALL_FLOATING_TYPES>>(
         {target->Dtype(), input_casted->Dtype()},
         [=]<typename Ttarget, typename Tinput>() {
-            grad_input->Fill(0.0);
+            // One sample block writes all of its num_classes gradient elements; no Fill is needed.
             const Tinput *output_grad_ptr = static_cast<const Tinput *>(grad_output->DataPtr());
             const Ttarget *target_ptr = static_cast<const Ttarget *>(target->DataPtr());
             const Tinput *input_ptr = static_cast<const Tinput *>(input_casted->DataPtr());

infini_train/src/kernels/cuda/layernorm.cu

@@ -89,8 +89,7 @@ LayerNormForward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Ten
     core::cuda::DispatchCudaFunc<INFINI_ALL_FLOATING_TYPES>(
         dtype,
         [=]<typename T>() {
-            mean->Fill(0.0);
-            rstd->Fill(0.0);
+            // Each token block writes its mean and rstd exactly once; no Fill is needed.
             LayerNormForwardKernel<BLOCK_SIZE><<<num_blocks, threads_per_block, 0, cuda_stream>>>(
                 static_cast<const T *>(input->DataPtr()), static_cast<const T *>(weight->DataPtr()),
                 static_cast<const T *>(bias->DataPtr()), static_cast<float *>(mean->DataPtr()),
@@ -183,7 +182,7 @@ LayerNormBackward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Te
     core::cuda::DispatchCudaFunc<INFINI_ALL_FLOATING_TYPES>(
         dtype,
         [=]<typename T>() {
-            grad_input->Fill(0.0);
+            // Each token block writes its complete grad_input slice; no Fill is needed.
             grad_weight->Fill(0.0);
             grad_bias->Fill(0.0);
             LayerNormBackwardKernel<BLOCK_SIZE><<<num_blocks, threads_per_block, 0, cuda_stream>>>(

infini_train/src/kernels/cuda/linear.cu

@@ -65,6 +65,7 @@ std::shared_ptr<Tensor> LinearForward(const std::shared_ptr<Tensor> &input, cons
                                  infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device))
                                  ->cuda_stream();
 
+    const float beta = bias ? 1.0f : 0.0f;
     if (bias) {
         CHECK_EQ(bias->Dims().size(), 1);
         CHECK_EQ(bias->Dims()[0], out_features);
@@ -78,9 +79,8 @@ std::shared_ptr<Tensor> LinearForward(const std::shared_ptr<Tensor> &input, cons
                     static_cast<T *>(output->DataPtr()), static_cast<const T *>(bias->DataPtr()), bs, out_features);
             },
             "CUDA LinearForward");
-    } else {
-        output->Fill(0.0);
     }
+    // In the no-bias path, beta=0 makes cuBLAS fully overwrite output; no Fill is needed.
 
     // When bs==1 and fp32, use cublasSgemv (more efficient than GEMM for matrix-vector).
     // cublasSgemv does not support bf16, so bf16 falls through to Gemm.
@@ -94,7 +94,7 @@ std::shared_ptr<Tensor> LinearForward(const std::shared_ptr<Tensor> &input, cons
                               .x = static_cast<const float *>(input->DataPtr()),
                               .y = static_cast<float *>(output->DataPtr()),
                               .alpha = 1.0f,
-                              .beta = 1.0f, // output already initialized with bias or zero above
+                              .beta = beta,
                           });
     } else {
         // cuBLAS is colmun-major
@@ -125,7 +125,7 @@ std::shared_ptr<Tensor> LinearForward(const std::shared_ptr<Tensor> &input, cons
                 .C = output->DataPtr(),
                 .ldc = static_cast<int>(out_features),
                 .alpha = 1.0f,
-                .beta = 1.0f, // bias already written into output; beta=1 accumulates
+                .beta = beta,
                 .batch_count = 1,
                 .input_dtype = dtype,
                 .output_dtype = dtype,

infini_train/src/kernels/cuda/reduction.cu

@@ -181,7 +181,7 @@ std::shared_ptr<Tensor> ReduceOpBackward(const std::shared_ptr<Tensor> &grad_out
     core::cuda::DispatchCudaFunc<INFINI_ALL_FLOATING_TYPES>(
         dtype,
         [=]<typename T>() {
-            grad_input->Fill(0.0);
+            // The backward kernel assigns every grad_input element on all reduction branches; no Fill is needed.
             GenericReduceBackwardKernel<<<num_blocks, threads_per_block, 0, cuda_stream>>>(
                 static_cast<T *>(grad_input->DataPtr()), static_cast<const T *>(grad_output->DataPtr()),
                 input ? static_cast<const T *>(input->DataPtr()) : nullptr,

infini_train/src/kernels/cuda/slice.cu

@@ -48,8 +48,7 @@ std::shared_ptr<Tensor> SliceForward(const std::shared_ptr<Tensor> &input, const
 
     auto dtype = input->Dtype();
     auto new_tensor = std::make_shared<Tensor>(new_dims, dtype, input->GetDevice());
-    // NOTE(zbl): must initialize with 0
-    new_tensor->Fill(0.0);
+    // SliceForwardKernel writes every output index in [0, total_elements); no Fill is needed.
 
     std::vector<int64_t> src_strides(dims.size(), 0), dst_strides(new_dims.size(), 0);
     int64_t stride = 1;

infini_train/src/kernels/cuda/softmax.cu

@@ -198,7 +198,7 @@ std::shared_ptr<Tensor> SoftmaxBackward(const std::shared_ptr<Tensor> &grad_outp
     CHECK(dim >= 0 && dim < output->Dims().size());
 
     auto grad_input = std::make_shared<Tensor>(output_dims, promoted_type, output->GetDevice());
-    grad_input->Fill(0.0);
+    // For non-empty tensors, the grid covers every outer/axis/inner index exactly once; no Fill is needed.
 
     switch (promoted_type) {
         DISPATCH_CASE(WRAP(LaunchBackward<256, float>(grad_input, grad_output_promoted, output_promoted, dim);),

infini_train/src/kernels/cuda/split.cu

@@ -114,6 +114,7 @@ std::shared_ptr<Tensor> LaunchSplitBackward(const std::vector<int64_t> &input_di
     const auto &grad = grad_outputs[0];
     auto dtype = grad->Dtype();
     auto grad_input = std::make_shared<Tensor>(input_dims, dtype, grad->GetDevice());
+    // Keep initialization: defensive early returns in SplitBackwardKernel can leave elements unwritten.
     grad_input->Fill(0.0);
 
     int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies<int64_t>());

infini_train/src/kernels/cuda/stack.cu

@@ -113,7 +113,7 @@ std::vector<std::shared_ptr<Tensor>> StackBackward(const std::vector<int64_t> &i
     std::vector<std::shared_ptr<Tensor>> grads;
     for (int i = 0; i < num_inputs; ++i) {
         auto t = std::make_shared<Tensor>(base_dims, dtype, grad_output->GetDevice());
-        t->Fill(0.0);
+        // StackBackwardKernel writes every element of every grad tensor exactly once; no Fill is needed.
         grads.push_back(t);
     }
 

infini_train/src/nn/parallel/pp/pipeline_schedule.cc

@@ -255,9 +255,8 @@ float PipelineSchedule::StepMicroBatches(const std::vector<std::shared_ptr<Tenso
                         {activations[task.local_chunk_idx][mb][0], std::make_shared<Tensor>(target_on_device)})[0];
                     loss = loss / n;
                 }
-                total_loss += static_cast<const float *>(loss->To(Device()).DataPtr())[0];
-
                 loss->Backward();
+                total_loss += static_cast<const float *>(loss->To(Device()).DataPtr())[0];
             } else {
                 auto out_tensor = activations[task.local_chunk_idx][mb][0];
 

tests/autograd/test_autograd_linear_forward.cc

@@ -3,6 +3,7 @@
 #include "gtest/gtest.h"
 
 #include "infini_train/include/autograd/linear.h"
+#include "infini_train/include/core/runtime/device_guard.h"
 #include "infini_train/include/nn/parallel/global.h"
 #include "infini_train/include/tensor.h"
 
@@ -12,17 +13,30 @@ using namespace infini_train;
 
 class AutogradLinearForwardTest : public infini_train::test::InfiniTrainTest {};
 
+namespace {
+std::shared_ptr<Tensor> CopyToCPU(const std::shared_ptr<Tensor> &tensor) {
+    auto cpu = std::make_shared<Tensor>(tensor->Dims(), tensor->Dtype(), Device());
+    cpu->CopyFrom(tensor);
+    core::GetDeviceGuardImpl(tensor->GetDevice().type())->SynchronizeDevice(tensor->GetDevice());
+    return cpu;
+}
+} // namespace
+
 TEST_P(AutogradLinearForwardTest, LinearForward) {
     auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
     input->Fill(1.0f);
     auto weight = std::make_shared<Tensor>(std::vector<int64_t>{4, 3}, DataType::kFLOAT32, GetDevice(), true);
     weight->Fill(1.0f);
     auto bias = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice(), true);
-    bias->Fill(0.0f);
+    bias->Fill(2.0f);
     auto linear_fn = std::make_shared<autograd::Linear>();
     auto result = linear_fn->Apply({input, weight, bias});
     EXPECT_EQ(result.size(), 1);
     EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{2, 4}));
+
+    auto output = CopyToCPU(result[0]);
+    const float *actual = static_cast<const float *>(output->DataPtr());
+    for (int64_t i = 0; i < output->NumElements(); ++i) { EXPECT_FLOAT_EQ(actual[i], 5.0f); }
 }
 
 TEST_P(AutogradLinearForwardTest, LinearNoBias) {
@@ -34,6 +48,10 @@ TEST_P(AutogradLinearForwardTest, LinearNoBias) {
     auto result = linear_fn->Apply({input, weight});
     EXPECT_EQ(result.size(), 1);
     EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{2, 4}));
+
+    auto output = CopyToCPU(result[0]);
+    const float *actual = static_cast<const float *>(output->DataPtr());
+    for (int64_t i = 0; i < output->NumElements(); ++i) { EXPECT_FLOAT_EQ(actual[i], 3.0f); }
 }
 
 TEST_P(AutogradLinearForwardTest, LinearBatch) {

tests/autograd/test_autograd_loss.cc

@@ -0,0 +1,63 @@
+#include <cmath>
+#include <vector>
+
+#include "gtest/gtest.h"
+
+#include "infini_train/include/autograd/loss.h"
+#include "infini_train/include/core/runtime/device_guard.h"
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/tensor.h"
+
+#include "tests/common/test_utils.h"
+
+using namespace infini_train;
+
+class AutogradLossTest : public infini_train::test::InfiniTrainTest {};
+
+namespace {
+std::shared_ptr<Tensor> CopyToCPU(const std::shared_ptr<Tensor> &tensor) {
+    auto cpu = std::make_shared<Tensor>(tensor->Dims(), tensor->Dtype(), Device());
+    cpu->CopyFrom(tensor);
+    core::GetDeviceGuardImpl(tensor->GetDevice().type())->SynchronizeDevice(tensor->GetDevice());
+    return cpu;
+}
+} // namespace
+
+TEST_P(AutogradLossTest, CrossEntropyForwardAndBackwardValues) {
+    std::vector<float> logits_values{1.0f, 2.0f, 3.0f, 2.0f, 1.0f, 0.0f};
+    auto logits
+        = std::make_shared<Tensor>(logits_values.data(), std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    auto target = std::make_shared<Tensor>(std::vector<int64_t>{2}, DataType::kINT64, GetDevice());
+    target->Fill(0);
+
+    auto cross_entropy = std::make_shared<autograd::CrossEntropy>();
+    auto result = cross_entropy->Apply({logits, target});
+    ASSERT_EQ(result.size(), 1);
+    ASSERT_TRUE(result[0]->Dims().empty());
+
+    const float row0_sum = std::exp(1.0f) + std::exp(2.0f) + std::exp(3.0f);
+    const float row1_sum = std::exp(2.0f) + std::exp(1.0f) + std::exp(0.0f);
+    const float expected_loss = 0.5f * ((std::log(row0_sum) - 1.0f) + (std::log(row1_sum) - 2.0f));
+    auto loss_cpu = CopyToCPU(result[0]);
+    EXPECT_NEAR(static_cast<const float *>(loss_cpu->DataPtr())[0], expected_loss, 1e-5f);
+
+    auto grad_output = std::make_shared<Tensor>(std::vector<int64_t>{}, DataType::kFLOAT32, GetDevice());
+    grad_output->Fill(1.0f);
+    auto grad_inputs = cross_entropy->Backward({grad_output});
+    ASSERT_EQ(grad_inputs.size(), 2);
+    ASSERT_NE(grad_inputs[0], nullptr);
+    EXPECT_EQ(grad_inputs[1], nullptr);
+
+    auto grad_cpu = CopyToCPU(grad_inputs[0]);
+    const float *actual_grad = static_cast<const float *>(grad_cpu->DataPtr());
+    for (int row = 0; row < 2; ++row) {
+        const float sum = row == 0 ? row0_sum : row1_sum;
+        for (int col = 0; col < 3; ++col) {
+            const float probability = std::exp(logits_values[row * 3 + col]) / sum;
+            const float expected_grad = 0.5f * (probability - (col == 0 ? 1.0f : 0.0f));
+            EXPECT_NEAR(actual_grad[row * 3 + col], expected_grad, 1e-5f);
+        }
+    }
+}
+
+INFINI_TRAIN_REGISTER_TEST(AutogradLossTest);

tests/autograd/test_autograd_normalization_backward.cc

@@ -3,6 +3,7 @@
 #include "gtest/gtest.h"
 
 #include "infini_train/include/autograd/normalization.h"
+#include "infini_train/include/core/runtime/device_guard.h"
 #include "infini_train/include/nn/parallel/global.h"
 #include "infini_train/include/tensor.h"
 
@@ -12,6 +13,17 @@ using namespace infini_train;
 
 class AutogradNormalizationBackwardTest : public infini_train::test::InfiniTrainTest {};
 
+namespace {
+void ExpectValues(const std::shared_ptr<Tensor> &tensor, float expected) {
+    auto cpu = std::make_shared<Tensor>(tensor->Dims(), tensor->Dtype(), Device());
+    cpu->CopyFrom(tensor);
+    core::GetDeviceGuardImpl(tensor->GetDevice().type())->SynchronizeDevice(tensor->GetDevice());
+
+    const float *actual = static_cast<const float *>(cpu->DataPtr());
+    for (int64_t i = 0; i < cpu->NumElements(); ++i) { EXPECT_NEAR(actual[i], expected, 1e-5f); }
+}
+} // namespace
+
 TEST_P(AutogradNormalizationBackwardTest, LayerNormBackward) {
     auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3, 4}, DataType::kFLOAT32, GetDevice(), true);
     a->Fill(1.0f);
@@ -25,6 +37,9 @@ TEST_P(AutogradNormalizationBackwardTest, LayerNormBackward) {
     grad->Fill(1.0f);
     auto grad_inputs = layernorm_fn->Backward({grad});
     EXPECT_EQ(grad_inputs.size(), 3);
+    ExpectValues(grad_inputs[0], 0.0f);
+    ExpectValues(grad_inputs[1], 0.0f);
+    ExpectValues(grad_inputs[2], 6.0f);
 }
 
 TEST_P(AutogradNormalizationBackwardTest, LayerNormBackwardZeroBias) {