refactor: clean up legacy load_state_dict for linear layers

src/CMakeLists.txt

@@ -3,9 +3,7 @@ add_subdirectory(common)
 add_subdirectory(handlers)
 add_subdirectory(kernels)
 add_subdirectory(tokenizer)
-add_subdirectory(module)
 add_subdirectory(layers)
-add_subdirectory(quantization)
 add_subdirectory(models)
 add_subdirectory(model_loader)
 add_subdirectory(model_parallel)

src/engine/llm_engine.h

@@ -5,8 +5,8 @@
 #include "batch.h"
 #include "common/macros.h"
 #include "engine.h"
+#include "layers/quantization/quant_args.h"
 #include "memory/block_manager.h"
-#include "quantization/quant_args.h"
 #include "tokenizer/tokenizer.h"
 #include "tokenizer/tokenizer_args.h"
 #include "worker.h"

src/engine/worker.h

@@ -4,14 +4,14 @@
 #include <torch/torch.h>
 
 #include "common/threadpool.h"
+#include "layers/quantization/quant_args.h"
 #include "model_loader/state_dict.h"
 #include "model_parallel/parallel_args.h"
 #include "model_runner.h"
 #include "models/causal_lm.h"
 #include "models/model_args.h"
 #include "models/parameters.h"
 #include "parameters.h"
-#include "quantization/quant_args.h"
 
 namespace llm {
 

src/engine/worker_test.cpp

@@ -4,9 +4,9 @@
 
 #include "engine/batch.h"
 #include "engine/utils.h"
+#include "layers/quantization/quant_args.h"
 #include "memory/block_manager.h"
 #include "models/simple_model.h"
-#include "quantization/quant_args.h"
 namespace llm {
 
 class TestableWorker {

src/layers/CMakeLists.txt

@@ -1,32 +1,6 @@
 include(cc_library)
 include(cc_test)
 
-cc_library(
-  NAME
-    linear
-  HDRS
-    linear.h
-    qkv_linear.h
-    linear_impl.h
-    fused_linear.h
-    weight_utils.h
-  SRCS
-    linear.cpp
-    qkv_linear.cpp
-    linear_impl.cpp
-    fused_linear.cpp
-    weight_utils.cpp
-  DEPS
-    :state_dict
-    :model_parallel
-    :quantization
-    :kernels
-    :module
-    glog::glog
-    gflags::gflags
-    torch
-)
-
 cc_library(
   NAME
     pos_embedding
@@ -73,14 +47,15 @@ cc_test(
     activation_test.cpp
     pos_embedding_test.cpp
     normalization_test.cpp
-    linear_test.cpp
-    qkv_linear_test.cpp
   DEPS
     :layers
     :state_dict
     absl::random_random
     :gtest_main
 )
 
+add_subdirectory(module)
+add_subdirectory(linear)
+add_subdirectory(quantization)
 add_subdirectory(attention)
 add_subdirectory(moe)

src/layers/attention/attention.h

@@ -3,10 +3,10 @@
 #include <torch/torch.h>
 
 #include "layers/attention/handler.h"
+#include "layers/module/module.h"
+#include "layers/module/module_holder.h"
 #include "memory/kv_cache.h"
 #include "models/parameters.h"
-#include "module/module.h"
-#include "module/module_holder.h"
 
 namespace llm {
 

src/layers/embedding.h

@@ -5,10 +5,10 @@
 
 #include <cstdint>
 
+#include "layers/module/module.h"
+#include "layers/module/module_holder.h"
 #include "model_loader/state_dict.h"
 #include "model_parallel/model_parallel.h"
-#include "module/module.h"
-#include "module/module_holder.h"
 
 namespace llm {