Revert "feat: add ERNIE 4.5 VL support"

This reverts commit 985e10a.

Author: qinyiqun

csrc/engine/compiler/paged_compiler.cpp

@@ -4,17 +4,6 @@
 
 namespace infinilm::engine {
 
-namespace {
-bool same_shape_and_dtype(const infinicore::Tensor &lhs, const infinicore::Tensor &rhs) {
-    return lhs->shape() == rhs->shape() && lhs->dtype() == rhs->dtype();
-}
-
-bool is_mrope_decode_position_ids(const infinicore::Tensor &position_ids, size_t batch_size) {
-    const auto shape = position_ids->shape();
-    return shape.size() == 2 && shape[0] == batch_size && shape[1] == 3;
-}
-} // namespace
-
 PagedCompiler::PagedCompiler(const std::shared_ptr<InfinilmModel> &model, RankBarrier *barrier)
     : GraphCompiler(model, barrier) {
     for (size_t b = 1; b < 64; ++b) {
@@ -36,19 +25,14 @@ void PagedCompiler::compile() {
         size_t nblocks = dynamic_cast<const cache::PagedKVCacheConfig *>(model_->get_cache_config())->num_blocks();
         size_t max_batch_size = *std::max_element(decode_batch_sizes_.begin(), decode_batch_sizes_.end());
         compiled_map_decode_.clear();
-        compiled_map_decode_mrope_.clear();
         block_tables_holder_ = infinicore::Tensor::empty(
             {nblocks * max_batch_size}, infinicore::DataType::I32, infinicore::context::getDevice());
         set_zeros(block_tables_holder_);
 
-        auto make_decode_input = [&](size_t b, bool mrope_position_ids) {
+        auto make_decode_input = [&](size_t b) {
             InfinilmModel::Input input;
             input.input_ids = infinicore::Tensor::empty({1, b}, infinicore::DataType::I64, infinicore::context::getDevice());
-            if (mrope_position_ids) {
-                input.position_ids = infinicore::Tensor::empty({b, 3}, infinicore::DataType::I64, infinicore::context::getDevice());
-            } else {
-                input.position_ids = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice());
-            }
+            input.position_ids = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice());
             input.total_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::I32, infinicore::context::getDevice());
             set_zeros(input.input_ids.value());
             set_zeros(input.position_ids.value());
@@ -82,7 +66,7 @@ void PagedCompiler::compile() {
 
         {
             const size_t warmup_batch_size = std::min(max_batch_size, static_cast<size_t>(64));
-            auto input = make_decode_input(warmup_batch_size, false);
+            auto input = make_decode_input(warmup_batch_size);
             model_->forward(input);
             infinicore::context::syncStream();
             // Warmup runs the eager Marlin path and may leave per-layer lock
@@ -92,10 +76,8 @@ void PagedCompiler::compile() {
             infinicore::context::syncStream();
         }
 
-        auto capture_decode_graph = [&](size_t b,
-                                        bool mrope_position_ids,
-                                        std::unordered_map<size_t, CompiledResult> &compiled_map) {
-            auto input = make_decode_input(b, mrope_position_ids);
+        for (size_t b : decode_batch_sizes_) {
+            auto input = make_decode_input(b);
 
             barrier_->wait();
             (void)model_->forward(input);
@@ -114,15 +96,7 @@ void PagedCompiler::compile() {
             auto shared_output = std::shared_ptr<InfinilmModel::Output>(
                 new InfinilmModel::Output{infinicore::graph::GraphTensor(output.logits)});
 
-            compiled_map[b] = CompiledResult{std::move(input), std::make_tuple(graph, shared_output)};
-        };
-
-        const bool compile_mrope_decode = model_->supports_mrope_position_ids();
-        for (size_t b : decode_batch_sizes_) {
-            capture_decode_graph(b, false, compiled_map_decode_);
-            if (compile_mrope_decode) {
-                capture_decode_graph(b, true, compiled_map_decode_mrope_);
-            }
+            compiled_map_decode_[b] = CompiledResult{std::move(input), std::make_tuple(graph, shared_output)};
         }
     }
 }
@@ -136,23 +110,12 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input &
         if (batch_size != input.input_ids.value()->size(1)) {
             return {nullptr, nullptr};
         } else {
-            const bool mrope_position_ids = is_mrope_decode_position_ids(input.position_ids.value(), batch_size);
-            auto &compiled_map = mrope_position_ids ? compiled_map_decode_mrope_ : compiled_map_decode_;
-            auto result = compiled_map.find(batch_size);
-            if (result == compiled_map.end()) {
+            auto result = compiled_map_decode_.find(batch_size);
+            if (result == compiled_map_decode_.end()) {
                 return {nullptr, nullptr};
             }
             auto &graph_input = result->second.input;
 
-            if (!same_shape_and_dtype(graph_input.input_ids.value(), input.input_ids.value()) ||
-                !same_shape_and_dtype(graph_input.position_ids.value(), input.position_ids.value()) ||
-                !same_shape_and_dtype(graph_input.total_sequence_lengths.value(), input.total_sequence_lengths.value()) ||
-                !same_shape_and_dtype(graph_input.input_offsets.value(), input.input_offsets.value()) ||
-                !same_shape_and_dtype(graph_input.cu_seqlens.value(), input.cu_seqlens.value()) ||
-                !same_shape_and_dtype(graph_input.slot_mapping.value(), input.slot_mapping.value())) {
-                return {nullptr, nullptr};
-            }
-
             graph_input.input_ids.value()->copy_from(input.input_ids.value());
             graph_input.position_ids.value()->copy_from(input.position_ids.value());
             graph_input.total_sequence_lengths.value()->copy_from(input.total_sequence_lengths.value());
@@ -169,9 +132,6 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input &
             // runtime logical region. Avoid clearing the full preallocated
             // holder on every decode token.
             auto &graph_block_tables = graph_input.block_tables.value();
-            if (graph_block_tables->dtype() != input.block_tables.value()->dtype()) {
-                return {nullptr, nullptr};
-            }
             set_minus_one_device_async(graph_block_tables);
             graph_block_tables->narrow({{1, 0, block_per_req}})->copy_from(input.block_tables.value());
             graph_input.slot_mapping.value()->copy_from(input.slot_mapping.value());

csrc/engine/compiler/paged_compiler.hpp

@@ -27,10 +27,5 @@ class PagedCompiler : public GraphCompiler {
         size_t, // num_requests
         CompiledResult>
         compiled_map_decode_;
-
-    std::unordered_map<
-        size_t, // num_requests
-        CompiledResult>
-        compiled_map_decode_mrope_;
 };
 } // namespace infinilm::engine

csrc/layers/linear/fused_linear.cpp

@@ -151,10 +151,9 @@ GateUpParallelLinear::GateUpParallelLinear(size_t hidden_size, size_t intermedia
 
 std::tuple<infinicore::Tensor, infinicore::Tensor> GateUpParallelLinear::forward_split(infinicore::Tensor &input) {
     auto output = this->forward(input);
-    const size_t split_dim = output->ndim() - 1;
-    auto cols = output->size(split_dim);
-    auto gate_output = output->narrow({{split_dim, 0, cols / 2}});
-    auto up_output = output->narrow({{split_dim, cols / 2, cols / 2}});
+    auto cols = output->shape()[2];
+    auto gate_output = output->narrow({{2, 0, cols / 2}});
+    auto up_output = output->narrow({{2, cols / 2, cols / 2}});
     return std::make_tuple(gate_output, up_output);
 }
 

csrc/layers/moe/router/topk_router.cpp

@@ -1,6 +1,5 @@
 #include "topk_router.hpp"
 
-#include "../../../utils.hpp"
 #include "infinicore/ops.hpp"
 
 #include <stdexcept>
@@ -13,9 +12,6 @@ TopKRouterBackend parse_router_backend(const std::string &backend) {
     if (backend == "softmax") {
         return TopKRouterBackend::Softmax;
     }
-    if (backend == "topksoftmax" || backend == "legacy_softmax") {
-        return TopKRouterBackend::LegacySoftmax;
-    }
     if (backend == "sigmoid") {
         return TopKRouterBackend::Sigmoid;
     }
@@ -57,13 +53,9 @@ TopKRouter::TopKRouter(std::shared_ptr<infinilm::config::ModelConfig> model_conf
     use_correction_bias_ = model_config->get_or<bool>({"e_score_correction_bias", "moe_router_use_correction_bias"}, false) || router_backend_ == TopKRouterBackend::FusedGate;
     ASSERT((num_experts_ > 0) && (num_experts_per_tok_ > 0) && (num_experts_per_tok_ <= num_experts_));
 
-    const auto router_dtype_name = model_config->get_or<std::string>("moe_router_dtype", "");
-    const auto router_dtype = router_dtype_name.empty()
-                                ? model_config->get_dtype()
-                                : parse_dtype(router_dtype_name);
     INFINICORE_NN_PARAMETER_INIT(
         weight,
-        ({num_experts_, model_config->get<size_t>("hidden_size")}, router_dtype, device));
+        ({num_experts_, model_config->get<size_t>("hidden_size")}, model_config->get_dtype(), device));
 
     if (use_correction_bias_) {
         INFINICORE_NN_PARAMETER_INIT(e_score_correction_bias, ({num_experts_}, infinicore::DataType::F32, device));
@@ -87,17 +79,9 @@ TopKRouter::TopKRouter(std::shared_ptr<infinilm::config::ModelConfig> model_conf
 
 std::tuple<infinicore::Tensor, infinicore::Tensor> TopKRouter::forward(const infinicore::Tensor &hidden_states) const {
     ASSERT(hidden_states->ndim() == 2);
+
     size_t ntoken = hidden_states->shape()[0];
-    auto router_input = hidden_states->is_contiguous() ? hidden_states : hidden_states->contiguous();
-    auto router_weight = static_cast<infinicore::Tensor>(weight_);
-    router_weight = router_weight->is_contiguous() ? router_weight : router_weight->contiguous();
-    auto logits_shape = router_input->shape();
-    logits_shape[logits_shape.size() - 1] = num_experts_;
-    auto router_logits = infinicore::Tensor::empty(
-        logits_shape,
-        router_weight->dtype(),
-        router_input->device());
-    infinicore::op::linear_(router_logits, router_input, router_weight, std::nullopt, 1.0f);
+    auto router_logits = infinicore::op::linear(hidden_states, weight_, std::nullopt, 1.0f);
 
     auto router_scores = infinicore::Tensor::empty({ntoken, num_experts_per_tok_}, infinicore::DataType::F32, hidden_states->device());
     auto router_indices = infinicore::Tensor::empty({ntoken, num_experts_per_tok_}, infinicore::DataType::I32, hidden_states->device());
@@ -114,14 +98,6 @@ std::tuple<infinicore::Tensor, infinicore::Tensor> TopKRouter::forward(const inf
             norm_topk_prob_,
             moe_softcapping_);
         break;
-    case TopKRouterBackend::LegacySoftmax:
-        infinicore::op::topksoftmax(
-            router_scores,
-            router_indices,
-            router_logits,
-            num_experts_per_tok_,
-            norm_topk_prob_);
-        break;
     case TopKRouterBackend::Sigmoid:
         infinicore::op::moe_topk_sigmoid_(
             router_scores,

csrc/layers/moe/router/topk_router.hpp

@@ -11,7 +11,6 @@ namespace infinilm::layers::moe {
 
 enum class TopKRouterBackend {
     Softmax,
-    LegacySoftmax,
     Sigmoid,
     FusedGate,
 };