[init] baseline7 from project_6

ex_engine/xllm_layers/common/dense_mlp.cpp

@@ -0,0 +1,141 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "dense_mlp.h"
+
+#include <glog/logging.h>
+
+#include "kernels/ops_api.h"
+#include "platform/device.h"
+
+namespace xllm {
+namespace layer {
+
+DenseMLPImpl::DenseMLPImpl(int64_t hidden_size,
+                           int64_t intermediate_size,
+                           bool is_gated,
+                           bool has_bias,
+                           const std::string& hidden_act,
+                           bool enable_result_reduction,
+                           const QuantArgs& quant_args,
+                           ProcessGroup* process_group,
+                           const torch::TensorOptions& options,
+                           const std::string& module_prefix)
+    : is_gated_(is_gated),
+      intermediate_size_(intermediate_size),
+      process_group_(process_group),
+      hidden_act_(hidden_act) {
+  // Check if using w8a8 smoothquant quantization
+  is_smoothquant_ = quant_args.quant_method() == kQuantMethodSmoothquant;
+
+  if (is_smoothquant_) {
+    // Safety check: only w8a8 smoothquant is supported
+    if (quant_args.bits() != 8 || !quant_args.activation_dynamic()) {
+      LOG(FATAL)
+          << "DenseMLP w8a8 mode only supports w8a8 smoothquant quantization. "
+          << "Got bits=" << quant_args.bits()
+          << ", activation_dynamic=" << quant_args.activation_dynamic();
+    }
+  }
+
+  // Determine extra args based on quantization mode
+  LinearExtraArgs gate_up_proj_extra_args("none", false);
+  LinearExtraArgs down_proj_extra_args("none", false);
+  if (is_smoothquant_) {
+    // For per-token smoothquant, use specific args
+    down_proj_extra_args = LinearExtraArgs(hidden_act_, is_gated_);
+  }
+
+  // 1. gate + up
+  int64_t out_feature = is_gated_ ? intermediate_size_ * 2 : intermediate_size_;
+  gate_up_proj_ =
+      register_module("gate_up_proj",
+                      ColumnParallelLinear(hidden_size,
+                                           out_feature,
+                                           /*bias=*/has_bias,
+                                           /*gather_output=*/false,
+                                           quant_args,
+                                           process_group_,
+                                           options,
+                                           gate_up_proj_extra_args));
+
+  act_ = register_module("act", Activation(hidden_act_, is_gated_));
+
+  // 2. down
+  const auto down_proj_quant_args =
+      module_prefix.empty()
+          ? quant_args
+          : quant_args.for_module(module_prefix + ".down_proj");
+  down_proj_ = register_module("down_proj",
+                               RowParallelLinear(intermediate_size_,
+                                                 hidden_size,
+                                                 /*bias=*/has_bias,
+                                                 /*input_is_parallelized=*/true,
+                                                 enable_result_reduction,
+                                                 down_proj_quant_args,
+                                                 process_group_,
+                                                 options,
+                                                 down_proj_extra_args));
+}
+
+torch::Tensor DenseMLPImpl::forward(const torch::Tensor& hidden_states) {
+  // input shape: [num_tokens, hidden_size]
+  auto gate_up = gate_up_proj_->forward(hidden_states);
+
+  if (is_smoothquant_) {
+    // For w8a8 quantization, the active operation is fused with the down_proj
+    return down_proj_->forward(gate_up);
+  } else {
+    torch::Tensor output;
+    if (Device::type_str() != "npu") {
+      int64_t batch_size = gate_up.sizes()[0];
+      output = torch::empty(
+          {batch_size, intermediate_size_ / process_group_->world_size()},
+          gate_up.options());
+    }
+
+    act_->forward(gate_up, output);
+    return down_proj_->forward(output);
+  }
+}
+
+void DenseMLPImpl::load_state_dict(const StateDict& state_dict) {
+  gate_up_proj_->load_state_dict(state_dict, {"gate_proj.", "up_proj."});
+  down_proj_->load_state_dict(state_dict.get_dict_with_prefix("down_proj."));
+}
+
+void DenseMLPImpl::load_state_dict(const StateDict& state_dict,
+                                   const std::vector<std::string>& gate_up_name,
+                                   const std::string& down_name) {
+  if (is_gated_) {
+    CHECK_EQ(gate_up_name.size(), 2);
+    gate_up_proj_->load_state_dict(state_dict, gate_up_name);
+  } else {
+    CHECK_EQ(gate_up_name.size(), 1);
+    gate_up_proj_->load_state_dict(
+        state_dict.get_dict_with_prefix(gate_up_name[0]));
+  }
+  down_proj_->load_state_dict(state_dict.get_dict_with_prefix(down_name));
+}
+
+std::optional<torch::Tensor> DenseMLPImpl::get_fp8_input_scale() const {
+  if (gate_up_proj_) {
+    return gate_up_proj_->get_input_scale();
+  }
+  return std::nullopt;
+}
+
+}  // namespace layer
+}  // namespace xllm