diff --git a/.gitmodules b/.gitmodules
index d33d0a64d..bc4569d66 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -1,12 +1,12 @@
 [submodule "third_party/glog"]
 	path = third_party/glog
-	url = git@github.com:google/glog.git
+	url = https://github.com/google/glog.git
 [submodule "third_party/gflags"]
 	path = third_party/gflags
-	url = git@github.com:gflags/gflags.git
+	url = https://github.com/gflags/gflags.git
 [submodule "third_party/eigen"]
 	path = third_party/eigen
-	url = git@github.com:InfiniTensor/eigen-mirror.git
+	url = https://github.com/InfiniTensor/eigen-mirror.git
 [submodule "third_party/googletest"]
 	path = third_party/googletest
-	url = git@github.com:google/googletest.git
+	url = https://github.com/google/googletest.git
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 6bd8069d4..69dfd35f3 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -104,7 +104,8 @@ if(USE_CUDA)
   file(GLOB_RECURSE CUDA_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/*.cu)
 
   add_library(infini_train_cuda_kernels STATIC ${CUDA_KERNELS})
-  set_target_properties(infini_train_cuda_kernels PROPERTIES CUDA_ARCHITECTURES "75;80;90")
+  set(INFINITRAIN_CUDA_ARCHITECTURES "75;80;90" CACHE STRING "CUDA architectures for kernel compilation")
+  set_target_properties(infini_train_cuda_kernels PROPERTIES CUDA_ARCHITECTURES "${INFINITRAIN_CUDA_ARCHITECTURES}")
 
   target_link_libraries(infini_train_cuda_kernels
     PUBLIC
diff --git a/bench/binary_backward_bench.cc b/bench/binary_backward_bench.cc
new file mode 100644
index 000000000..623297885
--- /dev/null
+++ b/bench/binary_backward_bench.cc
@@ -0,0 +1,115 @@
+// Microbenchmark for BinaryBackward (Mul/Add backward) kernels: BF16 vs FP32.
+// Mirrors bench/torch_binary_backward_bench.py for PyTorch comparison.
+#include <chrono>
+#include <cmath>
+#include <cstdio>
+#include <memory>
+#include <vector>
+
+#include <cuda_runtime_api.h>
+
+#include "infini_train/include/autograd/elementwise.h"
+#include "infini_train/include/device.h"
+#include "infini_train/include/tensor.h"
+
+using namespace infini_train;
+
+namespace {
+// Pin the CUDA async mempool so it never releases cached blocks back to the OS. With the default
+// release threshold (0), every profiler stream-sync (PROFILE_MODE build) unmaps all cached memory
+// and the next iteration re-maps ~200MB of physical pages — the benchmark would measure WSL2 page
+// mapping speed (~7ms/iter) instead of the kernels. PyTorch's caching allocator never unmaps, so
+// this also keeps the comparison fair.
+void PinMemPool() {
+    cudaMemPool_t pool = nullptr;
+    if (cudaDeviceGetDefaultMemPool(&pool, 0) != cudaSuccess || pool == nullptr) {
+        return;
+    }
+    uint64_t threshold = ~0ull;
+    cudaMemPoolSetAttribute(pool, cudaMemPoolAttrReleaseThreshold, &threshold);
+}
+
+// Check that every element of t equals expect (relative error within tol). Returns max rel error.
+double MaxRelError(const std::shared_ptr<Tensor> &t, double expect) {
+    auto host = t->To(DataType::kFLOAT32).To(Device());
+    const float *data = static_cast<const float *>(host.DataPtr());
+    double max_err = 0.0;
+    for (size_t i = 0; i < host.NumElements(); ++i) {
+        const double err = std::abs(static_cast<double>(data[i]) - expect) / std::max(std::abs(expect), 1e-12);
+        max_err = std::max(max_err, err);
+    }
+    return max_err;
+}
+} // namespace
+
+static float TimeBackward(std::function<void()> fn, int warmup, int iters) {
+    for (int i = 0; i < warmup; ++i) { fn(); }
+    cudaDeviceSynchronize();
+    cudaEvent_t start, stop;
+    cudaEventCreate(&start);
+    cudaEventCreate(&stop);
+    cudaEventRecord(start);
+    for (int i = 0; i < iters; ++i) { fn(); }
+    cudaEventRecord(stop);
+    cudaEventSynchronize(stop);
+    float ms = 0.0f;
+    cudaEventElapsedTime(&ms, start, stop);
+    cudaEventDestroy(start);
+    cudaEventDestroy(stop);
+    return ms * 1000.0f / iters; // us per iter
+}
+
+template <typename Op>
+static void RunCase(const char *op_name, std::vector<int64_t> a_dims, std::vector<int64_t> b_dims, DataType dtype) {
+    constexpr bool kIsMul = std::is_same_v<Op, autograd::Mul>;
+    auto dev = Device(Device::DeviceType::kCUDA, 0);
+    auto a = std::make_shared<Tensor>(a_dims, dtype, dev, true);
+    a->Fill(2.0f);
+    auto b = std::make_shared<Tensor>(b_dims, dtype, dev, true);
+    b->Fill(3.0f);
+    auto op = std::make_shared<Op>();
+    auto out = op->Apply({a, b});
+    auto grad = std::make_shared<Tensor>(a_dims, dtype, dev, true);
+    grad->Fill(1.0f);
+
+    // One-shot correctness check. a=2, b=3, grad=1:
+    //   mul: ga = g*b = 3;        gb = g*a = 2 per use (row-bcast: rows*2, col-bcast: cols*2)
+    //   add: ga = g = 1;          gb = g = 1 per use   (row-bcast: rows,   col-bcast: cols)
+    {
+        auto grads = op->Backward({grad});
+        const double ga_expect = kIsMul ? 3.0 : 1.0;
+        const double unit = kIsMul ? 2.0 : 1.0;
+        double gb_expect = unit;
+        if (a_dims != b_dims) {
+            gb_expect = unit * (b_dims.size() == 1 ? a_dims[0] : a_dims[1]);
+        }
+        const double tol = dtype == DataType::kFLOAT32 ? 1e-5 : 1e-2;
+        const double err_a = MaxRelError(grads[0], ga_expect);
+        const double err_b = MaxRelError(grads[1], gb_expect);
+        const bool pass = err_a <= tol && err_b <= tol;
+        printf("  correctness: %s (ga max rel err %.2e, gb max rel err %.2e, tol %.0e)\n", pass ? "PASS" : "FAIL",
+               err_a, err_b, tol);
+    }
+
+    const float us = TimeBackward([&]() { auto g = op->Backward({grad}); }, 20, 200);
+    const double bytes = 5.0 * a->NumElements() * (dtype == DataType::kFLOAT32 ? 4 : 2); // g,a,b in; ga,gb out (approx)
+    printf("[%-5s] a=[%ld,%ld] b_bcast=%-5s %-8s %8.1f us  (~%.0f GB/s)\n", op_name, a_dims[0], a_dims[1],
+           a_dims == b_dims ? "false" : "true", dtype == DataType::kFLOAT32 ? "float32" : "bfloat16", us,
+           bytes / (us * 1e-6) / 1e9);
+}
+
+int main() {
+    google::InitGoogleLogging("binary_backward_bench");
+    PinMemPool();
+    const std::vector<std::pair<int64_t, int64_t>> shapes = {{65536, 768}, {8192, 3072}};
+    for (auto [r, c] : shapes) {
+        for (DataType dt : {DataType::kBFLOAT16, DataType::kFLOAT32}) {
+            RunCase<autograd::Mul>("mul", {r, c}, {r, c}, dt);
+            RunCase<autograd::Mul>("mul", {r, c}, {c}, dt);    // row-broadcast (bias style)
+            RunCase<autograd::Mul>("mul", {r, c}, {r, 1}, dt); // col-broadcast
+            RunCase<autograd::Add>("add", {r, c}, {r, c}, dt);
+            RunCase<autograd::Add>("add", {r, c}, {c}, dt);
+        }
+    }
+    return 0;
+}
diff --git a/bench/torch_binary_backward_bench.py b/bench/torch_binary_backward_bench.py
new file mode 100644
index 000000000..bd81f353c
--- /dev/null
+++ b/bench/torch_binary_backward_bench.py
@@ -0,0 +1,51 @@
+"""PyTorch baseline for binary-op backward kernels, BF16 vs FP32.
+
+Covers the patterns that hit InfiniTrain's BinaryBackward:
+  1. no-broadcast elementwise (mul / add) on [B*T, C]
+  2. broadcast backward where B is a row vector [C] (bias-add style)
+"""
+
+import torch
+import torch.utils.benchmark as tb
+
+DEV = "cuda"
+
+
+def bench(fn, n=200):
+    for _ in range(20):
+        fn()
+    torch.cuda.synchronize()
+    t = tb.Timer(stmt="fn()", globals={"fn": fn}).timeit(n)
+    return t.median * 1e6  # us
+
+
+def make_case(rows, cols, dtype, op, broadcast_b):
+    a = torch.randn(rows, cols, device=DEV, dtype=dtype, requires_grad=True)
+    if broadcast_b:
+        b = torch.randn(cols, device=DEV, dtype=dtype, requires_grad=True)
+    else:
+        b = torch.randn(rows, cols, device=DEV, dtype=dtype, requires_grad=True)
+    g = torch.randn(rows, cols, device=DEV, dtype=dtype)
+
+    def fn():
+        if a.grad is not None:
+            a.grad = None
+        if b.grad is not None:
+            b.grad = None
+        out = (a * b) if op == "mul" else (a + b)
+        out.backward(g)
+
+    return fn
+
+
+print(f"device: {torch.cuda.get_device_name(0)}")
+print(f"{'case':<28}{'dtype':<10}{'op':<6}{'bcastB':<8}{'us':>10}")
+for rows, cols in [(65536, 768), (8192, 3072), (262144, 768)]:
+    for dtype in (torch.bfloat16, torch.float32):
+        for op in ("mul", "add"):
+            for bcast in (False, True):
+                fn = make_case(rows, cols, dtype, op, bcast)
+                us = bench(fn)
+                print(
+                    f"[{rows:>6},{cols:>4}]  {str(dtype).split('.')[-1]:<10}{op:<6}{str(bcast):<8}{us:>10.1f}"
+                )
diff --git a/infini_train/include/autocast.h b/infini_train/include/autocast.h
index 0ce83e98d..9b833512f 100644
--- a/infini_train/include/autocast.h
+++ b/infini_train/include/autocast.h
@@ -76,6 +76,31 @@ inline constexpr std::array<DataType, static_cast<size_t>(Device::DeviceType::kC
     DataType::kFLOAT16,  // CUDA.
 };
 
+// Thread-local cache of autocast-casted leaf parameters (e.g. FP32 master weights
+// demoted to BF16/FP16 by the kLowerPrecision policy). Keyed by the source tensor's
+// address; the weak_ptr guards against address reuse after the source is freed.
+//
+// Threading model: training is single-threaded per rank, so the cache is
+// thread_local (same as tls_autocast_context) and needs no locking. The
+// invalidation entry points below must be called on the thread that mutates the
+// parameter (all in-place parameter mutations in this codebase -- optimizer
+// kernels, CopyFrom/Fill/SetData, checkpoint LoadStateDict -- run on the training
+// thread). Raw writes through Tensor::DataPtr() bypass these hooks, as do
+// initializer helpers; both only happen before training starts in practice.
+struct AutocastWeightCacheEntry {
+    std::weak_ptr<Tensor> source; // detects source destruction / pointer reuse
+    DataType target_dtype;
+    std::shared_ptr<Tensor> casted;
+};
+
+inline thread_local std::unordered_map<const Tensor *, AutocastWeightCacheEntry> tls_autocast_weight_cache;
+
+// Drop the cached cast of one tensor (call after any in-place mutation of it).
+inline void InvalidateAutocastWeightCacheEntry(const Tensor *tensor) { tls_autocast_weight_cache.erase(tensor); }
+
+// Drop the whole cache (e.g. after bulk parameter replacement).
+inline void ClearAutocastWeightCache() { tls_autocast_weight_cache.clear(); }
+
 // Thread-local context to track autocast state
 struct AutocastContext {
     bool enabled = false;                                      // Whether autocast is active in the current thread
@@ -113,6 +138,12 @@ struct AutocastContext {
             }
         };
 
+        // Only kLowerPrecision casts of FP32 leaf parameters are cacheable: those
+        // tensors are owned by the module (stable address) and are re-cast every
+        // forward without this cache. Non-leaf activations are short-lived, so keying
+        // on their address would risk stale hits after allocator reuse.
+        const bool cache_weights = policy == CastPolicy::kLowerPrecision;
+
         auto cast_arg = [&](auto &arg) {
             using T = std::decay_t<decltype(arg)>;
             if constexpr (std::is_same_v<T, std::shared_ptr<Tensor>>) {
@@ -121,7 +152,27 @@ struct AutocastContext {
                     if (is_floating_point(current_dtype)) {
                         DataType target_dtype = get_target_dtype();
                         if (current_dtype != target_dtype) {
-                            arg = std::make_shared<Tensor>(arg->To(target_dtype));
+                            if (cache_weights && current_dtype == DataType::kFLOAT32 && arg->is_leaf()
+                                && arg->requires_grad()) {
+                                auto it = tls_autocast_weight_cache.find(arg.get());
+                                if (it != tls_autocast_weight_cache.end()) {
+                                    auto source = it->second.source.lock();
+                                    if (source && source.get() == arg.get()
+                                        && it->second.target_dtype == target_dtype) {
+                                        arg = it->second.casted;
+                                        return;
+                                    }
+                                    // Stale entry (source freed or target changed): drop it.
+                                    tls_autocast_weight_cache.erase(it);
+                                }
+                                auto casted = std::make_shared<Tensor>(arg->To(target_dtype));
+                                tls_autocast_weight_cache.emplace(
+                                    arg.get(),
+                                    AutocastWeightCacheEntry{std::weak_ptr<Tensor>(arg), target_dtype, casted});
+                                arg = std::move(casted);
+                            } else {
+                                arg = std::make_shared<Tensor>(arg->To(target_dtype));
+                            }
                         }
                     }
                 }
diff --git a/infini_train/include/autograd/activations.h b/infini_train/include/autograd/activations.h
index a63977263..acf866e15 100644
--- a/infini_train/include/autograd/activations.h
+++ b/infini_train/include/autograd/activations.h
@@ -21,4 +21,16 @@ class Sigmoid : public Function {
                       const std::vector<std::shared_ptr<Tensor>> &output_tensors) override;
     std::vector<std::shared_ptr<Tensor>> Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) override;
 };
+
+class NewGELU : public Function {
+public:
+    static constexpr char kType[] = "NewGELUFunction";
+
+    NewGELU() : Function(kType) {}
+
+    std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
+    void SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
+                      const std::vector<std::shared_ptr<Tensor>> &output_tensors) override;
+    std::vector<std::shared_ptr<Tensor>> Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) override;
+};
 } // namespace infini_train::autograd
diff --git a/infini_train/src/autograd/accumulate.cc b/infini_train/src/autograd/accumulate.cc
index 0c34819f6..71cd0aa0e 100644
--- a/infini_train/src/autograd/accumulate.cc
+++ b/infini_train/src/autograd/accumulate.cc
@@ -26,10 +26,12 @@ AccumulateGrad::Backward(const std::vector<std::shared_ptr<Tensor>> &grad_output
 
     if (grad_output) {
         if (grad_output->Dtype() != tensor_->Dtype()) {
-            LOG(WARNING) << "AccumulateGrad: grad dtype (" << kDataTypeToDesc.at(grad_output->Dtype())
-                         << ") does not match parameter dtype (" << kDataTypeToDesc.at(tensor_->Dtype())
-                         << "). This indicates a dtype mismatch in the autograd graph (e.g. autocast "
-                            "running before autograd). The grad is not cast and will be used as-is.";
+            // GEMM backwards emit lower-precision (autocast) gradients while the
+            // master weight and its grad buffer stay in fp32. Cast the grad back to
+            // the parameter dtype here -- the AccumulateGrad kernel is single-dtype
+            // and would otherwise reinterpret the buffer incorrectly. This matches
+            // PyTorch's AccumulateGrad semantics for autocast training.
+            grad_output = std::make_shared<Tensor>(grad_output->To(tensor_->Dtype()));
         }
 
         const bool overwrite = tensor_->ConsumeGradOverwriteFlag();
diff --git a/infini_train/src/autograd/activations.cc b/infini_train/src/autograd/activations.cc
index bb8b8e5ea..063a00e0a 100644
--- a/infini_train/src/autograd/activations.cc
+++ b/infini_train/src/autograd/activations.cc
@@ -30,4 +30,29 @@ std::vector<std::shared_ptr<Tensor>> Sigmoid::Backward(const std::vector<std::sh
     auto device = output->GetDevice().type();
     return {Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({device, "SigmoidBackward"}, output, grad_output)};
 }
+
+std::vector<std::shared_ptr<Tensor>> NewGELU::Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {
+    CHECK_EQ(input_tensors.size(), 1);
+    const auto &input = input_tensors[0];
+
+    auto device = input->GetDevice().type();
+    return {Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({device, "NewGELUForward"}, input)};
+}
+
+void NewGELU::SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
+                           const std::vector<std::shared_ptr<Tensor>> &) {
+    // Save the forward input x; the backward kernel recomputes tanh(beta * (x + kappa * x^3)) from it.
+    ctx_.SaveForBackward({input_tensors[0]});
+}
+
+std::vector<std::shared_ptr<Tensor>> NewGELU::Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) {
+    auto saved_tensors = ctx_.GetSavedTensors();
+    CHECK_EQ(saved_tensors.size(), 1);
+    const auto &input = saved_tensors[0];
+    CHECK_EQ(grad_outputs.size(), 1);
+    const auto &grad_output = grad_outputs[0];
+
+    auto device = grad_output->GetDevice().type();
+    return {Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({device, "NewGELUBackward"}, grad_output, input)};
+}
 } // namespace infini_train::autograd
diff --git a/infini_train/src/autograd/normalization.cc b/infini_train/src/autograd/normalization.cc
index eca830b06..dffb9039c 100644
--- a/infini_train/src/autograd/normalization.cc
+++ b/infini_train/src/autograd/normalization.cc
@@ -45,10 +45,19 @@ std::vector<std::shared_ptr<Tensor>> LayerNorm::Backward(const std::vector<std::
     const auto &grad_output = grad_outputs[0];
 
     auto device = input->GetDevice().type();
+
+    // GEMM backwards emit gradients in the compute dtype (bf16 under autocast)
+    // while this op runs in fp32 (autocast kFP32 policy); promote the incoming
+    // grad back to the saved input dtype. This preserves the kernel's pre-existing
+    // single-dtype behavior bit-for-bit.
+    auto grad_output_promoted = grad_output->Dtype() == input->Dtype()
+                                  ? grad_output
+                                  : std::make_shared<Tensor>(grad_output->To(input->Dtype()));
+
     auto [grad_input, grad_weight, grad_bias]
         = Dispatcher::Instance()
               .Call<std::tuple<std::shared_ptr<Tensor>, std::shared_ptr<Tensor>, std::shared_ptr<Tensor>>>(
-                  {device, "LayerNormBackward"}, input, weight, bias, mean, rstd, grad_output);
+                  {device, "LayerNormBackward"}, input, weight, bias, mean, rstd, grad_output_promoted);
     return {grad_input, grad_weight, grad_bias};
 }
 } // namespace infini_train::autograd
diff --git a/infini_train/src/core/runtime/cuda/cuda_guard_impl.cc b/infini_train/src/core/runtime/cuda/cuda_guard_impl.cc
index 60f60b619..433a32d0e 100644
--- a/infini_train/src/core/runtime/cuda/cuda_guard_impl.cc
+++ b/infini_train/src/core/runtime/cuda/cuda_guard_impl.cc
@@ -1,6 +1,7 @@
 #include "infini_train/src/core/runtime/cuda/cuda_guard_impl.h"
 
 #include <array>
+#include <cstdint>
 #include <memory>
 #include <mutex>
 
@@ -50,6 +51,17 @@ void CudaGuardImpl::InitSingleStream(Device device) {
 
     cuda_streams[device.index()] = std::make_unique<CudaStream>();
 
+    // Keep the default memory pool's cached blocks mapped across stream/device
+    // synchronizations. With the default release threshold (0), every sync returns
+    // cached blocks to the OS and the next cudaMallocAsync re-maps physical pages,
+    // which is extremely slow on WSL2 (~5ms per allocation) and is hit on every
+    // profiled kernel launch (the profiler synchronizes the stream per kernel).
+    // On native Linux this only keeps the pool from shrinking back at sync points.
+    cudaMemPool_t default_pool;
+    CUDA_CHECK(cudaDeviceGetDefaultMemPool(&default_pool, device.index()));
+    cuuint64_t release_threshold = UINT64_MAX;
+    CUDA_CHECK(cudaMemPoolSetAttribute(default_pool, cudaMemPoolAttrReleaseThreshold, &release_threshold));
+
     CUDA_CHECK(cudaSetDevice(current_device));
 }
 
diff --git a/infini_train/src/kernels/cpu/gelu.cc b/infini_train/src/kernels/cpu/gelu.cc
new file mode 100644
index 000000000..1889098ea
--- /dev/null
+++ b/infini_train/src/kernels/cpu/gelu.cc
@@ -0,0 +1,58 @@
+#include <cmath>
+#include <memory>
+
+#include "glog/logging.h"
+
+#include "infini_train/include/dispatcher.h"
+#include "infini_train/include/tensor.h"
+
+namespace infini_train::kernels::cpu {
+namespace {
+// Constants of the NewGELU tanh approximation: beta = sqrt(2/pi), kappa = 0.044715.
+constexpr float kGeluBeta = 0.7978845608028654f;
+constexpr float kGeluKappa = 0.044715f;
+} // namespace
+
+std::shared_ptr<Tensor> NewGELUForward(const std::shared_ptr<Tensor> &input) {
+    auto output = std::make_shared<Tensor>(input->Dims(), DataType::kFLOAT32);
+    const float *input_ptr = static_cast<const float *>(input->DataPtr());
+    float *output_ptr = static_cast<float *>(output->DataPtr());
+
+    const int64_t numel = input->NumElements();
+    for (int64_t idx = 0; idx < numel; ++idx) {
+        const float x = input_ptr[idx];
+        const float inner = kGeluBeta * (x + kGeluKappa * x * x * x);
+        output_ptr[idx] = 0.5f * x * (1.0f + tanhf(inner));
+    }
+    return output;
+}
+
+std::shared_ptr<Tensor> NewGELUBackward(const std::shared_ptr<Tensor> &grad_output,
+                                        const std::shared_ptr<Tensor> &input) {
+    auto grad_input = std::make_shared<Tensor>(grad_output->Dims(), DataType::kFLOAT32);
+    const float *grad_output_ptr = static_cast<const float *>(grad_output->DataPtr());
+    const float *input_ptr = static_cast<const float *>(input->DataPtr());
+    float *grad_input_ptr = static_cast<float *>(grad_input->DataPtr());
+
+    const int64_t numel = grad_output->NumElements();
+    for (int64_t idx = 0; idx < numel; ++idx) {
+        const float x = input_ptr[idx];
+        const float x_sq = x * x;
+        const float inner = kGeluBeta * (x + kGeluKappa * x_sq * x);
+        const float tanh_inner = tanhf(inner);
+        const float left_derivative = 0.5f * (1.0f + tanh_inner);
+        const float right_derivative
+            = 0.5f * x * (1.0f - tanh_inner * tanh_inner) * kGeluBeta * (1.0f + 3.0f * kGeluKappa * x_sq);
+        grad_input_ptr[idx] = grad_output_ptr[idx] * (left_derivative + right_derivative);
+    }
+    return grad_input;
+}
+} // namespace infini_train::kernels::cpu
+
+#define REGISTER_CPU_GELU_KERNEL(kernel_name)                                                                          \
+    REGISTER_KERNEL(infini_train::Device::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name)
+
+REGISTER_CPU_GELU_KERNEL(NewGELUForward)
+REGISTER_CPU_GELU_KERNEL(NewGELUBackward)
+
+#undef REGISTER_CPU_GELU_KERNEL
diff --git a/infini_train/src/kernels/cuda/cast.cu b/infini_train/src/kernels/cuda/cast.cu
index 96a70ae28..a9bf6c826 100644
--- a/infini_train/src/kernels/cuda/cast.cu
+++ b/infini_train/src/kernels/cuda/cast.cu
@@ -1,3 +1,6 @@
+#include <algorithm>
+#include <cstdint>
+#include <cstring>
 #include <memory>
 
 #include "infini_train/include/common/common.h"
@@ -13,15 +16,70 @@
 
 namespace infini_train::kernels::cuda {
 
-template <typename Tdst, typename Tsrc>
-__global__ void CastKernel(Tdst *dst, const Tsrc *src, size_t num_elements, size_t offset) {
-    size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset;
+namespace {
 
-    if (idx < num_elements) {
+constexpr int kCastThreadsPerBlock = 256;
+// Cap the grid so the vectorized kernel uses a grid-stride loop on large tensors
+// instead of launching an unbounded number of blocks.
+constexpr size_t kCastMaxBlocks = 4096;
+
+// Unsigned chunk types for wide (up to 128-bit) vectorized loads/stores.
+template <int Bytes> struct CastChunk;
+template <> struct CastChunk<16> {
+    using type = uint4;
+};
+template <> struct CastChunk<8> {
+    using type = uint2;
+};
+template <> struct CastChunk<4> {
+    using type = uint32_t;
+};
+template <> struct CastChunk<2> {
+    using type = uint16_t;
+};
+template <> struct CastChunk<1> {
+    using type = uint8_t;
+};
+
+// Scalar grid-stride fallback for misaligned buffers.
+template <typename Tdst, typename Tsrc> __global__ void CastKernel(Tdst *dst, const Tsrc *src, size_t num_elements) {
+    const size_t stride = static_cast<size_t>(gridDim.x) * blockDim.x;
+    for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < num_elements; idx += stride) {
         dst[idx] = common::cuda::Cast<Tdst>(src[idx]);
     }
 }
 
+// Vectorized cast: each thread converts kElems elements through one wide load and
+// one wide store (each up to 128-bit). The trailing (< kElems) elements are
+// handled by the first few threads. Elementwise conversion goes through
+// common::cuda::Cast, identical to the scalar kernel.
+template <typename Tdst, typename Tsrc, int kElems>
+__global__ void CastKernelVec(Tdst *__restrict__ dst, const Tsrc *__restrict__ src, size_t num_elements) {
+    using LoadChunk = typename CastChunk<kElems * sizeof(Tsrc)>::type;
+    using StoreChunk = typename CastChunk<kElems * sizeof(Tdst)>::type;
+
+    const size_t num_vecs = num_elements / kElems;
+    const size_t stride = static_cast<size_t>(gridDim.x) * blockDim.x;
+    for (size_t v = blockIdx.x * blockDim.x + threadIdx.x; v < num_vecs; v += stride) {
+        const LoadChunk load_chunk = reinterpret_cast<const LoadChunk *>(src)[v];
+        Tsrc in[kElems];
+        Tdst out[kElems];
+        memcpy(in, &load_chunk, sizeof(in));
+#pragma unroll
+        for (int i = 0; i < kElems; ++i) { out[i] = common::cuda::Cast<Tdst>(in[i]); }
+        StoreChunk store_chunk;
+        memcpy(&store_chunk, out, sizeof(store_chunk));
+        reinterpret_cast<StoreChunk *>(dst)[v] = store_chunk;
+    }
+
+    const size_t tail_idx = num_vecs * kElems + blockIdx.x * blockDim.x + threadIdx.x;
+    if (tail_idx < num_elements) {
+        dst[tail_idx] = common::cuda::Cast<Tdst>(src[tail_idx]);
+    }
+}
+
+} // namespace
+
 std::shared_ptr<Tensor> Cast(std::shared_ptr<Tensor> input, DataType dtype) {
     auto dst_tensor = std::make_shared<Tensor>(input->Dims(), dtype, input->GetDevice());
     auto device = input->GetDevice();
@@ -30,9 +88,6 @@ std::shared_ptr<Tensor> Cast(std::shared_ptr<Tensor> input, DataType dtype) {
                                   ->cuda_stream();
 
     const size_t num_elements = input->NumElements();
-    dim3 block_dims(256);
-    dim3 grid_dims(CEIL_DIV(num_elements, block_dims.x));
-    const size_t step = grid_dims.x * block_dims.x;
 
     core::cuda::DispatchCudaFunc<DataTypeList<INFINI_ALL_NUMERIC_TYPES, INFINI_LOGICAL_TYPES>,
                                  DataTypeList<INFINI_ALL_NUMERIC_TYPES, INFINI_LOGICAL_TYPES>>(
@@ -41,8 +96,22 @@ std::shared_ptr<Tensor> Cast(std::shared_ptr<Tensor> input, DataType dtype) {
             auto dst = static_cast<Tdst *>(dst_tensor->DataPtr());
             auto src = static_cast<const Tsrc *>(input->DataPtr());
 
-            for (size_t offset = 0; offset < num_elements; offset += step) {
-                CastKernel<<<grid_dims, block_dims, 0, cuda_stream>>>(dst, src, num_elements, offset);
+            // Wide enough that both the load and the store chunk are at most 128-bit.
+            constexpr int kElems = static_cast<int>(16 / std::max(sizeof(Tdst), sizeof(Tsrc)));
+            const bool aligned
+                = (reinterpret_cast<uintptr_t>(dst) % 16 == 0) && (reinterpret_cast<uintptr_t>(src) % 16 == 0);
+            if (aligned && num_elements >= static_cast<size_t>(kElems)) {
+                const size_t num_vecs = num_elements / kElems;
+                const size_t blocks = std::max<size_t>(
+                    std::min((num_vecs + kCastThreadsPerBlock - 1) / kCastThreadsPerBlock, kCastMaxBlocks), 1);
+                CastKernelVec<Tdst, Tsrc, kElems>
+                    <<<static_cast<unsigned int>(blocks), kCastThreadsPerBlock, 0, cuda_stream>>>(dst, src,
+                                                                                                  num_elements);
+            } else {
+                const size_t blocks = std::max<size_t>(
+                    std::min((num_elements + kCastThreadsPerBlock - 1) / kCastThreadsPerBlock, kCastMaxBlocks), 1);
+                CastKernel<<<static_cast<unsigned int>(blocks), kCastThreadsPerBlock, 0, cuda_stream>>>(dst, src,
+                                                                                                        num_elements);
             }
         },
         "CUDA Cast");
diff --git a/infini_train/src/kernels/cuda/elementwise.cu b/infini_train/src/kernels/cuda/elementwise.cu
index fc423b35f..bd2aa536c 100644
--- a/infini_train/src/kernels/cuda/elementwise.cu
+++ b/infini_train/src/kernels/cuda/elementwise.cu
@@ -296,188 +296,197 @@ __global__ void UnaryBackwardKernel(T *output, Func fn, size_t num_elements, siz
     }
 }
 
-enum class BF16Path { NoBroadcast, TwoPassHist, BlockReduce };
+// Broadcast-pattern classification for the binary backward fast paths.
+// In backward the forward output always has a's dims, so a is never broadcast; meta.a_shape is the
+// right-aligned output shape and b is the only operand that may be reduced over.
+enum class BackwardBcastPattern { kRow, kCol, kGeneric };
+
+struct BackwardBcastInfo {
+    BackwardBcastPattern pattern = BackwardBcastPattern::kGeneric;
+    // kRow: K = b_numel, bin(idx) = idx % K. kCol: inner = numel / b_numel, bin(idx) = idx / inner.
+    int64_t param = 0;
+};
 
-// Lightweight and stable selector for bf16/half execution paths.
-inline BF16Path DecideBF16Path(const std::vector<int64_t> &b_shape, const std::vector<int64_t> &out_shape,
-                               size_t b_num_elements) {
-    if (ShapesEqual(b_shape, out_shape)) {
-        return BF16Path::NoBroadcast;
+// - kRow: b's non-1 dims form a suffix of the output shape (e.g. out=[rows, cols], b=[cols]).
+// - kCol: b's non-1 dims form a prefix of the output shape (e.g. out=[rows, cols], b=[rows, 1]).
+// Everything else (mixed middle-dim broadcasts) falls back to the generic kernels.
+inline BackwardBcastInfo ClassifyBackwardBroadcast(const BroadcastMeta &meta, size_t numel, size_t b_numel) {
+    BackwardBcastInfo info;
+    if (b_numel == 0 || numel % b_numel != 0) {
+        return info;
     }
-    const bool varies_last = (b_shape.back() > 1);
-    if (varies_last) {
-        if (b_num_elements <= 4096) {
-            return BF16Path::TwoPassHist; // shared histogram two-pass path
+    const int ndim = meta.ndim;
+    bool row_ok = true;
+    bool in_suffix = true;
+    for (int i = ndim - 1; i >= 0; --i) {
+        if (in_suffix && meta.b_shape[i] == meta.a_shape[i]) {
+            continue;
         }
+        if (meta.b_shape[i] == 1) {
+            in_suffix = false;
+            continue;
+        }
+        row_ok = false;
+        break;
+    }
+    if (row_ok) {
+        info.pattern = BackwardBcastPattern::kRow;
+        info.param = static_cast<int64_t>(b_numel);
+        return info;
     }
-    return BF16Path::BlockReduce; // fallback to block reduction kernel otherwise
+    bool col_ok = true;
+    bool in_prefix = true;
+    for (int i = 0; i < ndim; ++i) {
+        if (in_prefix && meta.b_shape[i] == meta.a_shape[i]) {
+            continue;
+        }
+        if (meta.b_shape[i] == 1) {
+            in_prefix = false;
+            continue;
+        }
+        col_ok = false;
+        break;
+    }
+    if (col_ok) {
+        info.pattern = BackwardBcastPattern::kCol;
+        info.param = static_cast<int64_t>(numel / b_numel);
+    }
+    return info;
 }
 
-// Each B element is used exactly once, so gradients can be written directly without reduction.
-template <typename T, typename FuncA, typename FuncB>
-__global__ void BinaryBackwardKernelNoBroadcast(T *__restrict__ outA, T *__restrict__ outB, FuncA fn_a, FuncB fn_b,
-                                                BroadcastMeta meta, size_t numel, const T *__restrict__ grad_out,
-                                                const T *__restrict__ inA, const T *__restrict__ inB) {
-    const size_t grid_stride = static_cast<size_t>(gridDim.x) * blockDim.x;
-    for (size_t idx = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x; idx < numel; idx += grid_stride) {
-        const int64_t a_off = CalcOffset(idx, meta.ndim, meta.a_strides, meta.a_shape, meta.out_strides);
-        const int64_t b_off = CalcOffset(idx, meta.ndim, meta.b_strides, meta.b_shape, meta.out_strides);
-
-        const T a = inA ? inA[a_off] : T(0);
-        const T b = inB ? inB[b_off] : T(0);
+// Maximum number of B bins held in shared memory by the row-broadcast backward kernel.
+// 12288 floats = 48 KB, the default dynamic shared-memory limit per block.
+constexpr int64_t kMaxRowBcastBins = 12288;
 
-        // Gradient for A has a one-to-one mapping, so we write directly.
-        outA[a_off] = Mul<T>(grad_out[idx], fn_a(a, b));
-
-        // Gradient for B also maps one-to-one; no atomics or reductions are required.
-        outB[b_off] = common::cuda::Cast<T>(Mul<T>(grad_out[idx], fn_b(a, b)));
-    }
+inline int CudaSmCount() {
+    int dev = 0;
+    CUDA_CHECK(cudaGetDevice(&dev));
+    int sm_count = 0;
+    CUDA_CHECK(cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev));
+    return sm_count;
 }
 
-// First pass of histogram two-pass strategy: per-block accumulation in shared memory.
-template <typename T, typename FuncA, typename FuncB>
-__global__ void BinaryBackwardBhistPass1Kernel(T *__restrict__ outA, float *__restrict__ work, FuncA fn_a, FuncB fn_b,
-                                               BroadcastMeta meta, size_t numel, int K, const T *__restrict__ grad_out,
-                                               const T *__restrict__ inA, const T *__restrict__ inB) {
-    extern __shared__ float s_hist[]; // dynamic shared memory: K bins plus padding for every 32 buckets
-    const int pad = K >> 5;           // insert one padding slot for every 32 buckets
-    const int hist_len = K + pad;
-
-    // Zero the shared histogram buffer.
-    for (int t = threadIdx.x; t < hist_len; t += blockDim.x) { s_hist[t] = 0.0f; }
+// Row-broadcast backward (b's non-1 dims are a suffix, e.g. out=[rows, K], b=[K]).
+// Each block keeps a private fp32 histogram of the K bins in shared memory and sweeps the whole
+// tensor with a grid-stride loop, then flushes to the fp32 global accumulator with one atomicAdd
+// per bin per block. The host guarantees numel % VecSize == 0 and K % VecSize == 0.
+template <typename T, int VecSize, typename FuncA, typename FuncB>
+__global__ void BinaryBackwardRowBcastKernel(T *__restrict__ outA, float *__restrict__ outB_accum, FuncA fn_a,
+                                             FuncB fn_b, size_t numel, int K, const T *__restrict__ grad_out,
+                                             const T *__restrict__ inA, const T *__restrict__ inB) {
+    extern __shared__ float s_hist[];
+    for (int k = threadIdx.x; k < K; k += blockDim.x) { s_hist[k] = 0.0f; }
     __syncthreads();
 
-    const size_t total_threads = (size_t)gridDim.x * blockDim.x;
-    for (size_t idx = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x; idx < numel; idx += total_threads) {
-        // Linearized offset for B under general broadcasting.
-        const int64_t b_off = CalcOffset(idx, meta.ndim, meta.b_strides, meta.b_shape, meta.out_strides);
-        const int bin = static_cast<int>(b_off); // assume K fits in a 32-bit int
-        const int pbin = bin + (bin >> 5);       // apply padding mapping
-
-        // Compute the offset for A under broadcasting.
-        const int64_t a_off = CalcOffset(idx, meta.ndim, meta.a_strides, meta.a_shape, meta.out_strides);
+    using VecT = aligned_vector<T, VecSize>;
+    const size_t num_vecs = numel / VecSize;
+    const int vecs_per_k = K / VecSize;
+    const size_t grid_stride = static_cast<size_t>(gridDim.x) * blockDim.x;
+    for (size_t vid = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x; vid < num_vecs; vid += grid_stride) {
+        const size_t base = vid * VecSize;
+        const int bin0 = static_cast<int>(vid % static_cast<size_t>(vecs_per_k)) * VecSize;
 
-        const T a = inA ? inA[a_off] : T(0);
-        const T b = inB ? inB[bin] : T(0); // B is indexed via the flattened bin
+        VecT g_vec = *reinterpret_cast<const VecT *>(&grad_out[base]);
+        VecT a_vec, b_vec;
+        if (inA) {
+            a_vec = *reinterpret_cast<const VecT *>(&inA[base]);
+        } else {
+#pragma unroll
+            for (int i = 0; i < VecSize; ++i) { a_vec.val[i] = T(0); }
+        }
+        if (inB) {
+            b_vec = *reinterpret_cast<const VecT *>(&inB[bin0]);
+        } else {
+#pragma unroll
+            for (int i = 0; i < VecSize; ++i) { b_vec.val[i] = T(0); }
+        }
 
-        // A is not broadcast, so gradients can be written directly.
-        outA[a_off] = Mul<T>(grad_out[idx], fn_a(a, b));
+        VecT outA_vec;
+#pragma unroll
+        for (int i = 0; i < VecSize; ++i) { outA_vec.val[i] = Mul<T>(g_vec.val[i], fn_a(a_vec.val[i], b_vec.val[i])); }
+        *reinterpret_cast<VecT *>(&outA[base]) = outA_vec;
 
-        // Accumulate B's contribution into the shared histogram using float precision.
-        const float g = common::cuda::Cast<float>(Mul<T>(grad_out[idx], fn_b(a, b)));
-        atomicAdd(&s_hist[pbin], g);
+        // Accumulate B's contribution into the shared histogram in float precision.
+#pragma unroll
+        for (int i = 0; i < VecSize; ++i) {
+            atomicAdd(&s_hist[bin0 + i],
+                      common::cuda::Cast<float>(Mul<T>(g_vec.val[i], fn_b(a_vec.val[i], b_vec.val[i]))));
+        }
     }
     __syncthreads();
 
-    // Write this block's histogram back to the global workspace: work[block, :].
-    float *dst = work + static_cast<size_t>(blockIdx.x) * static_cast<size_t>(K);
-    for (int bin = threadIdx.x; bin < K; bin += blockDim.x) {
-        const int pbin = bin + (bin >> 5);
-        dst[bin] = s_hist[pbin];
-    }
+    for (int k = threadIdx.x; k < K; k += blockDim.x) { atomicAdd(&outB_accum[k], s_hist[k]); }
 }
 
-// Second pass for histogram path: tile the workspace along CTA dimension and atomically add into float buffer.
-template <typename T>
-__global__ void BinaryBackwardBhistPass2Reduce2D(const float *__restrict__ work, float *__restrict__ outB_accum,
-                                                 size_t numBlocks, int K, int tile_height) {
-    const int k = blockIdx.x * blockDim.x + threadIdx.x;
-    if (k >= K) {
-        return;
+// Col-broadcast backward (b's non-1 dims are a prefix, e.g. out=[rows, inner], b=[rows, 1]).
+// One warp per row: reduces the row's B contribution in fp32 with shuffle and writes outB[row]
+// directly — every row is visited exactly once, so no atomics or zero-init are required.
+template <typename T, int VecSize, typename FuncA, typename FuncB>
+__global__ void BinaryBackwardColBcastKernel(T *__restrict__ outA, T *__restrict__ outB, FuncA fn_a, FuncB fn_b,
+                                             size_t rows, size_t inner, const T *__restrict__ grad_out,
+                                             const T *__restrict__ inA, const T *__restrict__ inB) {
+    using VecT = aligned_vector<T, VecSize>;
+    const size_t num_warps = (static_cast<size_t>(gridDim.x) * blockDim.x) / kWarpSize;
+    const size_t warp = (static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x) / kWarpSize;
+    const int lane = threadIdx.x % kWarpSize;
+    const size_t num_vecs = inner / VecSize;
+
+    for (size_t row = warp; row < rows; row += num_warps) {
+        const T b_val = inB ? inB[row] : T(0);
+        const size_t base = row * inner;
+        float acc = 0.0f;
+        for (size_t v = lane; v < num_vecs; v += kWarpSize) {
+            const size_t off = base + v * VecSize;
+            VecT g_vec = *reinterpret_cast<const VecT *>(&grad_out[off]);
+            VecT a_vec, outA_vec;
+            if (inA) {
+                a_vec = *reinterpret_cast<const VecT *>(&inA[off]);
+            } else {
+#pragma unroll
+                for (int i = 0; i < VecSize; ++i) { a_vec.val[i] = T(0); }
+            }
+#pragma unroll
+            for (int i = 0; i < VecSize; ++i) {
+                outA_vec.val[i] = Mul<T>(g_vec.val[i], fn_a(a_vec.val[i], b_val));
+                acc += common::cuda::Cast<float>(Mul<T>(g_vec.val[i], fn_b(a_vec.val[i], b_val)));
+            }
+            *reinterpret_cast<VecT *>(&outA[off]) = outA_vec;
+        }
+#pragma unroll
+        for (int d = kWarpSize / 2; d > 0; d >>= 1) { acc += __shfl_down_sync(0xFFFFFFFF, acc, d); }
+        if (lane == 0) {
+            outB[row] = common::cuda::Cast<T>(acc);
+        }
     }
-
-    const size_t begin_row = static_cast<size_t>(blockIdx.y) * static_cast<size_t>(tile_height);
-    const size_t end_row = min(begin_row + static_cast<size_t>(tile_height), numBlocks);
-
-    float acc = 0.0f;
-    for (size_t row = begin_row; row < end_row; ++row) { acc += work[row * static_cast<size_t>(K) + k]; }
-
-    atomicAdd(outB_accum + k, acc);
 }
 
-// Convert the accumulated float buffer back to the target type (bf16/half/float).
-template <typename T> __global__ void CastFloatToTBhist(const float *__restrict__ src, T *__restrict__ dst, int K) {
-    const int k = blockIdx.x * blockDim.x + threadIdx.x;
-    if (k < K) {
+// Cast the fp32 accumulator of the row-broadcast kernel back to the output dtype.
+template <typename T> __global__ void CastFloatToT(const float *__restrict__ src, T *__restrict__ dst, int64_t n) {
+    const int64_t k = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+    if (k < n) {
         dst[k] = common::cuda::Cast<T>(src[k]);
     }
 }
 
-// Legacy single-dimensional reduction fallback for small grids where atomic tiling is unnecessary.
-template <typename T>
-__global__ void BinaryBackwardBhistPass2Reduce1D(const float *__restrict__ work, T *__restrict__ outB, size_t numBlocks,
-                                                 int K) {
-    const size_t k = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
-    if (k >= static_cast<size_t>(K)) {
-        return;
-    }
-
-    float acc = 0.0f;
-    for (size_t b = 0; b < numBlocks; ++b) { acc += work[b * static_cast<size_t>(K) + k]; }
-    outB[k] = common::cuda::Cast<T>(acc);
-}
-
-// Helper that materializes the two-pass histogram path for bf16/half B gradients.
+// Each B element is used exactly once, so gradients can be written directly without reduction.
 template <typename T, typename FuncA, typename FuncB>
-void BinaryBackwardBhistLaunch(FuncA fn_a, FuncB fn_b, T *outA, T *outB, const T *grad_out, const BroadcastMeta &meta,
-                               size_t numel, int K, const T *inA, const T *inB, cudaStream_t stream) {
-    const int kBlockSize = 256;
-    int grid = static_cast<int>((numel + kBlockSize - 1) / kBlockSize);
-    if (grid < 1) {
-        grid = 1;
-    }
-
-    // Workspace layout: [grid, K] floats.
-    float *work = nullptr;
-    CUDA_CHECK(cudaMallocAsync(&work, static_cast<size_t>(grid) * static_cast<size_t>(K) * sizeof(float), stream));
-
-    // Pass 1: per-block histogram accumulation.
-    const size_t smem_bytes = static_cast<size_t>(K + (K >> 5)) * sizeof(float);
-    BinaryBackwardBhistPass1Kernel<T, FuncA, FuncB>
-        <<<grid, kBlockSize, smem_bytes, stream>>>(outA, work, fn_a, fn_b, meta, numel, K, grad_out, inA, inB);
-    CUDA_CHECK(cudaGetLastError());
-
-    // Pass 2: choose between 1D and 2D reductions depending on workload shape.
-    int dev = 0;
-    int sm_count = 0;
-    CUDA_CHECK(cudaGetDevice(&dev));
-    CUDA_CHECK(cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev));
-
-    const int RED_THREADS = 256;
-    const int oneD_blocks = (K + RED_THREADS - 1) / RED_THREADS;
-
-    // Use the 2D path when the 1D kernel underutilizes the SMs and there are many partial histograms to merge.
-    const bool use2D = (oneD_blocks < sm_count) && (grid > 4 * sm_count);
-
-    if (!use2D) {
-        // Fallback: reuse the legacy 1D kernel without atomics.
-        const dim3 rgrid(oneD_blocks);
-        const dim3 rblock(RED_THREADS);
-        BinaryBackwardBhistPass2Reduce1D<T><<<rgrid, rblock, 0, stream>>>(work, outB, static_cast<size_t>(grid), K);
-        CUDA_CHECK(cudaGetLastError());
-    } else {
-        // 2D tiling path: slice the workspace and accumulate using float atomics.
-        constexpr int kTileHeight = 128; // rows per CTA; tune between 128 and 256 if needed
-        float *outB_accum = nullptr;
-        CUDA_CHECK(cudaMallocAsync(&outB_accum, static_cast<size_t>(K) * sizeof(float), stream));
-        CUDA_CHECK(cudaMemsetAsync(outB_accum, 0, static_cast<size_t>(K) * sizeof(float), stream));
-
-        const dim3 rblock(RED_THREADS, 1, 1);
-        const dim3 rgrid2((K + RED_THREADS - 1) / RED_THREADS, (grid + kTileHeight - 1) / kTileHeight, 1);
+__global__ void BinaryBackwardKernelNoBroadcast(T *__restrict__ outA, T *__restrict__ outB, FuncA fn_a, FuncB fn_b,
+                                                BroadcastMeta meta, size_t numel, const T *__restrict__ grad_out,
+                                                const T *__restrict__ inA, const T *__restrict__ inB) {
+    const size_t grid_stride = static_cast<size_t>(gridDim.x) * blockDim.x;
+    for (size_t idx = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x; idx < numel; idx += grid_stride) {
+        const int64_t a_off = CalcOffset(idx, meta.ndim, meta.a_strides, meta.a_shape, meta.out_strides);
+        const int64_t b_off = CalcOffset(idx, meta.ndim, meta.b_strides, meta.b_shape, meta.out_strides);
 
-        BinaryBackwardBhistPass2Reduce2D<T>
-            <<<rgrid2, rblock, 0, stream>>>(work, outB_accum, static_cast<size_t>(grid), K, kTileHeight);
-        CUDA_CHECK(cudaGetLastError());
+        const T a = inA ? inA[a_off] : T(0);
+        const T b = inB ? inB[b_off] : T(0);
 
-        // Convert accumulated floats back to the target dtype.
-        const dim3 cgrid((K + RED_THREADS - 1) / RED_THREADS);
-        CastFloatToTBhist<T><<<cgrid, RED_THREADS, 0, stream>>>(outB_accum, outB, K);
-        CUDA_CHECK(cudaGetLastError());
+        // Gradient for A has a one-to-one mapping, so we write directly.
+        outA[a_off] = Mul<T>(grad_out[idx], fn_a(a, b));
 
-        CUDA_CHECK(cudaFreeAsync(outB_accum, stream));
+        // Gradient for B also maps one-to-one; no atomics or reductions are required.
+        outB[b_off] = common::cuda::Cast<T>(Mul<T>(grad_out[idx], fn_b(a, b)));
     }
-
-    CUDA_CHECK(cudaFreeAsync(work, stream));
 }
 
 // Backward kernel for binary operators
@@ -687,6 +696,80 @@ void LaunchBackward(FuncA fun_a, FuncB fun_b, const std::shared_ptr<Tensor> &out
     // dominated the host-side jitter floor (especially under LoRA training).
     BroadcastMeta meta = MakeBroadcastMeta(a_dims, b_dims, out_dims);
 
+    auto extract_ptrs
+        = [](const auto &...ts) { return std::make_tuple(static_cast<const T *>(ts ? ts->DataPtr() : nullptr)...); };
+    auto [input_a_ptr, input_b_ptr] = extract_ptrs(inputs...);
+
+    const size_t b_num_elements = output_b->NumElements();
+    if (b_num_elements != num_elements) {
+        const BackwardBcastInfo bcast = ClassifyBackwardBroadcast(meta, num_elements, b_num_elements);
+        static const int sm_count = CudaSmCount();
+
+        if (bcast.pattern == BackwardBcastPattern::kRow && bcast.param <= kMaxRowBcastBins) {
+            // Row broadcast (e.g. out=[rows, K], b=[K]): shared-memory fp32 histogram per block,
+            // fixed grid, grid-stride sweep; one global fp32 atomicAdd per bin per block.
+            const int K = static_cast<int>(bcast.param);
+            constexpr int kVec = kVecSize<T>;
+            const bool aligned16 = (reinterpret_cast<uintptr_t>(output_a_ptr) % 16 == 0)
+                                && (reinterpret_cast<uintptr_t>(grad_output_ptr) % 16 == 0)
+                                && (!input_a_ptr || reinterpret_cast<uintptr_t>(input_a_ptr) % 16 == 0)
+                                && (!input_b_ptr || reinterpret_cast<uintptr_t>(input_b_ptr) % 16 == 0);
+            const bool vec_ok = aligned16 && num_elements % kVec == 0 && K % kVec == 0;
+            constexpr int kBlock = 256;
+            const size_t work_items = vec_ok ? num_elements / kVec : num_elements;
+            const int grid = static_cast<int>(
+                std::max<size_t>(1, std::min(static_cast<size_t>(sm_count) * 4, CEIL_DIV(work_items, kBlock))));
+            const size_t smem = static_cast<size_t>(K) * sizeof(float);
+
+            float *outB_accum = nullptr; // fp32 scratch for low-precision dtypes
+            if constexpr (!std::is_same_v<T, float>) {
+                CUDA_CHECK(cudaMallocAsync(&outB_accum, static_cast<size_t>(K) * sizeof(float), stream));
+                CUDA_CHECK(cudaMemsetAsync(outB_accum, 0, static_cast<size_t>(K) * sizeof(float), stream));
+            }
+            // float accumulates straight into the caller-zeroed output_b.
+            float *accum_ptr = std::is_same_v<T, float> ? reinterpret_cast<float *>(output_b_ptr) : outB_accum;
+            if (vec_ok) {
+                BinaryBackwardRowBcastKernel<T, kVec><<<grid, kBlock, smem, stream>>>(
+                    output_a_ptr, accum_ptr, fun_a, fun_b, num_elements, K, grad_output_ptr, input_a_ptr, input_b_ptr);
+            } else {
+                BinaryBackwardRowBcastKernel<T, 1><<<grid, kBlock, smem, stream>>>(
+                    output_a_ptr, accum_ptr, fun_a, fun_b, num_elements, K, grad_output_ptr, input_a_ptr, input_b_ptr);
+            }
+            CUDA_CHECK(cudaGetLastError());
+            if constexpr (!std::is_same_v<T, float>) {
+                CastFloatToT<T><<<CEIL_DIV(K, kBlock), kBlock, 0, stream>>>(outB_accum, output_b_ptr, K);
+                CUDA_CHECK(cudaGetLastError());
+                CUDA_CHECK(cudaFreeAsync(outB_accum, stream));
+            }
+            return;
+        }
+
+        if (bcast.pattern == BackwardBcastPattern::kCol) {
+            // Col broadcast (e.g. out=[rows, inner], b=[rows, 1]): one warp per row, fp32 shuffle
+            // reduction, outB written directly (no atomics, no zero-init dependency).
+            const size_t inner = static_cast<size_t>(bcast.param);
+            const size_t rows = b_num_elements;
+            constexpr int kVec = kVecSize<T>;
+            const bool aligned16 = (reinterpret_cast<uintptr_t>(output_a_ptr) % 16 == 0)
+                                && (reinterpret_cast<uintptr_t>(grad_output_ptr) % 16 == 0)
+                                && (!input_a_ptr || reinterpret_cast<uintptr_t>(input_a_ptr) % 16 == 0);
+            const bool vec_ok = aligned16 && num_elements % kVec == 0 && inner % kVec == 0;
+            constexpr int kBlock = 256;
+            const size_t warps_needed = rows;
+            const int grid = static_cast<int>(std::max<size_t>(
+                1, std::min(static_cast<size_t>(sm_count) * 4, CEIL_DIV(warps_needed, kBlock / kWarpSize))));
+            if (vec_ok) {
+                BinaryBackwardColBcastKernel<T, kVec><<<grid, kBlock, 0, stream>>>(
+                    output_a_ptr, output_b_ptr, fun_a, fun_b, rows, inner, grad_output_ptr, input_a_ptr, input_b_ptr);
+            } else {
+                BinaryBackwardColBcastKernel<T, 1><<<grid, kBlock, 0, stream>>>(
+                    output_a_ptr, output_b_ptr, fun_a, fun_b, rows, inner, grad_output_ptr, input_a_ptr, input_b_ptr);
+            }
+            CUDA_CHECK(cudaGetLastError());
+            return;
+        }
+    }
+
     if constexpr (std::is_same_v<T, float>) {
         LaunchKernel<T>(
             [=](dim3 grid, dim3 block, size_t /*offset*/, auto... ptrs) {
@@ -698,21 +781,7 @@ void LaunchBackward(FuncA fun_a, FuncB fun_b, const std::shared_ptr<Tensor> &out
             },
             output_a, inputs...);
     } else if constexpr (std::is_same_v<T, __half> || std::is_same_v<T, __nv_bfloat16>) {
-        // Dynamically choose the most efficient bf16/half strategy based on broadcast pattern.
-        // Reconstruct right-aligned b_shape (stack-only, no device allocations) for
-        // DecideBF16Path which still operates on std::vector.
-        const int ndim = meta.ndim;
-        std::vector<int64_t> b_shape(meta.b_shape, meta.b_shape + ndim);
-        const std::vector<int64_t> &out_shape = out_dims;
-
-        size_t b_num_elements = 1;
-        for (auto v : b_shape) { b_num_elements *= static_cast<size_t>(v); }
-        const int K_linear = static_cast<int>(b_num_elements);
-
-        // Select the execution path.
-        const BF16Path path = DecideBF16Path(b_shape, out_shape, b_num_elements);
-
-        if (path == BF16Path::NoBroadcast) {
+        if (ShapesEqual(b_dims, out_dims)) {
             // No broadcast: write gradients directly without shared memory or atomics.
             LaunchKernel<T>(
                 [=](dim3 grid, dim3 block, size_t /*offset*/, auto... ptrs) {
@@ -723,19 +792,6 @@ void LaunchBackward(FuncA fun_a, FuncB fun_b, const std::shared_ptr<Tensor> &out
             return;
         }
 
-        if (path == BF16Path::TwoPassHist) {
-            // Small K with variation in the innermost dimension: use two-pass histogram strategy.
-            LaunchKernel<T>(
-                [=](dim3 /*grid*/, dim3 /*block*/, size_t /*offset*/, const T *input_a_ptr, const T *input_b_ptr) {
-                    BinaryBackwardBhistLaunch<T, FuncA, FuncB>(fun_a, fun_b, output_a_ptr, output_b_ptr,
-                                                               grad_output_ptr, meta, num_elements, K_linear,
-                                                               input_a_ptr, input_b_ptr, stream);
-                },
-                output_a, inputs...);
-
-            return;
-        }
-
         // Otherwise fall back to the block-reduction kernel with SoA layout and fast atomics.
         LaunchKernel<T>(
             [=](dim3 grid, dim3 block, size_t /*offset*/, auto... ptrs) {
diff --git a/infini_train/src/kernels/cuda/embedding.cu b/infini_train/src/kernels/cuda/embedding.cu
index 89361e03a..1172ddfcc 100644
--- a/infini_train/src/kernels/cuda/embedding.cu
+++ b/infini_train/src/kernels/cuda/embedding.cu
@@ -1,4 +1,6 @@
+#include <cstdint>
 #include <memory>
+#include <type_traits>
 
 #include "infini_train/include/common/cuda/common_cuda.h"
 #include "infini_train/include/core/runtime/device_guard.h"
@@ -63,22 +65,48 @@ std::shared_ptr<Tensor> EmbeddingForward(const std::shared_ptr<Tensor> &input, c
     return output;
 }
 
+// One thread block per token: threads in the block split the embedding dimension and atomically accumulate the
+// token's gradient row into grad_weight. This lifts the grid-level parallelism from ceil(num_tokens/256) blocks
+// to num_tokens blocks. Rows untouched by any token keep the zero value written by the Fill below.
 template <typename T>
 __global__ void EmbeddingBackwardKernel(const int64_t *input_ptr, const T *grad_output_ptr, T *grad_weight_ptr,
                                         int num_tokens, int embedding_dim, int vocab_size) {
-    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    const int idx = blockIdx.x;
     if (idx >= num_tokens) {
         return;
     }
 
-    int token_id = static_cast<int>(input_ptr[idx]);
+    const int token_id = static_cast<int>(input_ptr[idx]);
     if (token_id < 0 || token_id >= vocab_size) {
         return;
     }
 
-    for (int j = 0; j < embedding_dim; ++j) {
-        atomicAdd(&grad_weight_ptr[token_id * embedding_dim + j], grad_output_ptr[idx * embedding_dim + j]);
+    const T *grad_row = grad_output_ptr + static_cast<int64_t>(idx) * embedding_dim;
+    T *weight_row = grad_weight_ptr + static_cast<int64_t>(token_id) * embedding_dim;
+
+    if constexpr (std::is_same_v<T, float>) {
+        // 128-bit fast path for fp32: vectorized loads plus one vector atomicAdd per 16B chunk (sm_90+).
+        if ((embedding_dim & 3) == 0 && (reinterpret_cast<uintptr_t>(grad_row) & 0xF) == 0
+            && (reinterpret_cast<uintptr_t>(weight_row) & 0xF) == 0) {
+            const float4 *grad_row4 = reinterpret_cast<const float4 *>(grad_row);
+            float4 *weight_row4 = reinterpret_cast<float4 *>(weight_row);
+            for (int j = threadIdx.x; j < embedding_dim / 4; j += blockDim.x) {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+                atomicAdd(&weight_row4[j], grad_row4[j]);
+#else
+                const float4 grad = grad_row4[j];
+                float *weight = &weight_row[j * 4];
+                atomicAdd(&weight[0], grad.x);
+                atomicAdd(&weight[1], grad.y);
+                atomicAdd(&weight[2], grad.z);
+                atomicAdd(&weight[3], grad.w);
+#endif
+            }
+            return;
+        }
     }
+
+    for (int j = threadIdx.x; j < embedding_dim; j += blockDim.x) { atomicAdd(&weight_row[j], grad_row[j]); }
 }
 
 std::shared_ptr<Tensor> EmbeddingBackward(const std::shared_ptr<Tensor> &input, const std::vector<int64_t> &weight_dims,
@@ -100,15 +128,18 @@ std::shared_ptr<Tensor> EmbeddingBackward(const std::shared_ptr<Tensor> &input,
     auto grad_weight = std::make_shared<Tensor>(weight_dims, dtype, grad_output->GetDevice());
     const int num_tokens = input->NumElements();
     const int threads_per_block = 256;
-    const int num_blocks = (num_tokens + threads_per_block - 1) / threads_per_block;
+    // One block per token; each block cooperatively accumulates a whole gradient row.
+    const int num_blocks = num_tokens;
 
     core::cuda::DispatchCudaFunc<INFINI_ALL_FLOATING_TYPES>(
         dtype,
         [=]<typename T>() {
             grad_weight->Fill(0.0);
-            EmbeddingBackwardKernel<<<num_blocks, threads_per_block, 0, cuda_stream>>>(
-                static_cast<const int64_t *>(input->DataPtr()), static_cast<const T *>(grad_output->DataPtr()),
-                static_cast<T *>(grad_weight->DataPtr()), num_tokens, embedding_dim, vocab_size);
+            if (num_tokens > 0) {
+                EmbeddingBackwardKernel<<<num_blocks, threads_per_block, 0, cuda_stream>>>(
+                    static_cast<const int64_t *>(input->DataPtr()), static_cast<const T *>(grad_output->DataPtr()),
+                    static_cast<T *>(grad_weight->DataPtr()), num_tokens, embedding_dim, vocab_size);
+            }
         },
         "CUDA EmbeddingBackward");
 
diff --git a/infini_train/src/kernels/cuda/gelu.cu b/infini_train/src/kernels/cuda/gelu.cu
new file mode 100644
index 000000000..0aba0e409
--- /dev/null
+++ b/infini_train/src/kernels/cuda/gelu.cu
@@ -0,0 +1,219 @@
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <type_traits>
+
+#include "infini_train/include/common/common.h"
+#include "infini_train/include/common/cuda/kernel_helper.cuh"
+#include "infini_train/include/core/runtime/device_guard.h"
+#include "infini_train/include/datatype.h"
+#include "infini_train/include/device.h"
+#include "infini_train/include/dispatcher.h"
+#include "infini_train/include/tensor.h"
+
+#include "infini_train/src/core/runtime/cuda/cuda_dispatch.h"
+#include "infini_train/src/core/runtime/cuda/cuda_runtime_common.h"
+
+namespace infini_train::kernels::cuda {
+namespace {
+
+constexpr int kGeluThreadsPerBlock = 256;
+// Cap the grid so the vectorized kernels use a grid-stride loop on large tensors
+// instead of launching an unbounded number of blocks.
+constexpr size_t kGeluMaxBlocks = 4096;
+
+// Constants of the NewGELU tanh approximation (same as PyTorch's GELU(tanh) CUDA kernel):
+//   beta = sqrt(2/pi), kappa = 0.044715.
+template <typename T> constexpr T kGeluBeta = T(0.7978845608028654);
+template <typename T> constexpr T kGeluKappa = T(0.044715);
+
+// All arithmetic is carried out in opmath precision (double for double, float for the
+// float/half/bf16 cases), matching PyTorch. The single rounding to the storage dtype
+// happens at the output store (round-to-nearest-even via common::cuda::Cast).
+template <typename T> using GeluOpMath = std::conditional_t<std::is_same_v<T, double>, double, float>;
+
+// Forward: y = 0.5 * x * (1 + tanh(beta * (x + kappa * x^3)))
+template <typename T, typename OpMathT> __device__ __forceinline__ T NewGeluForwardOp(const T &x) {
+    const OpMathT xf = common::cuda::Cast<OpMathT>(x);
+    const OpMathT x_cube = xf * xf * xf;
+    const OpMathT inner = kGeluBeta<OpMathT> * (xf + kGeluKappa<OpMathT> * x_cube);
+    const OpMathT tanh_inner = common::cuda::Tanh(inner);
+    return common::cuda::Cast<T>(OpMathT(0.5) * xf * (OpMathT(1) + tanh_inner));
+}
+
+// Backward: dx = dy * [0.5 * (1 + t) + 0.5 * x * (1 - t^2) * beta * (1 + 3 * kappa * x^2)]
+// where t = tanh(beta * (x + kappa * x^3)). Same analytic derivative as PyTorch.
+template <typename T, typename OpMathT>
+__device__ __forceinline__ T NewGeluBackwardOp(const T &grad_output, const T &x) {
+    const OpMathT dy = common::cuda::Cast<OpMathT>(grad_output);
+    const OpMathT xf = common::cuda::Cast<OpMathT>(x);
+    const OpMathT x_sq = xf * xf;
+    const OpMathT inner = kGeluBeta<OpMathT> * (xf + kGeluKappa<OpMathT> * x_sq * xf);
+    const OpMathT tanh_inner = common::cuda::Tanh(inner);
+
+    const OpMathT left = OpMathT(0.5) * xf;
+    const OpMathT left_derivative = OpMathT(0.5) * (OpMathT(1) + tanh_inner);
+    const OpMathT tanh_derivative = OpMathT(1) - tanh_inner * tanh_inner;
+    const OpMathT inner_derivative = kGeluBeta<OpMathT> * (OpMathT(1) + OpMathT(3) * kGeluKappa<OpMathT> * x_sq);
+    const OpMathT right_derivative = left * tanh_derivative * inner_derivative;
+
+    return common::cuda::Cast<T>(dy * (left_derivative + right_derivative));
+}
+
+// Aligned vector type for vectorized loads/stores (128-bit).
+template <typename T, int N> struct __align__(sizeof(T) * N) aligned_vector { T val[N]; };
+
+// Scalar grid-stride fallback for misaligned buffers.
+template <typename T>
+__global__ void NewGeluForwardKernel(T *__restrict__ output, const T *__restrict__ input, size_t num_elements) {
+    const size_t stride = static_cast<size_t>(gridDim.x) * blockDim.x;
+    for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < num_elements; idx += stride) {
+        output[idx] = NewGeluForwardOp<T, GeluOpMath<T>>(input[idx]);
+    }
+}
+
+template <typename T>
+__global__ void NewGeluBackwardKernel(T *__restrict__ grad_input, const T *__restrict__ grad_output,
+                                      const T *__restrict__ input, size_t num_elements) {
+    const size_t stride = static_cast<size_t>(gridDim.x) * blockDim.x;
+    for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < num_elements; idx += stride) {
+        grad_input[idx] = NewGeluBackwardOp<T, GeluOpMath<T>>(grad_output[idx], input[idx]);
+    }
+}
+
+// Vectorized kernels: each thread processes kElems elements through one 128-bit load and
+// one 128-bit store. The trailing (< kElems) elements are handled by the first few threads.
+template <typename T, int kElems>
+__global__ void NewGeluForwardKernelVec(T *__restrict__ output, const T *__restrict__ input, size_t num_elements) {
+    using VecT = aligned_vector<T, kElems>;
+    const size_t num_vecs = num_elements / kElems;
+    const size_t stride = static_cast<size_t>(gridDim.x) * blockDim.x;
+    for (size_t v = blockIdx.x * blockDim.x + threadIdx.x; v < num_vecs; v += stride) {
+        const VecT in_vec = reinterpret_cast<const VecT *>(input)[v];
+        VecT out_vec;
+#pragma unroll
+        for (int i = 0; i < kElems; ++i) { out_vec.val[i] = NewGeluForwardOp<T, GeluOpMath<T>>(in_vec.val[i]); }
+        reinterpret_cast<VecT *>(output)[v] = out_vec;
+    }
+
+    const size_t tail_idx = num_vecs * kElems + blockIdx.x * blockDim.x + threadIdx.x;
+    if (tail_idx < num_elements) {
+        output[tail_idx] = NewGeluForwardOp<T, GeluOpMath<T>>(input[tail_idx]);
+    }
+}
+
+template <typename T, int kElems>
+__global__ void NewGeluBackwardKernelVec(T *__restrict__ grad_input, const T *__restrict__ grad_output,
+                                         const T *__restrict__ input, size_t num_elements) {
+    using VecT = aligned_vector<T, kElems>;
+    const size_t num_vecs = num_elements / kElems;
+    const size_t stride = static_cast<size_t>(gridDim.x) * blockDim.x;
+    for (size_t v = blockIdx.x * blockDim.x + threadIdx.x; v < num_vecs; v += stride) {
+        const VecT grad_vec = reinterpret_cast<const VecT *>(grad_output)[v];
+        const VecT in_vec = reinterpret_cast<const VecT *>(input)[v];
+        VecT out_vec;
+#pragma unroll
+        for (int i = 0; i < kElems; ++i) {
+            out_vec.val[i] = NewGeluBackwardOp<T, GeluOpMath<T>>(grad_vec.val[i], in_vec.val[i]);
+        }
+        reinterpret_cast<VecT *>(grad_input)[v] = out_vec;
+    }
+
+    const size_t tail_idx = num_vecs * kElems + blockIdx.x * blockDim.x + threadIdx.x;
+    if (tail_idx < num_elements) {
+        grad_input[tail_idx] = NewGeluBackwardOp<T, GeluOpMath<T>>(grad_output[tail_idx], input[tail_idx]);
+    }
+}
+
+inline size_t GeluGridSize(size_t work_items) {
+    return std::max<size_t>(std::min(CEIL_DIV(work_items, kGeluThreadsPerBlock), kGeluMaxBlocks), 1);
+}
+
+} // namespace
+
+std::shared_ptr<Tensor> NewGELUForward(const std::shared_ptr<Tensor> &input) {
+    CHECK(input->IsContiguous()) << "CUDA NewGELUForward: only contiguous input is supported";
+    auto output = std::make_shared<Tensor>(input->Dims(), input->Dtype(), input->GetDevice());
+    auto device = input->GetDevice();
+    const auto &cuda_stream = dynamic_cast<infini_train::core::cuda::CudaStream *>(
+                                  infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device))
+                                  ->cuda_stream();
+
+    const size_t num_elements = input->NumElements();
+    core::cuda::DispatchCudaFunc<INFINI_ALL_FLOATING_TYPES>(
+        input->Dtype(),
+        [=]<typename T>() {
+            T *output_ptr = static_cast<T *>(output->DataPtr());
+            const T *input_ptr = static_cast<const T *>(input->DataPtr());
+
+            // 128-bit vectorized path: float -> 4 elems, half/bf16 -> 8 elems, double -> 2 elems.
+            constexpr int kElems = static_cast<int>(16 / sizeof(T));
+            const bool aligned = (reinterpret_cast<uintptr_t>(output_ptr) % 16 == 0)
+                              && (reinterpret_cast<uintptr_t>(input_ptr) % 16 == 0);
+            if (aligned && num_elements >= static_cast<size_t>(kElems)) {
+                NewGeluForwardKernelVec<T, kElems>
+                    <<<static_cast<unsigned int>(GeluGridSize(num_elements / kElems)), kGeluThreadsPerBlock, 0,
+                       cuda_stream>>>(output_ptr, input_ptr, num_elements);
+            } else {
+                NewGeluForwardKernel<T>
+                    <<<static_cast<unsigned int>(GeluGridSize(num_elements)), kGeluThreadsPerBlock, 0, cuda_stream>>>(
+                        output_ptr, input_ptr, num_elements);
+            }
+        },
+        "CUDA NewGELUForward");
+
+    return output;
+}
+
+std::shared_ptr<Tensor> NewGELUBackward(const std::shared_ptr<Tensor> &grad_output,
+                                        const std::shared_ptr<Tensor> &input) {
+    CHECK(input->IsContiguous() && grad_output->IsContiguous())
+        << "CUDA NewGELUBackward: only contiguous tensors are supported";
+    CHECK_EQ(input->NumElements(), grad_output->NumElements());
+    // Forward output dtype equals input dtype, so the incoming grad matches it; convert the
+    // saved input in the rare case of a dtype-mismatched graph instead of failing.
+    auto input_matched
+        = input->Dtype() == grad_output->Dtype() ? input : std::make_shared<Tensor>(input->To(grad_output->Dtype()));
+
+    auto grad_input = std::make_shared<Tensor>(grad_output->Dims(), grad_output->Dtype(), grad_output->GetDevice());
+    auto device = grad_output->GetDevice();
+    const auto &cuda_stream = dynamic_cast<infini_train::core::cuda::CudaStream *>(
+                                  infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device))
+                                  ->cuda_stream();
+
+    const size_t num_elements = grad_output->NumElements();
+    core::cuda::DispatchCudaFunc<INFINI_ALL_FLOATING_TYPES>(
+        grad_output->Dtype(),
+        [=]<typename T>() {
+            T *grad_input_ptr = static_cast<T *>(grad_input->DataPtr());
+            const T *grad_output_ptr = static_cast<const T *>(grad_output->DataPtr());
+            const T *input_ptr = static_cast<const T *>(input_matched->DataPtr());
+
+            constexpr int kElems = static_cast<int>(16 / sizeof(T));
+            const bool aligned = (reinterpret_cast<uintptr_t>(grad_input_ptr) % 16 == 0)
+                              && (reinterpret_cast<uintptr_t>(grad_output_ptr) % 16 == 0)
+                              && (reinterpret_cast<uintptr_t>(input_ptr) % 16 == 0);
+            if (aligned && num_elements >= static_cast<size_t>(kElems)) {
+                NewGeluBackwardKernelVec<T, kElems>
+                    <<<static_cast<unsigned int>(GeluGridSize(num_elements / kElems)), kGeluThreadsPerBlock, 0,
+                       cuda_stream>>>(grad_input_ptr, grad_output_ptr, input_ptr, num_elements);
+            } else {
+                NewGeluBackwardKernel<T>
+                    <<<static_cast<unsigned int>(GeluGridSize(num_elements)), kGeluThreadsPerBlock, 0, cuda_stream>>>(
+                        grad_input_ptr, grad_output_ptr, input_ptr, num_elements);
+            }
+        },
+        "CUDA NewGELUBackward");
+
+    return grad_input;
+}
+} // namespace infini_train::kernels::cuda
+
+#define REGISTER_CUDA_GELU_KERNEL(kernel_name)                                                                         \
+    REGISTER_KERNEL(infini_train::Device::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name)
+
+REGISTER_CUDA_GELU_KERNEL(NewGELUForward)
+REGISTER_CUDA_GELU_KERNEL(NewGELUBackward)
+
+#undef REGISTER_CUDA_GELU_KERNEL
diff --git a/infini_train/src/kernels/cuda/linear.cu b/infini_train/src/kernels/cuda/linear.cu
index 1b4c18190..00cd82ecd 100644
--- a/infini_train/src/kernels/cuda/linear.cu
+++ b/infini_train/src/kernels/cuda/linear.cu
@@ -151,7 +151,7 @@ __global__ void ReduceColumnsKernel(const TIn *__restrict__ input, TOut *__restr
     float reduced = BlockReduce(temp_storage).Sum(sum);
 
     if (threadIdx.x == 0) {
-        output[row] = reduced;
+        output[row] = common::cuda::Cast<TOut>(reduced);
     }
 }
 
@@ -168,8 +168,10 @@ std::shared_ptr<Tensor> LinearBackwardInput(const std::shared_ptr<Tensor> &weigh
     auto grad_output_promoted
         = grad_output_dtype == compute_dtype ? grad_output : std::make_shared<Tensor>(grad_output->To(compute_dtype));
 
-    // FIXME(cx): output dtype promotion is a temporary hack; revisit when autograd/autocast is fixed.
-    auto output_dtype = (compute_dtype == DataType::kBFLOAT16) ? DataType::kFLOAT32 : compute_dtype;
+    // GEMM runs with fp32 compute but writes the gradient directly in the compute
+    // dtype (bf16 under autocast), matching PyTorch autocast backward semantics;
+    // the fp32 master-weight accumulation casts it back in AccumulateGrad.
+    auto output_dtype = compute_dtype;
     // No Fill(0) needed: cuBLAS beta=0.0f fully overwrites output.
     auto grad_input = std::make_shared<Tensor>(input_dims, output_dtype, grad_output->GetDevice());
 
@@ -240,8 +242,8 @@ std::shared_ptr<Tensor> LinearBackwardWeight(const std::shared_ptr<Tensor> &inpu
     auto grad_output_promoted
         = grad_output_dtype == compute_dtype ? grad_output : std::make_shared<Tensor>(grad_output->To(compute_dtype));
 
-    // FIXME(cx): output dtype promotion is a temporary hack; revisit when autograd/autocast is fixed.
-    auto output_dtype = (compute_dtype == DataType::kBFLOAT16) ? DataType::kFLOAT32 : compute_dtype;
+    // See LinearBackwardInput: bf16 GEMM writes the gradient directly in bf16.
+    auto output_dtype = compute_dtype;
     const std::vector<int64_t> weight_dims
         = transpose ? std::vector<int64_t>{out_features, in_features} : std::vector<int64_t>{in_features, out_features};
     // No Fill(0) needed: cuBLAS beta=0.0f fully overwrites output.
@@ -292,8 +294,9 @@ std::shared_ptr<Tensor> LinearBackwardBias(const std::shared_ptr<Tensor> &grad_o
     const int64_t bs = std::accumulate(dims.rbegin() + 1, dims.rend(), 1, std::multiplies<int64_t>{});
 
     auto compute_dtype = grad_output->Dtype();
-    // FIXME(cx): output dtype promotion is a temporary hack; revisit when autograd/autocast is fixed.
-    auto output_dtype = (compute_dtype == DataType::kBFLOAT16) ? DataType::kFLOAT32 : compute_dtype;
+    // Same as the GEMM backwards: emit the gradient in the compute dtype (bf16
+    // under autocast); the fp32 master-weight accumulation casts it back.
+    auto output_dtype = compute_dtype;
     auto grad_bias
         = std::make_shared<Tensor>(std::vector<int64_t>{out_features}, output_dtype, grad_output->GetDevice());
 
@@ -315,7 +318,7 @@ std::shared_ptr<Tensor> LinearBackwardBias(const std::shared_ptr<Tensor> &grad_o
         DISPATCH_CASE(WRAP({
                           ReduceColumnsKernel<BLOCK_SIZE><<<out_features, BLOCK_SIZE, 0, cuda_stream>>>(
                               static_cast<const nv_bfloat16 *>(grad_output->DataPtr()),
-                              static_cast<float *>(grad_bias->DataPtr()), out_features, bs);
+                              static_cast<nv_bfloat16 *>(grad_bias->DataPtr()), out_features, bs);
                       }),
                       DataType::kBFLOAT16)
     }
diff --git a/infini_train/src/kernels/cuda/matmul.cu b/infini_train/src/kernels/cuda/matmul.cu
index b1c6e381b..295d27709 100644
--- a/infini_train/src/kernels/cuda/matmul.cu
+++ b/infini_train/src/kernels/cuda/matmul.cu
@@ -105,8 +105,9 @@ std::shared_ptr<Tensor> MatmulBackwardInput(const std::shared_ptr<Tensor> &other
     auto grad_output_promoted
         = grad_output_dtype == compute_dtype ? grad_output : std::make_shared<Tensor>(grad_output->To(compute_dtype));
 
-    // For bf16 compute, output in fp32 to preserve accumulation precision.
-    auto output_dtype = (compute_dtype == DataType::kBFLOAT16) ? DataType::kFLOAT32 : compute_dtype;
+    // GEMM runs with fp32 compute but writes the gradient directly in the compute
+    // dtype (bf16 under autocast), matching PyTorch autocast backward semantics.
+    auto output_dtype = compute_dtype;
     auto grad_input = std::make_shared<Tensor>(input_dims, output_dtype, grad_output->GetDevice());
 
     // No Fill(0) needed: cuBLAS beta=0.0f means C is fully overwritten, never read.
@@ -175,8 +176,8 @@ std::shared_ptr<Tensor> MatmulBackwardOther(const std::shared_ptr<Tensor> &input
     auto grad_output_promoted
         = grad_output_dtype == compute_dtype ? grad_output : std::make_shared<Tensor>(grad_output->To(compute_dtype));
 
-    // For bf16 compute, output in fp32 to preserve accumulation precision.
-    auto output_dtype = (compute_dtype == DataType::kBFLOAT16) ? DataType::kFLOAT32 : compute_dtype;
+    // See MatmulBackwardInput: bf16 GEMM writes the gradient directly in bf16.
+    auto output_dtype = compute_dtype;
     auto grad_other = std::make_shared<Tensor>(other_dims, output_dtype, grad_output->GetDevice());
 
     // No Fill(0) needed: cuBLAS beta=0.0f means C is fully overwritten, never read.
diff --git a/infini_train/src/nn/modules/activations.cc b/infini_train/src/nn/modules/activations.cc
index d1bbc9da8..1d6f974d8 100644
--- a/infini_train/src/nn/modules/activations.cc
+++ b/infini_train/src/nn/modules/activations.cc
@@ -13,9 +13,7 @@ std::vector<std::shared_ptr<Tensor>> Sigmoid::Forward(const std::vector<std::sha
 }
 
 std::vector<std::shared_ptr<Tensor>> NewGELU::Forward(const std::vector<std::shared_ptr<Tensor>> &x) {
-    auto &input = x[0];
-    return {0.5 * input
-            * (1.0 + function::Tanh(std::sqrt(2.0 / M_PI) * (input + 0.044715 * function::Pow(input, 3.0))))};
+    return std::make_shared<autograd::NewGELU>()->Apply(x);
 }
 
 std::vector<std::shared_ptr<Tensor>> SwiGLU::Forward(const std::vector<std::shared_ptr<Tensor>> &x) {
diff --git a/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc
index 523bcf2d7..8a4cdd7ee 100644
--- a/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc
+++ b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc
@@ -2,6 +2,7 @@
 
 #include "glog/logging.h"
 
+#include "infini_train/include/autocast.h"
 #include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h"
 #include "infini_train/include/tensor.h"
 
@@ -182,6 +183,11 @@ void DistributedOptimizer::Step() {
     StartParamSync(/*force_sync=*/false);
     // TODO(zbl): Delay sync call until param is actually used in next step
     FinishParamSync(/*skip_next_bucket_dispatch=*/true);
+
+    // Shards are views into bucket storage, so the in-place base-optimizer update
+    // and the param-sync gather do not hit the per-tensor autocast cache
+    // invalidation hooks of the full params. Invalidate them explicitly.
+    for (const auto &param : params_) { InvalidateAutocastWeightCacheEntry(param.get()); }
 }
 
 std::unordered_map<std::string, std::shared_ptr<Tensor>> DistributedOptimizer::StateDict() const {
diff --git a/infini_train/src/optimizer.cc b/infini_train/src/optimizer.cc
index 39b999c77..19dc082e1 100644
--- a/infini_train/src/optimizer.cc
+++ b/infini_train/src/optimizer.cc
@@ -3,6 +3,7 @@
 #include <unordered_map>
 #include <vector>
 
+#include "infini_train/include/autocast.h"
 #include "infini_train/include/core/runtime/device_guard.h"
 #include "infini_train/include/device.h"
 #include "infini_train/include/dispatcher.h"
@@ -63,6 +64,8 @@ void SGD::Step() {
         core::DeviceGuard guard(device);
         auto kernel = Dispatcher::Instance().GetKernel({device.type(), "AccumulateGrad"});
         kernel.Call<void>(param->grad(), -learning_rate_, param);
+        // The parameter was updated in place; any autocast-cached demoted copy is stale.
+        InvalidateAutocastWeightCacheEntry(param.get());
     }
 }
 
@@ -116,6 +119,8 @@ void Adam::Step() {
         core::DeviceGuard guard(device);
         auto kernel = Dispatcher::Instance().GetKernel({device.type(), "AdamAccumulateGrad"});
         kernel.Call<void>(grad, param, m, v, learning_rate_, beta1_, beta2_, eps_, t_);
+        // The parameter was updated in place; any autocast-cached demoted copy is stale.
+        InvalidateAutocastWeightCacheEntry(param.get());
     }
 }
 
diff --git a/infini_train/src/tensor.cc b/infini_train/src/tensor.cc
index 18ca3d22b..0c848d052 100644
--- a/infini_train/src/tensor.cc
+++ b/infini_train/src/tensor.cc
@@ -9,6 +9,7 @@
 #include "Eigen/Dense"
 #include "glog/logging.h"
 
+#include "infini_train/include/autocast.h"
 #include "infini_train/include/autograd/accumulate.h"
 #include "infini_train/include/autograd/elementwise.h"
 #include "infini_train/include/autograd/function.h"
@@ -79,6 +80,9 @@ void Tensor::SetData(const Tensor &tensor, size_t offset, bool preserve_data) {
     CHECK(tensor.Dtype() == Dtype());
     CHECK_LE(tensor.offset_ + offset + SizeInBytes(), tensor.buffer_->Size());
 
+    // The storage backing this tensor is being replaced/rebound.
+    InvalidateAutocastWeightCacheEntry(this);
+
     if (preserve_data) {
         // Create a view of original tensor buffer
         auto new_tensor = Tensor(tensor, offset, Dims());
@@ -107,6 +111,7 @@ DataType Tensor::Dtype() const { return dtype_; }
 std::shared_ptr<Tensor> Tensor::Detach() const { return std::make_shared<Tensor>(*this, 0, dims_); }
 
 void Tensor::Fill(Scalar value) {
+    InvalidateAutocastWeightCacheEntry(this);
     auto device = GetDevice();
     core::DeviceGuard guard(device);
     auto kernel = Dispatcher::Instance().GetKernel({device.type(), "Fill"});
@@ -202,6 +207,9 @@ void Tensor::CopyFrom(const Tensor &src) {
     CHECK_EQ(NumElements(), src.NumElements()) << "Tensor::CopyFrom element count mismatch";
     CHECK(Dims() == src.Dims()) << "Tensor::CopyFrom shape mismatch";
 
+    // In-place overwrite of this tensor's data (also used by checkpoint loading).
+    InvalidateAutocastWeightCacheEntry(this);
+
     const size_t nbytes = SizeInBytes();
     const Device dst_dev = GetDevice();
     const Device src_dev = src.GetDevice();
diff --git a/tests/autograd/test_autograd_activations.cc b/tests/autograd/test_autograd_activations.cc
new file mode 100644
index 000000000..82b3eb9f8
--- /dev/null
+++ b/tests/autograd/test_autograd_activations.cc
@@ -0,0 +1,31 @@
+#include <memory>
+#include <vector>
+
+#include "gtest/gtest.h"
+
+#include "infini_train/include/autograd/activations.h"
+#include "infini_train/include/tensor.h"
+
+#include "tests/common/test_utils.h"
+
+using namespace infini_train;
+
+class AutogradActivationTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(AutogradActivationTest, NewGELUForwardAndBackward) {
+    auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
+    input->Fill(1.0f);
+
+    auto gelu = std::make_shared<autograd::NewGELU>();
+    auto outputs = gelu->Apply({input});
+    ASSERT_EQ(outputs.size(), 1);
+    test::ExpectTensorNear(outputs[0], 0.84119199f, 1e-6f);
+
+    auto grad_output = std::make_shared<Tensor>(input->Dims(), DataType::kFLOAT32, GetDevice());
+    grad_output->Fill(1.0f);
+    auto grad_inputs = gelu->Backward({grad_output});
+    ASSERT_EQ(grad_inputs.size(), 1);
+    test::ExpectTensorNear(grad_inputs[0], 1.08296408f, 1e-6f);
+}
+
+INFINI_TRAIN_REGISTER_TEST(AutogradActivationTest);