diff --git a/example/common/tiny_shakespeare_dataset.cc b/example/common/tiny_shakespeare_dataset.cc
index 3bc5f1b..20daa16 100644
--- a/example/common/tiny_shakespeare_dataset.cc
+++ b/example/common/tiny_shakespeare_dataset.cc
@@ -3,6 +3,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <cstdlib>
+#include <cstring>
 #include <filesystem>
 #include <fstream>
 #include <functional>
@@ -61,15 +62,52 @@ TinyShakespeareFile ReadTinyShakespeareFile(const std::string &path, size_t sequ
     | magic(4B) | version(4B) | num_toks(4B) | reserved(1012B) | token数据           |
     ----------------------------------------------------------------------------------
        =================================== 作业 =================================== */
+    std::ifstream ifs(path, std::ios::binary);
+    CHECK(ifs.is_open()) << "Failed to open file: " << path;
+
+    auto header_bytes = ReadSeveralBytesFromIfstream(1024, &ifs);
+    const uint32_t magic = BytesToType<uint32_t>(header_bytes, 0);
+    const uint32_t version = BytesToType<uint32_t>(header_bytes, 4);
+    const uint32_t num_toks = BytesToType<uint32_t>(header_bytes, 8);
+    (void)version;
+
+    CHECK(kTypeMap.contains(magic)) << "Unsupported magic number: " << magic;
+    const TinyShakespeareType type = kTypeMap.at(magic);
+    const size_t token_size = kTypeToSize.at(type);
+
+    CHECK_GT(num_toks, sequence_length) << "Not enough tokens for one sequence";
+    infini_train::Tensor tensor(std::vector<int64_t>{static_cast<int64_t>(num_toks)}, DataType::kINT64);
+    auto *dst = static_cast<int64_t *>(tensor.DataPtr());
+
+    if (token_size == 2) {
+        std::vector<uint16_t> tokens(num_toks);
+        ifs.read(reinterpret_cast<char *>(tokens.data()), static_cast<std::streamsize>(num_toks * token_size));
+        for (uint32_t i = 0; i < num_toks; ++i) {
+            dst[i] = static_cast<int64_t>(tokens[i]);
+        }
+    } else {
+        std::vector<uint32_t> tokens(num_toks);
+        ifs.read(reinterpret_cast<char *>(tokens.data()), static_cast<std::streamsize>(num_toks * token_size));
+        for (uint32_t i = 0; i < num_toks; ++i) {
+            dst[i] = static_cast<int64_t>(tokens[i]);
+        }
+    }
+
+    // dims[0] is used by operator[] as CHECK_LT(idx, dims[0] - 1)
+    // Non-overlapping windows (llm.c semantics): sample idx covers tokens[idx*seq_len : (idx+1)*seq_len]
+    const int64_t num_samples = static_cast<int64_t>(num_toks) / static_cast<int64_t>(sequence_length);
+    std::vector<int64_t> dims{num_samples + 1, static_cast<int64_t>(sequence_length)};
+    return TinyShakespeareFile{type, dims, std::move(tensor)};
 }
 } // namespace
 
-TinyShakespeareDataset::TinyShakespeareDataset(const std::string &filepath, size_t sequence_length) {
+TinyShakespeareDataset::TinyShakespeareDataset(const std::string &filepath, size_t sequence_length)
     // =================================== 作业 ===================================
     // TODO：初始化数据集实例
     // HINT: 调用ReadTinyShakespeareFile加载数据文件
     // =================================== 作业 ===================================
-}
+    : text_file_(ReadTinyShakespeareFile(filepath, sequence_length)), sequence_length_(sequence_length),
+      sequence_size_in_bytes_(sequence_length * sizeof(int64_t)), num_samples_(text_file_.dims[0] - 1) {}
 
 std::pair<std::shared_ptr<infini_train::Tensor>, std::shared_ptr<infini_train::Tensor>>
 TinyShakespeareDataset::operator[](size_t idx) const {
diff --git a/example/common/tokenizer.cc b/example/common/tokenizer.cc
index 23b9537..9b230b2 100644
--- a/example/common/tokenizer.cc
+++ b/example/common/tokenizer.cc
@@ -2,9 +2,11 @@
 
 #include <cctype>
 #include <cstdint>
+#include <cstring>
 #include <filesystem>
 #include <fstream>
 #include <iostream>
+#include <unordered_map>
 #include <vector>
 
 #include "glog/logging.h"
@@ -78,6 +80,26 @@ Tokenizer::Tokenizer(const std::string &filepath) {
     | magic(4B) | version(4B) | vocab_size(4B) | reserved(1012B) | token词表数据       |
     ----------------------------------------------------------------------------------
     ===================================== 作业 ===================================== */
+    std::ifstream ifs(filepath, std::ios::binary);
+    CHECK(ifs.is_open()) << "Failed to open tokenizer file: " << filepath;
+
+    auto header_bytes = ReadSeveralBytesFromIfstream(1024, &ifs);
+    magic_number_ = BytesToType<uint32_t>(header_bytes, 0);
+    const uint32_t version = BytesToType<uint32_t>(header_bytes, 4);
+    vocab_size_ = BytesToType<uint32_t>(header_bytes, 8);
+    (void)version;
+
+    CHECK(kEotMap.contains(magic_number_)) << "Unsupported tokenizer magic: " << magic_number_;
+    eot_token_ = kEotMap.at(magic_number_);
+
+    token_table_.resize(vocab_size_);
+    for (uint32_t i = 0; i < vocab_size_; ++i) {
+        uint8_t len = 0;
+        ifs.read(reinterpret_cast<char *>(&len), 1);
+        std::string token(len, '\0');
+        ifs.read(token.data(), len);
+        token_table_[i] = std::move(token);
+    }
 }
 
 std::string Tokenizer::Decode(uint32_t token_id) const {
@@ -85,6 +107,9 @@ std::string Tokenizer::Decode(uint32_t token_id) const {
     TODO：实现token_id到文本的转换
     功能描述：根据token_id返回对应的文本片段
     ===================================== 作业 ===================================== */
+    if (token_id < token_table_.size()) {
+        return token_table_[token_id];
+    }
     return "";
 }
 
@@ -104,13 +129,36 @@ void Tokenizer::GenerateText(infini_train::nn::Module &model, uint32_t batch_siz
     std::cout << "The meaning of life is";
 
     auto x = std::make_shared<infini_train::Tensor>(x_tensor.To(device));
-    uint64_t kRngState = kRngState;
+    uint64_t rng_state = kRngState;
     LOG(INFO) << "start generate text:";
-    for (int t = prompt_len; t < text_length; t++) {
+    for (int t = prompt_len; t < static_cast<int>(text_length); t++) {
         /* ===================================== 作业 =====================================
         TODO：实现单步文本生成逻辑
         HINT：调用model.Forward推理获取logits，根据推理结果进行随机采样，调用Decode获取文本结果
         ===================================== 作业 ===================================== */
+        auto outputs = model.Forward({x});
+        auto logits = outputs[0];
+        const int last_pos = (t < static_cast<int>(sequence_length)) ? (t - 1) : (static_cast<int>(sequence_length) - 1);
+        auto last_logits = logits->Slice(1, last_pos, last_pos + 1, 1)->Squeeze(1);
+        auto probs = nn::function::Softmax(last_logits, -1)->To(Device(DeviceType::kCPU, 0));
+        float *probs_ptr = static_cast<float *>(probs.DataPtr());
+
+        for (uint32_t b = 0; b < batch_size; ++b) {
+            const float coin = RandomF32(rng_state);
+            const int next_token = SampleMult(probs_ptr + b * vocab_size_, static_cast<int>(vocab_size_), coin);
+            if (b == 0) {
+                std::cout << Decode(static_cast<uint32_t>(next_token)) << std::flush;
+            }
+            if (t < static_cast<int>(sequence_length)) {
+                x_buff[b * sequence_length + t] = next_token;
+            } else {
+                for (uint32_t i = 0; i + 1 < sequence_length; ++i) {
+                    x_buff[b * sequence_length + i] = x_buff[b * sequence_length + i + 1];
+                }
+                x_buff[b * sequence_length + sequence_length - 1] = next_token;
+            }
+        }
+        x = std::make_shared<infini_train::Tensor>(x_tensor.To(device));
     }
     std::cout << std::endl;
 }
diff --git a/example/gpt2/net.cc b/example/gpt2/net.cc
index 441b121..9ced3ad 100644
--- a/example/gpt2/net.cc
+++ b/example/gpt2/net.cc
@@ -22,6 +22,27 @@
 #include "infini_train/include/nn/modules/sparse.h"
 #include "infini_train/include/tensor.h"
 
+#include <string>
+#include <type_traits>
+
+namespace {
+template <typename T>
+std::string ToStr(const T &v) {
+    if constexpr (std::is_arithmetic_v<T>) {
+        return std::to_string(v);
+    } else {
+        return std::string(v);
+    }
+}
+template <typename... Args>
+std::string JoinDot(Args &&...args) {
+    std::string out;
+    size_t i = 0;
+    ((out += (i++ ? "." : "") + ToStr(args)), ...);
+    return out;
+}
+} // namespace
+
 namespace nn = infini_train::nn;
 
 namespace {
@@ -272,52 +293,52 @@ std::unique_ptr<GPT2> GPT2::FromLLMC(const std::string &filepath) {
     auto state_dict = gpt2->StateDict();
     // transformer.wte.weight
     // (padded_vocab_size, n_embd) -> un_pad -> (vocab_size, n_embd)
-    auto &transformer_wte_weight = state_dict[std::format("{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kWTELayerName,
+    auto &transformer_wte_weight = state_dict[JoinDot(GPT2::kTransformerLayerName, GPT2::kWTELayerName,
                                                           nn::Embedding::kParamWeightName)];
     ifs.read(reinterpret_cast<char *>(transformer_wte_weight->DataPtr()), transformer_wte_weight->SizeInBytes());
     ifs.ignore((padded_vocab_size - vocab_size) * n_embd * sizeof(float));
     // transformer.wpe.weight
-    auto &transformer_wpe_weight = state_dict[std::format("{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kWPELayerName,
+    auto &transformer_wpe_weight = state_dict[JoinDot(GPT2::kTransformerLayerName, GPT2::kWPELayerName,
                                                           nn::Embedding::kParamWeightName)];
     ifs.read(reinterpret_cast<char *>(transformer_wpe_weight->DataPtr()), transformer_wpe_weight->SizeInBytes());
     // transformer.h.{i}.ln_1.weight
     for (int idx = 0; idx < n_layer; idx++) {
         auto &tensor
-            = state_dict[std::format("{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName,
+            = state_dict[JoinDot(GPT2::kTransformerLayerName, GPT2::kHLayerName,
                                      std::to_string(idx), Block::kLn1LayerName, nn::LayerNorm::kParamWeightName)];
         ifs.read(reinterpret_cast<char *>(tensor->DataPtr()), tensor->SizeInBytes());
     }
     // transformer.h.{i}.ln_1.bias
     for (int idx = 0; idx < n_layer; idx++) {
         auto &tensor
-            = state_dict[std::format("{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName,
+            = state_dict[JoinDot(GPT2::kTransformerLayerName, GPT2::kHLayerName,
                                      std::to_string(idx), Block::kLn1LayerName, nn::LayerNorm::kParamBiasName)];
         ifs.read(reinterpret_cast<char *>(tensor->DataPtr()), tensor->SizeInBytes());
     }
     // transformer.h.{i}.attn.c_attn.weight
     for (int idx = 0; idx < n_layer; idx++) {
-        auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName,
+        auto &tensor = state_dict[JoinDot(GPT2::kTransformerLayerName, GPT2::kHLayerName,
                                               std::to_string(idx), Block::kAttnLayerName,
                                               CausalSelfAttention::kCAttnLayerName, GPT2Linear::kParamWeightName)];
         ifs.read(reinterpret_cast<char *>(tensor->DataPtr()), tensor->SizeInBytes());
     }
     // transformer.h.{i}.attn.c_attn.bias
     for (int idx = 0; idx < n_layer; idx++) {
-        auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName,
+        auto &tensor = state_dict[JoinDot(GPT2::kTransformerLayerName, GPT2::kHLayerName,
                                               std::to_string(idx), Block::kAttnLayerName,
                                               CausalSelfAttention::kCAttnLayerName, GPT2Linear::kParamBiasName)];
         ifs.read(reinterpret_cast<char *>(tensor->DataPtr()), tensor->SizeInBytes());
     }
     // transformer.h.{i}.attn.c_proj.weight
     for (int idx = 0; idx < n_layer; idx++) {
-        auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName,
+        auto &tensor = state_dict[JoinDot(GPT2::kTransformerLayerName, GPT2::kHLayerName,
                                               std::to_string(idx), Block::kAttnLayerName,
                                               CausalSelfAttention::kCProjLayerName, GPT2Linear::kParamWeightName)];
         ifs.read(reinterpret_cast<char *>(tensor->DataPtr()), tensor->SizeInBytes());
     }
     // transformer.h.{i}.attn.c_proj.bias
     for (int idx = 0; idx < n_layer; idx++) {
-        auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName,
+        auto &tensor = state_dict[JoinDot(GPT2::kTransformerLayerName, GPT2::kHLayerName,
                                               std::to_string(idx), Block::kAttnLayerName,
                                               CausalSelfAttention::kCProjLayerName, GPT2Linear::kParamBiasName)];
         ifs.read(reinterpret_cast<char *>(tensor->DataPtr()), tensor->SizeInBytes());
@@ -325,51 +346,51 @@ std::unique_ptr<GPT2> GPT2::FromLLMC(const std::string &filepath) {
     // transformer.h.{i}.ln_2.weight
     for (int idx = 0; idx < n_layer; idx++) {
         auto &tensor
-            = state_dict[std::format("{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName,
+            = state_dict[JoinDot(GPT2::kTransformerLayerName, GPT2::kHLayerName,
                                      std::to_string(idx), Block::kLn2LayerName, nn::LayerNorm::kParamWeightName)];
         ifs.read(reinterpret_cast<char *>(tensor->DataPtr()), tensor->SizeInBytes());
     }
     // transformer.h.{i}.ln_2.bias
     for (int idx = 0; idx < n_layer; idx++) {
         auto &tensor
-            = state_dict[std::format("{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName,
+            = state_dict[JoinDot(GPT2::kTransformerLayerName, GPT2::kHLayerName,
                                      std::to_string(idx), Block::kLn2LayerName, nn::LayerNorm::kParamBiasName)];
         ifs.read(reinterpret_cast<char *>(tensor->DataPtr()), tensor->SizeInBytes());
     }
     // transformer.h.{i}.mlp.c_fc.weight
     for (int idx = 0; idx < n_layer; idx++) {
-        auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName,
+        auto &tensor = state_dict[JoinDot(GPT2::kTransformerLayerName, GPT2::kHLayerName,
                                               std::to_string(idx), Block::kMlpLayerName, MLP::kCFclayerName,
                                               GPT2Linear::kParamWeightName)];
         ifs.read(reinterpret_cast<char *>(tensor->DataPtr()), tensor->SizeInBytes());
     }
     // transformer.h.{i}.mlp.c_fc.bias
     for (int idx = 0; idx < n_layer; idx++) {
-        auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName,
+        auto &tensor = state_dict[JoinDot(GPT2::kTransformerLayerName, GPT2::kHLayerName,
                                               std::to_string(idx), Block::kMlpLayerName, MLP::kCFclayerName,
                                               GPT2Linear::kParamBiasName)];
         ifs.read(reinterpret_cast<char *>(tensor->DataPtr()), tensor->SizeInBytes());
     }
     // transformer.h.{i}.mlp.c_proj.weight
     for (int idx = 0; idx < n_layer; idx++) {
-        auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName,
+        auto &tensor = state_dict[JoinDot(GPT2::kTransformerLayerName, GPT2::kHLayerName,
                                               std::to_string(idx), Block::kMlpLayerName, MLP::kCProjLayerName,
                                               GPT2Linear::kParamWeightName)];
         ifs.read(reinterpret_cast<char *>(tensor->DataPtr()), tensor->SizeInBytes());
     }
     // transformer.h.{i}.mlp.c_proj.bias
     for (int idx = 0; idx < n_layer; idx++) {
-        auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName,
+        auto &tensor = state_dict[JoinDot(GPT2::kTransformerLayerName, GPT2::kHLayerName,
                                               std::to_string(idx), Block::kMlpLayerName, MLP::kCProjLayerName,
                                               GPT2Linear::kParamBiasName)];
         ifs.read(reinterpret_cast<char *>(tensor->DataPtr()), tensor->SizeInBytes());
     }
     // transformer.ln_f.weight
-    auto &transformer_ln_f_weight = state_dict[std::format("{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kLnFLayerName,
+    auto &transformer_ln_f_weight = state_dict[JoinDot(GPT2::kTransformerLayerName, GPT2::kLnFLayerName,
                                                            nn::LayerNorm::kParamWeightName)];
     ifs.read(reinterpret_cast<char *>(transformer_ln_f_weight->DataPtr()), transformer_ln_f_weight->SizeInBytes());
     // transformer.ln_f.bias
-    auto &transformer_ln_f_bias = state_dict[std::format("{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kLnFLayerName,
+    auto &transformer_ln_f_bias = state_dict[JoinDot(GPT2::kTransformerLayerName, GPT2::kLnFLayerName,
                                                          nn::LayerNorm::kParamBiasName)];
     ifs.read(reinterpret_cast<char *>(transformer_ln_f_bias->DataPtr()), transformer_ln_f_bias->SizeInBytes());
 
diff --git a/infini_train/include/dispatcher.h b/infini_train/include/dispatcher.h
index 5b91d85..991e4b3 100644
--- a/infini_train/include/dispatcher.h
+++ b/infini_train/include/dispatcher.h
@@ -21,7 +21,8 @@ class KernelFunction {
         // =================================== 作业 ===================================
 
         using FuncT = RetT (*)(ArgsT...);
-        // TODO: 实现函数调用逻辑
+        auto func = reinterpret_cast<FuncT>(func_ptr_);
+        return func(std::forward<ArgsT>(args)...);
     }
 
 private:
@@ -48,6 +49,9 @@ class Dispatcher {
         // TODO：实现kernel注册机制
         // 功能描述：将kernel函数与设备类型、名称绑定
         // =================================== 作业 ===================================
+        CHECK(!key_to_kernel_map_.contains(key))
+            << "Kernel already registered: " << key.second << " on device: " << static_cast<int>(key.first);
+        key_to_kernel_map_.emplace(key, KernelFunction(std::forward<FuncT>(kernel)));
     }
 
 private:
@@ -55,8 +59,11 @@ class Dispatcher {
 };
 } // namespace infini_train
 
+#define REGISTER_KERNEL_INTERNAL(device, kernel_name, kernel_func, line)                                               \
+    static const bool register_##kernel_name##_##line [[maybe_unused]] = []() {                                        \
+        ::infini_train::Dispatcher::Instance().Register({device, #kernel_name}, kernel_func);                          \
+        return true;                                                                                                   \
+    }();
+
 #define REGISTER_KERNEL(device, kernel_name, kernel_func)                                                              \
-    // =================================== 作业 ===================================
-    // TODO：实现自动注册宏
-    // 功能描述：在全局静态区注册kernel，避免显式初始化代码
-    // =================================== 作业 ===================================
+    REGISTER_KERNEL_INTERNAL(device, kernel_name, kernel_func, __LINE__)
diff --git a/infini_train/src/autograd/elementwise.cc b/infini_train/src/autograd/elementwise.cc
index 5a790a5..65d65fb 100644
--- a/infini_train/src/autograd/elementwise.cc
+++ b/infini_train/src/autograd/elementwise.cc
@@ -10,8 +10,12 @@ std::vector<std::shared_ptr<Tensor>> Neg::Forward(const std::vector<std::shared_
     // TODO：通过Dispatcher获取设备专属kernel，对输入张量进行取反操作
     // NOTES: 依赖test_dispatcher，Neg kernel实现已给出
     // =================================== 作业 ===================================
+    CHECK_EQ(input_tensors.size(), 1);
+    const auto &input = input_tensors[0];
 
-    return std::vector<std::shared_ptr<Tensor>>();
+    auto device = input->GetDevice().Type();
+    auto kernel = Dispatcher::Instance().GetKernel({device, "NegForward"});
+    return {kernel.Call<std::shared_ptr<Tensor>>(input)};
 }
 
 std::vector<std::shared_ptr<Tensor>> Neg::Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) {
@@ -19,8 +23,12 @@ std::vector<std::shared_ptr<Tensor>> Neg::Backward(const std::vector<std::shared
     // TODO：通过Dispatcher获取设备专属的反向传播kernel，计算梯度
     // NOTES: 依赖test_dispatcher，Neg的kernel实现已给出
     // =================================== 作业 ===================================
+    CHECK_EQ(grad_outputs.size(), 1);
+    const auto &grad_output = grad_outputs[0];
 
-    return std::vector<std::shared_ptr<Tensor>>();
+    auto device = grad_output->GetDevice().Type();
+    auto kernel = Dispatcher::Instance().GetKernel({device, "NegBackward"});
+    return {kernel.Call<std::shared_ptr<Tensor>>(grad_output)};
 }
 
 std::vector<std::shared_ptr<Tensor>> Reciprocal::Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {
diff --git a/infini_train/src/kernels/cpu/accumulate_grad.cc b/infini_train/src/kernels/cpu/accumulate_grad.cc
index 55637cd..0e1a214 100644
--- a/infini_train/src/kernels/cpu/accumulate_grad.cc
+++ b/infini_train/src/kernels/cpu/accumulate_grad.cc
@@ -1,3 +1,4 @@
+#include <cmath>
 #include <cstddef>
 #include <memory>
 
@@ -18,6 +19,19 @@ void AdamAccumulateGrad(const std::shared_ptr<Tensor> &grad, const std::shared_p
     // TODO：实现Adam优化器的梯度累积和参数更新
     // REF:
     // =================================== 作业 ===================================
+    float *g_ptr = static_cast<float *>(grad->DataPtr());
+    float *p_ptr = static_cast<float *>(param->DataPtr());
+    float *m_ptr = static_cast<float *>(m->DataPtr());
+    float *v_ptr = static_cast<float *>(v->DataPtr());
+
+    const float step_size
+        = learning_rate * std::sqrt(1.0f - std::pow(beta2, static_cast<float>(t))) / (1.0f - std::pow(beta1, static_cast<float>(t)));
+
+    for (int64_t i = 0; i < grad->NumElements(); ++i) {
+        m_ptr[i] = beta1 * m_ptr[i] + (1.0f - beta1) * g_ptr[i];
+        v_ptr[i] = beta2 * v_ptr[i] + (1.0f - beta2) * g_ptr[i] * g_ptr[i];
+        p_ptr[i] -= step_size * m_ptr[i] / (std::sqrt(v_ptr[i]) + eps);
+    }
 }
 
 } // namespace infini_train::kernels::cpu
diff --git a/infini_train/src/kernels/cpu/linear.cc b/infini_train/src/kernels/cpu/linear.cc
index 140e756..c11cdf3 100644
--- a/infini_train/src/kernels/cpu/linear.cc
+++ b/infini_train/src/kernels/cpu/linear.cc
@@ -11,25 +11,111 @@
 
 namespace infini_train::kernels::cpu {
 std::shared_ptr<Tensor> MatmulForward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Tensor> &other) {
-    // =================================== 作业 ===================================
-    // TODO：实现CPU上的矩阵乘法前向计算
-    // REF:
-    // =================================== 作业 ===================================
+    const auto &input_dims = input->Dims();
+    const auto &other_dims = other->Dims();
+
+    CHECK_GE(input_dims.size(), 2);
+    CHECK_GE(other_dims.size(), 2);
+
+    int64_t M = input_dims[input_dims.size() - 2];
+    int64_t K = input_dims.back();
+    int64_t K2 = other_dims[other_dims.size() - 2];
+    int64_t N = other_dims.back();
+    CHECK_EQ(K, K2);
+
+    int64_t input_batches = 1;
+    for (size_t i = 0; i < input_dims.size() - 2; ++i) {
+        input_batches *= input_dims[i];
+    }
+    int64_t other_batches = 1;
+    for (size_t i = 0; i < other_dims.size() - 2; ++i) {
+        other_batches *= other_dims[i];
+    }
+
+    int64_t max_batches = std::max(input_batches, other_batches);
+    std::vector<int64_t> output_dims = (input_dims.size() >= other_dims.size()) ? input_dims : other_dims;
+    output_dims[output_dims.size() - 2] = M;
+    output_dims.back() = N;
 
-    auto output = std::make_shared<Tensor>();
-    return {output};
+    auto output = std::make_shared<Tensor>(output_dims, DataType::kFLOAT32);
+
+    float *input_ptr = static_cast<float *>(input->DataPtr());
+    float *other_ptr = static_cast<float *>(other->DataPtr());
+    float *output_ptr = static_cast<float *>(output->DataPtr());
+
+    for (int64_t b = 0; b < max_batches; ++b) {
+        float *cur_input_ptr = input_ptr + (input_batches == 1 ? 0 : b * M * K);
+        float *cur_other_ptr = other_ptr + (other_batches == 1 ? 0 : b * K * N);
+        float *cur_output_ptr = output_ptr + b * M * N;
+
+        auto input_map = Eigen::Map<const Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(
+            cur_input_ptr, M, K);
+        auto other_map = Eigen::Map<const Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(
+            cur_other_ptr, K, N);
+        auto output_map = Eigen::Map<Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(
+            cur_output_ptr, M, N);
+        output_map = input_map * other_map;
+    }
+    return output;
 }
 
 std::tuple<std::shared_ptr<Tensor>, std::shared_ptr<Tensor>>
 MatmulBackward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Tensor> &other,
                const std::shared_ptr<Tensor> &grad_output) {
-    // =================================== 作业 ===================================
-    // TODO：实现CPU上的矩阵乘法反向传播
-    // REF:
-    // =================================== 作业 ===================================
+    const auto &input_dims = input->Dims();
+    const auto &other_dims = other->Dims();
+    const auto &grad_output_dims = grad_output->Dims();
+
+    int64_t M = input_dims[input_dims.size() - 2];
+    int64_t K = input_dims.back();
+    int64_t N = other_dims.back();
+
+    int64_t input_batches = 1;
+    for (size_t i = 0; i < input_dims.size() - 2; ++i)
+        input_batches *= input_dims[i];
+    int64_t other_batches = 1;
+    for (size_t i = 0; i < other_dims.size() - 2; ++i)
+        other_batches *= other_dims[i];
+    int64_t grad_batches = 1;
+    for (size_t i = 0; i < grad_output_dims.size() - 2; ++i)
+        grad_batches *= grad_output_dims[i];
+
+    auto grad_input = std::make_shared<Tensor>(input_dims, DataType::kFLOAT32);
+    auto grad_other = std::make_shared<Tensor>(other_dims, DataType::kFLOAT32);
+    grad_input->Fill<float>(0.0f);
+    grad_other->Fill<float>(0.0f);
+
+    float *input_ptr = static_cast<float *>(input->DataPtr());
+    float *other_ptr = static_cast<float *>(other->DataPtr());
+    float *grad_output_ptr = static_cast<float *>(grad_output->DataPtr());
+    float *grad_input_ptr = static_cast<float *>(grad_input->DataPtr());
+    float *grad_other_ptr = static_cast<float *>(grad_other->DataPtr());
+
+    for (int64_t b = 0; b < grad_batches; ++b) {
+        auto grad_output_map = Eigen::Map<const Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(
+            grad_output_ptr + b * M * N, M, N);
+
+        if (grad_input_ptr) {
+            float *cur_grad_input_ptr = grad_input_ptr + (input_batches == 1 ? 0 : b * M * K);
+            float *cur_other_ptr = other_ptr + (other_batches == 1 ? 0 : b * K * N);
+            auto other_map = Eigen::Map<const Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(
+                cur_other_ptr, K, N);
+            auto grad_input_map = Eigen::Map<Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(
+                cur_grad_input_ptr, M, K);
+            grad_input_map += grad_output_map * other_map.transpose();
+        }
+
+        if (grad_other_ptr) {
+            float *cur_grad_other_ptr = grad_other_ptr + (other_batches == 1 ? 0 : b * K * N);
+            float *cur_input_ptr = input_ptr + (input_batches == 1 ? 0 : b * M * K);
+            auto input_map = Eigen::Map<const Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(
+                cur_input_ptr, M, K);
+            auto grad_other_map = Eigen::Map<Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(
+                cur_grad_other_ptr, K, N);
+            grad_other_map += input_map.transpose() * grad_output_map;
+        }
+    }
 
-    auto grad_input = std::make_shared<Tensor>();
-    auto grad_other = std::make_shared<Tensor>();
     return {grad_input, grad_other};
 }
 
diff --git a/infini_train/src/kernels/cuda/accumulate_grad.cu b/infini_train/src/kernels/cuda/accumulate_grad.cu
index 5f977c3..aac65ed 100644
--- a/infini_train/src/kernels/cuda/accumulate_grad.cu
+++ b/infini_train/src/kernels/cuda/accumulate_grad.cu
@@ -22,6 +22,18 @@ void AccumulateGrad(const std::shared_ptr<Tensor> &gradient, float rate, const s
     AccumulateGradKernel<<<num_blocks, threads_per_block>>>(grad_ptr, rate, tensor_ptr, num_elements);
 }
 
+__global__ void AdamAccumulateGradKernel(float *grad_ptr, float *param_ptr, float *m_ptr, float *v_ptr,
+                                         float step_size, float beta1, float beta2, float eps, size_t num_elements) {
+    size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= num_elements) {
+        return;
+    }
+    const float g = grad_ptr[idx];
+    m_ptr[idx] = beta1 * m_ptr[idx] + (1.0f - beta1) * g;
+    v_ptr[idx] = beta2 * v_ptr[idx] + (1.0f - beta2) * g * g;
+    param_ptr[idx] -= step_size * m_ptr[idx] / (sqrtf(v_ptr[idx]) + eps);
+}
+
 void AdamAccumulateGrad(const std::shared_ptr<Tensor> &grad, const std::shared_ptr<Tensor> &param,
                         const std::shared_ptr<Tensor> &m, const std::shared_ptr<Tensor> &v, float learning_rate,
                         float beta1, float beta2, float eps, int64_t t) {
@@ -29,6 +41,19 @@ void AdamAccumulateGrad(const std::shared_ptr<Tensor> &grad, const std::shared_p
     // TODO：实现Adam优化器的梯度累积和参数更新
     // REF:
     // =================================== 作业 ===================================
+    const size_t num_elements = grad->NumElements();
+    float *grad_ptr = static_cast<float *>(grad->DataPtr());
+    float *param_ptr = static_cast<float *>(param->DataPtr());
+    float *m_ptr = static_cast<float *>(m->DataPtr());
+    float *v_ptr = static_cast<float *>(v->DataPtr());
+
+    const float step_size
+        = learning_rate * sqrtf(1.0f - powf(beta2, static_cast<float>(t))) / (1.0f - powf(beta1, static_cast<float>(t)));
+
+    int threads_per_block = 256;
+    int num_blocks = (num_elements + threads_per_block - 1) / threads_per_block;
+    AdamAccumulateGradKernel<<<num_blocks, threads_per_block>>>(grad_ptr, param_ptr, m_ptr, v_ptr, step_size, beta1,
+                                                                beta2, eps, num_elements);
 }
 } // namespace infini_train::kernels::cuda
 
diff --git a/infini_train/src/kernels/cuda/linear.cu b/infini_train/src/kernels/cuda/linear.cu
index efaaaa6..9cea69b 100644
--- a/infini_train/src/kernels/cuda/linear.cu
+++ b/infini_train/src/kernels/cuda/linear.cu
@@ -1,4 +1,6 @@
 #include "cublas_v2.h"
+#include <algorithm>
+#include <vector>
 #include "glog/logging.h"
 #include <cub/block/block_reduce.cuh>
 
@@ -24,25 +26,113 @@ namespace infini_train::kernels::cuda {
     } while (0)
 
 std::shared_ptr<Tensor> MatmulForward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Tensor> &other) {
-    // =================================== 作业 ===================================
-    // TODO：实现CUDA上的矩阵乘法前向计算
-    // REF:
-    // =================================== 作业 ===================================
+    const auto &input_dims = input->Dims();
+    const auto &other_dims = other->Dims();
+
+    CHECK_GE(input_dims.size(), 2);
+    CHECK_GE(other_dims.size(), 2);
+
+    int64_t M = input_dims[input_dims.size() - 2];
+    int64_t K = input_dims.back();
+    int64_t K2 = other_dims[other_dims.size() - 2];
+    int64_t N = other_dims.back();
+    CHECK_EQ(K, K2);
+
+    int64_t input_batches = 1;
+    for (size_t i = 0; i < input_dims.size() - 2; ++i)
+        input_batches *= input_dims[i];
+    int64_t other_batches = 1;
+    for (size_t i = 0; i < other_dims.size() - 2; ++i)
+        other_batches *= other_dims[i];
+
+    int64_t max_batches = std::max(input_batches, other_batches);
+    std::vector<int64_t> output_dims = (input_dims.size() >= other_dims.size()) ? input_dims : other_dims;
+    output_dims[output_dims.size() - 2] = M;
+    output_dims.back() = N;
+
+    auto output = std::make_shared<Tensor>(output_dims, DataType::kFLOAT32, input->GetDevice());
+
+    cublasHandle_t handle;
+    CUBLAS_CHECK(cublasCreate(&handle));
+
+    const float alpha = 1.0f;
+    const float beta = 0.0f;
+
+    CUBLAS_CHECK(cublasSgemmStridedBatched(
+        handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, static_cast<const float *>(other->DataPtr()), N,
+        (other_batches == 1 ? 0 : K * N), static_cast<const float *>(input->DataPtr()), K,
+        (input_batches == 1 ? 0 : M * K), &beta, static_cast<float *>(output->DataPtr()), N, M * N, max_batches));
 
-    auto output = std::make_shared<Tensor>();
+    CUBLAS_CHECK(cublasDestroy(handle));
     return output;
 }
 
 std::tuple<std::shared_ptr<Tensor>, std::shared_ptr<Tensor>>
 MatmulBackward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Tensor> &other,
                const std::shared_ptr<Tensor> &grad_output) {
-    // =================================== 作业 ===================================
-    // TODO：实现CUDA上的矩阵乘法反向传播
-    // REF:
-    // =================================== 作业 ===================================
+    const auto &input_dims = input->Dims();
+    const auto &other_dims = other->Dims();
+    const auto &grad_output_dims = grad_output->Dims();
+
+    int64_t M = input_dims[input_dims.size() - 2];
+    int64_t K = input_dims.back();
+    int64_t N = other_dims.back();
+
+    int64_t input_batches = 1;
+    for (size_t i = 0; i < input_dims.size() - 2; ++i)
+        input_batches *= input_dims[i];
+    int64_t other_batches = 1;
+    for (size_t i = 0; i < other_dims.size() - 2; ++i)
+        other_batches *= other_dims[i];
+    int64_t grad_batches = 1;
+    for (size_t i = 0; i < grad_output_dims.size() - 2; ++i)
+        grad_batches *= grad_output_dims[i];
+
+    auto grad_input = std::make_shared<Tensor>(input_dims, DataType::kFLOAT32, input->GetDevice());
+    auto grad_other = std::make_shared<Tensor>(other_dims, DataType::kFLOAT32, other->GetDevice());
+    grad_input->Fill<float>(0.0f);
+    grad_other->Fill<float>(0.0f);
+
+    cublasHandle_t handle;
+    CUBLAS_CHECK(cublasCreate(&handle));
 
-    auto grad_input = std::make_shared<Tensor>();
-    auto grad_other = std::make_shared<Tensor>();
+    const float alpha = 1.0f;
+    const float beta = 0.0f;
+
+    // grad_input = grad_output * other^T --> grad_input^T = other * grad_output^T
+    // Dimensions: grad_input(M, K), other(K, N), grad_output(M, N)
+    // col-major: grad_input_T(K, M), other_T(N, K), grad_output_T(N, M)
+    // C(K, M) = A(K, N) * B(N, M) --> A = other_T^T, B = grad_output_T
+    CUBLAS_CHECK(cublasSgemmStridedBatched(
+        handle, CUBLAS_OP_T, CUBLAS_OP_N, (int)K, (int)M, (int)N, &alpha, static_cast<const float *>(other->DataPtr()),
+        (int)N, (other_batches == 1 ? 0 : (long long)(K * N)), static_cast<const float *>(grad_output->DataPtr()),
+        (int)N, (long long)(M * N), &beta, static_cast<float *>(grad_input->DataPtr()), (int)K,
+        (input_batches == 1 ? 0 : (long long)(M * K)), (int)grad_batches));
+
+    // grad_other = input^T * grad_output --> grad_other^T = grad_output^T * input
+    // Dimensions: grad_other(K, N), input(M, K), grad_output(M, N)
+    // col-major: grad_other_T(N, K), input_T(K, M), grad_output_T(N, M)
+    // C(N, K) = A(N, M) * B(M, K) --> A = grad_output_T, B = input_T^T
+    float beta_accum = 0.0f;
+    if (other_batches == 1 && grad_batches > 1) {
+        // We need to sum across batches. SgemmStridedBatched with stride 0 for C will NOT work correctly (race condition).
+        // For simplicity and safety, we use a loop here if we need to sum.
+        for (int b = 0; b < grad_batches; ++b) {
+            float cur_beta = (b == 0) ? 0.0f : 1.0f;
+            CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, (int)N, (int)K, (int)M, &alpha,
+                                     static_cast<const float *>(grad_output->DataPtr()) + b * M * N, (int)N,
+                                     static_cast<const float *>(input->DataPtr()) + (input_batches == 1 ? 0 : b * M * K),
+                                     (int)K, &cur_beta, static_cast<float *>(grad_other->DataPtr()), (int)N));
+        }
+    } else {
+        CUBLAS_CHECK(cublasSgemmStridedBatched(
+            handle, CUBLAS_OP_N, CUBLAS_OP_T, (int)N, (int)K, (int)M, &alpha,
+            static_cast<const float *>(grad_output->DataPtr()), (int)N, (long long)(M * N),
+            static_cast<const float *>(input->DataPtr()), (int)K, (input_batches == 1 ? 0 : (long long)(M * K)), &beta,
+            static_cast<float *>(grad_other->DataPtr()), (int)N, (long long)(K * N), (int)grad_batches));
+    }
+
+    CUBLAS_CHECK(cublasDestroy(handle));
     return {grad_input, grad_other};
 }
 
diff --git a/infini_train/src/tensor.cc b/infini_train/src/tensor.cc
index 8f8c744..8e9ba43 100644
--- a/infini_train/src/tensor.cc
+++ b/infini_train/src/tensor.cc
@@ -4,6 +4,7 @@
 #include <cstring>
 #include <fstream>
 #include <functional>
+#include <iomanip>
 #include <memory>
 #include <numeric>
 #include <unordered_map>
@@ -16,6 +17,7 @@
 #include "glog/logging.h"
 
 #include "infini_train/include/autograd/elementwise.h"
+#include "infini_train/include/autograd/function.h"
 #include "infini_train/include/autograd/matmul.h"
 #include "infini_train/include/autograd/misc.h"
 #include "infini_train/include/autograd/outer.h"
@@ -282,8 +284,31 @@ std::shared_ptr<Tensor> Tensor::Flatten(int64_t start, int64_t end) {
     // TODO：实现张量扁平化操作，将指定维度范围[start, end]内的所有维度合并为一个维度
     // HINT:
     // =================================== 作业 ===================================
+    const int64_t ndim = static_cast<int64_t>(dims_.size());
+    if (start < 0) {
+        start += ndim;
+    }
+    if (end < 0) {
+        end += ndim;
+    }
+    CHECK_GE(start, 0);
+    CHECK_GE(end, start);
+    CHECK_LT(end, ndim);
 
-    return std::make_shared<Tensor>();
+    std::vector<int64_t> new_shape;
+    new_shape.reserve(static_cast<size_t>(ndim - (end - start)));
+    for (int64_t i = 0; i < start; ++i) {
+        new_shape.push_back(dims_[i]);
+    }
+    int64_t flattened = 1;
+    for (int64_t i = start; i <= end; ++i) {
+        flattened *= dims_[i];
+    }
+    new_shape.push_back(flattened);
+    for (int64_t i = end + 1; i < ndim; ++i) {
+        new_shape.push_back(dims_[i]);
+    }
+    return Contiguous()->View(new_shape);
 }
 
 std::shared_ptr<Tensor> Tensor::Squeeze(int64_t dim) {
@@ -358,6 +383,30 @@ void Tensor::Backward(std::shared_ptr<Tensor> gradient, bool retain_graph, bool
     // TODO：实现自动微分反向传播
     // 功能描述：1. 计算当前张量对叶子节点的梯度    2. 支持多输出场景的梯度累加
     // =================================== 作业 ===================================
+    (void)retain_graph;
+    (void)create_graph;
+    if (!requires_grad_) {
+        return;
+    }
+
+    if (!gradient) {
+        if (NumElements() == 1) {
+            gradient = std::make_shared<Tensor>(dims_, dtype_, GetDevice());
+            gradient->Fill<float>(1.0f);
+        } else {
+            LOG(FATAL) << "grad can be implicitly created only for scalar outputs";
+        }
+    }
+
+    if (grad_fn_) {
+        grad_fn_->BackwardPartial(gradient, output_idx_);
+    } else if (is_leaf_) {
+        if (grad_) {
+            auto device = grad_->GetDevice().Type();
+            auto kernel = Dispatcher::Instance().GetKernel({device, "AccumulateGrad"});
+            kernel.Call<void>(gradient, 1.0f, grad_);
+        }
+    }
 }
 
 void Tensor::ZeroGrad() {
diff --git a/test/example/test_gpt2.cc b/test/example/test_gpt2.cc
index e7d038f..f594122 100644
--- a/test/example/test_gpt2.cc
+++ b/test/example/test_gpt2.cc
@@ -1,6 +1,5 @@
 #include <chrono>
 #include <cstdlib>
-#include <format>
 #include <memory>
 #include <optional>
 #include <unordered_map>