Feat/glm5 2 main

csrc/config/config_factory.cpp

@@ -12,7 +12,7 @@ std::shared_ptr<infinilm::config::ModelConfig> ConfigFactory::createConfig(const
     const auto &config_map = models::get_model_config_map();
     auto it = config_map.find(model_type);
     if (it != config_map.end()) {
-        it->second(model_config);
+        model_config = it->second(model_config);
     } else {
         throw std::invalid_argument("infinilm::config::ConfigFactory::createConfig: Unsupported model config type: " + model_type);
     }

csrc/config/quant_config.cpp

@@ -22,6 +22,10 @@ QuantConfig::get_quantization_method() const {
         return std::make_shared<infinilm::quantization::GPTQ>(quantization_config);
     } else if (quant_method == "quark") {
         return std::make_shared<infinilm::quantization::MXFP4>(quantization_config);
+    } else if (quant_method == "glm_w8a8" || quant_method == "w8a8") {
+        return std::make_shared<infinilm::quantization::GlmW8A8>(quantization_config);
+    } else if (quant_method == "glm_w4a8" || quant_method == "w4a8") {
+        return std::make_shared<infinilm::quantization::GlmW4A8>(quantization_config);
     } else {
         return std::make_shared<infinilm::quantization::NoneQuantization>(quantization_config);
     }

csrc/engine/compiler/paged_compiler.cpp

@@ -61,6 +61,7 @@ void PagedCompiler::compile() {
         size_t nblocks = dynamic_cast<const cache::PagedKVCacheConfig *>(model_->get_cache_config())->num_blocks();
         auto &forward_context = infinilm::global_state::get_forward_context();
         const bool has_mamba_state = has_mamba_cache(forward_context);
+        size_t block_size = dynamic_cast<const cache::PagedKVCacheConfig *>(model_->get_cache_config())->block_size();
 
         const auto &model_config = model_->get_model_config();
         const size_t position_id_axes = model_config == nullptr
@@ -96,6 +97,11 @@ void PagedCompiler::compile() {
                 input_offsets_vec[i] = i;
             }
             infinicore::context::memcpyH2D(input.input_offsets.value()->data(), input_offsets_vec.data(), (b + 1) * sizeof(int32_t), false);
+            input.request_ids = infinicore::Tensor::empty(
+                {b}, infinicore::DataType::I32, infinicore::context::getDevice());
+            infinicore::context::memcpyH2D(
+                input.request_ids.value()->data(), input_offsets_vec.data(),
+                b * sizeof(int32_t), false);
             input.cu_seqlens = infinicore::Tensor::empty({b + 1}, infinicore::DataType::I32, infinicore::context::getDevice());
             infinicore::context::memcpyH2D(input.cu_seqlens.value()->data(), input_offsets_vec.data(), (b + 1) * sizeof(int32_t), false);
             const size_t block_per_req = nblocks;
@@ -127,9 +133,13 @@ void PagedCompiler::compile() {
                 input.past_sequence_lengths,
                 input.total_sequence_lengths,
                 input.input_offsets,
+                input.request_ids,
                 input.cu_seqlens,
                 input.block_tables,
                 input.slot_mapping,
+                0,
+                0,
+                static_cast<int64_t>(nblocks * block_size),
             };
             // Hybrid linear-attention layers read cache indices from the same
             // thread-local context. These tensors remain alive in CompiledResult

csrc/engine/compiler/static_batching_compiler.cpp

@@ -27,6 +27,7 @@ void StaticBatchingCompiler::compile() {
             input.past_sequence_lengths,
             input.total_sequence_lengths,
             input.input_offsets,
+            input.request_ids,
             input.cu_seqlens,
             input.block_tables,
             input.slot_mapping,

csrc/engine/infer_engine.cpp

@@ -43,6 +43,53 @@ size_t max_length_from_offsets(
     return max_length;
 }
 
+std::optional<infinicore::Tensor> request_ids_from_offsets(
+    const std::optional<infinicore::Tensor> &offsets,
+    const infinicore::Device &device) {
+    if (!offsets.has_value()) {
+        return std::nullopt;
+    }
+
+    auto cpu_offsets = offsets.value();
+    if (cpu_offsets->device().getType() != infinicore::Device::Type::CPU) {
+        cpu_offsets = cpu_offsets->to(infinicore::Device::cpu());
+        infinicore::context::syncStream();
+    }
+    if (cpu_offsets->dtype() != infinicore::DataType::I32
+        || cpu_offsets->shape().size() != 1
+        || cpu_offsets->shape()[0] < 2) {
+        throw std::invalid_argument(
+            "input_offsets must be a one-dimensional int32 tensor with at least two entries");
+    }
+
+    const auto *values = reinterpret_cast<const int32_t *>(cpu_offsets->data());
+    const size_t num_requests = cpu_offsets->shape()[0] - 1;
+    if (values[0] != 0 || values[num_requests] < 0) {
+        throw std::invalid_argument(
+            "input_offsets must start at zero and be nonnegative");
+    }
+    std::vector<int32_t> request_ids(
+        static_cast<size_t>(values[num_requests]));
+    for (size_t request = 0; request < num_requests; ++request) {
+        if (values[request] > values[request + 1]) {
+            throw std::invalid_argument("input_offsets must be nondecreasing");
+        }
+        std::fill(
+            request_ids.begin() + values[request],
+            request_ids.begin() + values[request + 1],
+            static_cast<int32_t>(request));
+    }
+
+    auto result = infinicore::Tensor::empty(
+        {request_ids.size()}, infinicore::DataType::I32, device);
+    if (!request_ids.empty()) {
+        infinicore::context::memcpyH2D(
+            result->data(), request_ids.data(),
+            request_ids.size() * sizeof(int32_t), false);
+    }
+    return result;
+}
+
 } // namespace
 
 //------------------------------------------------------
@@ -230,16 +277,19 @@ InferEngine::Input::to_model_input(infinicore::Device device) const {
         visual_token_ranges,
         to_device(target_hidden_states),
         sample_all_positions};
+    input.request_ids = request_ids_from_offsets(input_offsets, device);
 
     infinilm::global_state::get_forward_context().attn_metadata = {
         input.past_sequence_lengths,
         input.total_sequence_lengths,
         input.input_offsets,
+        input.request_ids,
         input.cu_seqlens,
         input.block_tables,
         input.slot_mapping,
         max_query_length,
-        max_sequence_length};
+        max_sequence_length,
+        max_context_len};
 
     infinilm::global_state::get_forward_context().mamba_metadata = {
         input.input_offsets,

csrc/engine/rank_worker.hpp

@@ -72,6 +72,8 @@ class RankWorker {
         std::optional<infinicore::Tensor> target_hidden_states;
         /// Sample logits at every packed input position instead of one token per request.
         bool sample_all_positions{false};
+        /// Maximum total sequence length in the current request batch.
+        std::optional<int64_t> max_context_len;
 
         float temperature{1};
 

csrc/global_state/forward_context.hpp

@@ -11,6 +11,8 @@ struct AttentionMetadata {
     std::optional<infinicore::Tensor> total_sequence_lengths;
     /// Offsets of each request in a continous-batched sequence, of shape `[num_requests + 1]`.
     std::optional<infinicore::Tensor> input_offsets;
+    /// Request id for each flattened input token, of shape `[num_tokens]`.
+    std::optional<infinicore::Tensor> request_ids;
     /// Cumulative total sequence lengths for each request, of shape `[num_requests + 1]`.
     std::optional<infinicore::Tensor> cu_seqlens;
     /// Block ids for each request `[batch, max_block_table_length]`. Used for paged cache.
@@ -21,28 +23,35 @@ struct AttentionMetadata {
     size_t max_query_length{0};
     /// Maximum total sequence length in the current batch.
     size_t max_sequence_length{0};
+    /// Maximum total sequence length in the current request batch.
+    std::optional<int64_t> max_context_len;
 
     AttentionMetadata() = default;
 
     AttentionMetadata(std::optional<infinicore::Tensor> past_sequence_lengths,
                       std::optional<infinicore::Tensor> total_sequence_lengths,
                       std::optional<infinicore::Tensor> input_offsets,
+                      std::optional<infinicore::Tensor> request_ids,
                       std::optional<infinicore::Tensor> cu_seqlens,
                       std::optional<infinicore::Tensor> block_tables,
                       std::optional<infinicore::Tensor> slot_mapping,
                       size_t max_query_length = 0,
-                      size_t max_sequence_length = 0) : past_sequence_lengths(past_sequence_lengths),
+                      size_t max_sequence_length = 0,
+                      std::optional<int64_t> max_context_len = std::nullopt) : past_sequence_lengths(past_sequence_lengths),
                                                         total_sequence_lengths(total_sequence_lengths),
                                                         input_offsets(input_offsets),
+                                                        request_ids(request_ids),
                                                         cu_seqlens(cu_seqlens),
                                                         block_tables(block_tables),
                                                         slot_mapping(slot_mapping),
                                                         max_query_length(max_query_length),
-                                                        max_sequence_length(max_sequence_length) {}
+                                                        max_sequence_length(max_sequence_length),
+                                                        max_context_len(max_context_len) {}
 
     AttentionMetadata(const infinilm::InfinilmModel::Input &input) : AttentionMetadata(input.past_sequence_lengths,
                                                                                        input.total_sequence_lengths,
                                                                                        input.input_offsets,
+                                                                                       input.request_ids,
                                                                                        input.cu_seqlens,
                                                                                        input.block_tables,
                                                                                        input.slot_mapping) {}
@@ -68,6 +77,9 @@ struct ForwardContext {
     MambaMetadata mamba_metadata;
     MultiModalMetadata mm_metadata;
     std::vector<infinicore::Tensor> kv_cache_vec;
+    std::vector<infinicore::Tensor> mla_vendor_cache_vec;
+    std::vector<infinicore::Tensor> indexer_cache_vec;
+    std::optional<infinicore::Tensor> dsa_topk_indices;
     std::vector<infinicore::Tensor> conv_state_vec;
     std::vector<infinicore::Tensor> ssm_state_vec;
 };

csrc/layers/linear/base_linear.hpp

@@ -35,6 +35,7 @@ class BaseLinear : public infinicore::nn::Module {
     void set_alpha(float alpha) { alpha_ = alpha; }
 
     // Accessors for parameters (backward compatible)
+    void release_parameters() { parameters_.clear(); }
     infinicore::Tensor weight() const;
     infinicore::Tensor bias() const;
     infinicore::Tensor weight_scale() const;

csrc/layers/linear/fused_linear.cpp

@@ -3,6 +3,78 @@
 #include <spdlog/spdlog.h>
 
 namespace infinilm::layers::linear {
+namespace {
+
+size_t merged_output_size(const std::vector<size_t> &output_sizes) {
+    size_t total = 0;
+    for (const auto size : output_sizes) {
+        total += size;
+    }
+    return total;
+}
+
+} // namespace
+
+// ---------------------------------------------------------
+// Merged Replicated Linear
+// ---------------------------------------------------------
+MergedReplicatedLinear::MergedReplicatedLinear(
+    size_t input_size,
+    const std::vector<size_t> &output_sizes,
+    const std::vector<std::string> &param_names,
+    RegisterParamFn register_fn,
+    std::shared_ptr<infinilm::quantization::BaseQuantization> quantization,
+    bool bias,
+    const infinicore::DataType &dtype,
+    const infinicore::Device &device)
+    : infinilm::nn::Linear(
+        input_size,
+        merged_output_size(output_sizes),
+        quantization == nullptr
+            ? std::make_shared<infinilm::quantization::NoneQuantization>()
+            : quantization,
+        bias,
+        dtype,
+        device),
+      output_sizes_(output_sizes),
+      register_fn_(std::move(register_fn)) {
+    if (output_sizes_.empty() || output_sizes_.size() != param_names.size()) {
+        throw std::runtime_error(
+            "MergedReplicatedLinear expects non-empty, equally sized output_sizes and param_names");
+    }
+    size_t offset = 0;
+    for (size_t i = 0; i < output_sizes_.size(); ++i) {
+        split_infos_.push_back({param_names[i], offset, output_sizes_[i]});
+        offset += output_sizes_[i];
+    }
+    auto params = this->split_params(split_infos_, 0, 1, -1);
+    for (auto &sp : params) {
+        register_fn_(sp.full_name, std::move(sp.param));
+    }
+}
+
+std::vector<infinicore::Tensor>
+MergedReplicatedLinear::forward_split(infinicore::Tensor &input) const {
+    auto output = this->forward(input);
+    const size_t dim = output->ndim() - 1;
+    std::vector<infinicore::Tensor> result;
+    result.reserve(output_sizes_.size());
+    size_t offset = 0;
+    for (const auto size : output_sizes_) {
+        result.push_back(output->narrow({{dim, offset, size}}));
+        offset += size;
+    }
+    return result;
+}
+
+void MergedReplicatedLinear::process_weights_after_loading() {
+    BaseLinear::process_weights_after_loading();
+    auto params = this->split_params(split_infos_, 0, 1, -1);
+    for (auto &sp : params) {
+        register_fn_(sp.full_name, std::move(sp.param));
+    }
+}
+
 // ---------------------------------------------------------
 // QKV Parallel Linear
 // ---------------------------------------------------------

csrc/layers/linear/fused_linear.hpp

@@ -7,6 +7,29 @@
 namespace infinilm::layers::linear {
 using RegisterParamFn = std::function<void(const std::string &, infinicore::nn::Parameter)>;
 
+// A replicated equivalent of vLLM's MergedColumnParallelLinear with
+// disable_tp=True. Checkpoint shards keep their original names while sharing
+// one runtime GEMM and one fused output buffer.
+class MergedReplicatedLinear : public infinilm::nn::Linear {
+public:
+    MergedReplicatedLinear(size_t input_size,
+                           const std::vector<size_t> &output_sizes,
+                           const std::vector<std::string> &param_names,
+                           RegisterParamFn register_fn,
+                           std::shared_ptr<infinilm::quantization::BaseQuantization> quantization = nullptr,
+                           bool bias = false,
+                           const infinicore::DataType &dtype = infinicore::DataType::F32,
+                           const infinicore::Device &device = infinicore::Device());
+
+    std::vector<infinicore::Tensor> forward_split(infinicore::Tensor &input) const;
+    void process_weights_after_loading() override;
+
+private:
+    std::vector<size_t> output_sizes_;
+    RegisterParamFn register_fn_;
+    std::vector<infinilm::quantization::SplitInfo> split_infos_;
+};
+
 class QKVParallelLinear : public infinilm::nn::ColumnParallelLinear {
 public:
     explicit QKVParallelLinear(size_t hidden_size,

csrc/layers/mlp/mlp.cpp

@@ -15,6 +15,8 @@ MLP::MLP(std::shared_ptr<infinilm::config::ModelConfig> model_config,
     const engine::distributed::RankInfo &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info();
     int tp_rank = rank_info.tp_rank;
     int tp_size = rank_info.tp_size;
+    const bool reduce_results = model_config->get_or<bool>("reduce_results", true);
+    auto communicator = reduce_results ? rank_info.comm : nullptr;
 
     auto quantization_method = model_config->get_quantization_method();
     auto register_fn = [this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); };
@@ -23,7 +25,7 @@ MLP::MLP(std::shared_ptr<infinilm::config::ModelConfig> model_config,
         quantization_method, use_bias_, dtype, device, rank_info);
     down_proj_ = this->register_module<layers::linear::RowParallelLinear>(
         "down_proj", intermediate_size_, hidden_size_, quantization_method,
-        use_bias_, dtype, device, tp_rank, tp_size, rank_info.comm);
+        use_bias_, dtype, device, tp_rank, tp_size, communicator);
 }
 
 infinicore::Tensor MLP::forward(const infinicore::Tensor &hidden_states) const {

csrc/layers/mlp/mlp.hpp

@@ -38,10 +38,12 @@ class MLP : public infinicore::nn::Module {
 
     void process_weights_after_loading() override {
         gate_up_proj_->process_weights_after_loading();
+        down_proj_->process_weights_after_loading();
     }
 
     void reset_runtime_state() const override {
         gate_up_proj_->reset_runtime_state();
+        down_proj_->reset_runtime_state();
     }
 
     // Module information

csrc/layers/moe/legacy/moe_mlp.hpp

@@ -18,7 +18,15 @@ class MoeMLP : public infinicore::nn::Module {
     infinicore::Tensor gate_weight() const { return gate_proj_->weight(); }
     infinicore::Tensor up_weight() const { return up_proj_->weight(); }
     infinicore::Tensor down_weight() const { return down_proj_->weight(); }
+    infinicore::Tensor gate_weight_scale() const { return gate_proj_->weight_scale(); }
+    infinicore::Tensor up_weight_scale() const { return up_proj_->weight_scale(); }
+    infinicore::Tensor down_weight_scale() const { return down_proj_->weight_scale(); }
     void set_alpha(float alpha) { down_proj_->set_alpha(alpha); }
+    void release_parameters() {
+        gate_proj_->release_parameters();
+        up_proj_->release_parameters();
+        down_proj_->release_parameters();
+    }
 
 protected:
     std::shared_ptr<infinilm::layers::linear::ColumnParallelLinear> gate_proj_;