Pepe/glm5 2

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

@@ -20,6 +20,10 @@ QuantConfig::get_quantization_method() const {
         return std::make_shared<infinilm::quantization::AWQ>(quantization_config);
     } else if (quant_method == "gptq") {
         return std::make_shared<infinilm::quantization::GPTQ>(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.hpp

@@ -3,6 +3,7 @@
 #include "graph_compiler.hpp"
 
 #include <unordered_map>
+#include <unordered_set>
 
 namespace infinilm::engine {
 class PagedCompiler : public GraphCompiler {
@@ -14,18 +15,25 @@ class PagedCompiler : public GraphCompiler {
     Compiled get_compiled(const InfinilmModel::Input &input) override;
 
 private:
-    std::vector<size_t> decode_batch_sizes_;
+    InfinilmModel::Input make_decode_input(size_t batch_size, size_t block_per_req) const;
+    void compile_decode(size_t batch_size, size_t block_per_req);
 
-    infinicore::Tensor block_tables_holder_;
+    bool initialized_ = false;
+    size_t num_blocks_ = 0;
+    size_t block_size_ = 0;
 
     struct CompiledResult {
         InfinilmModel::Input input;
         Compiled compiled;
+        infinicore::Tensor padding_total_sequence_lengths;
+        infinicore::Tensor padding_request_ids;
+        std::unordered_map<size_t, infinicore::Tensor> block_tables_staging;
     };
 
     std::unordered_map<
-        size_t, // num_requests
+        size_t, // static graph batch bucket
         CompiledResult>
         compiled_map_decode_;
+    std::unordered_set<size_t> graph_disabled_batches_;
 };
 } // namespace infinilm::engine

csrc/engine/compiler/static_batching_compiler.cpp

@@ -23,6 +23,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/distributed/communication_group.cpp

@@ -5,23 +5,68 @@ namespace infinilm::engine::distributed {
 
 CommunicationGroup::CommunicationGroup(const DistConfig &dist_config, infinicore::Device::Type device_type)
     : dist_config_(dist_config), device_type_(device_type),
-      communicators_(std::vector<infinicclComm_t>(dist_config.tp_device_ids.size(), nullptr)) {
+      tp_communicators_(dist_config.world_size(), nullptr),
+      pp_communicators_(dist_config.world_size(), nullptr),
+      world_communicators_(dist_config.pipeline_parallel_size > 1
+                               ? dist_config.world_size()
+                               : 0,
+                           nullptr) {
 
-    size_t world_size = dist_config_.tp_device_ids.size();
+    const int tp_size = dist_config_.tensor_parallel_size;
+    const int pp_size = dist_config_.pipeline_parallel_size;
+    const int world_size = dist_config_.world_size();
+    if (tp_size < 1 || pp_size < 1
+        || static_cast<int>(dist_config_.tp_device_ids.size()) != world_size) {
+        throw std::runtime_error(
+            "DistConfig device count must equal tensor_parallel_size * pipeline_parallel_size");
+    }
     size_t device_count = infinicore::context::getDeviceCount(device_type);
-    if (device_count < world_size) {
+    if (device_count < static_cast<size_t>(world_size)) {
         throw std::runtime_error("infinilm::engine::distributed::CommunicationGroup error, world size is larger than the number of available GPUs. world size: " + std::to_string(world_size) + ", device count: " + std::to_string(device_count));
     }
 
     if (infinicore::context::getDevice().getType() != device_type_) {
         infinicore::context::setDevice(infinicore::Device(device_type_, 0));
     }
-    if (world_size > 1) {
+    if (pp_size == 1 && world_size > 1) {
+        RUN_INFINI(infinicclCommInitAll(
+            (infiniDevice_t)infinicore::context::getDevice().getType(),
+            tp_communicators_.data(), world_size,
+            dist_config.tp_device_ids.data()));
+        return;
+    }
+
+    if (pp_size > 1) {
         RUN_INFINI(infinicclCommInitAll(
             (infiniDevice_t)infinicore::context::getDevice().getType(),
-            communicators_.data(),
-            dist_config.tp_device_ids.size(),
+            world_communicators_.data(), world_size,
             dist_config.tp_device_ids.data()));
+
+        for (int stage = 0; stage < pp_size; ++stage) {
+            std::vector<infinicclComm_t> stage_comms(tp_size, nullptr);
+            const int offset = stage * tp_size;
+            RUN_INFINI(infinicclCommInitAll(
+                (infiniDevice_t)infinicore::context::getDevice().getType(),
+                stage_comms.data(), tp_size,
+                dist_config.tp_device_ids.data() + offset));
+            for (int lane = 0; lane < tp_size; ++lane) {
+                tp_communicators_[offset + lane] = stage_comms[lane];
+            }
+        }
+
+        for (int lane = 0; lane < tp_size; ++lane) {
+            std::vector<int> lane_devices(pp_size);
+            std::vector<infinicclComm_t> lane_comms(pp_size, nullptr);
+            for (int stage = 0; stage < pp_size; ++stage) {
+                lane_devices[stage] = dist_config.tp_device_ids[stage * tp_size + lane];
+            }
+            RUN_INFINI(infinicclCommInitAll(
+                (infiniDevice_t)infinicore::context::getDevice().getType(),
+                lane_comms.data(), pp_size, lane_devices.data()));
+            for (int stage = 0; stage < pp_size; ++stage) {
+                pp_communicators_[stage * tp_size + lane] = lane_comms[stage];
+            }
+        }
     }
 }
 
@@ -31,21 +76,37 @@ const DistConfig &CommunicationGroup::get_dist_config() const {
 
 RankInfo CommunicationGroup::get_rank_info(int rank) const {
     RankInfo info;
-    info.tp_size = dist_config_.tp_device_ids.size();
-    info.tp_rank = rank;
+    info.tp_size = dist_config_.tensor_parallel_size;
+    info.tp_rank = rank % info.tp_size;
+    info.pp_size = dist_config_.pipeline_parallel_size;
+    info.pp_rank = rank / info.tp_size;
+    info.world_size = dist_config_.world_size();
+    info.global_rank = rank;
     info.device = infinicore::Device(device_type_, dist_config_.tp_device_ids[rank]);
-    info.comm = communicators_[rank];
+    info.comm = tp_communicators_[rank];
+    info.pp_comm = pp_communicators_[rank];
+    info.world_comm = world_communicators_.empty()
+                        ? tp_communicators_[rank]
+                        : world_communicators_[rank];
     return info;
 }
 
 int CommunicationGroup::get_world_size() const {
-    return dist_config_.tp_device_ids.size();
+    return dist_config_.world_size();
+}
+
+int CommunicationGroup::get_output_rank() const {
+    return (dist_config_.pipeline_parallel_size - 1)
+         * dist_config_.tensor_parallel_size;
 }
 
 CommunicationGroup::~CommunicationGroup() {
-    if (communicators_.size() > 1) {
-        for (auto &comm : communicators_) {
-            infinicclCommDestroy(comm);
+    for (auto *group : {&pp_communicators_, &tp_communicators_,
+                        &world_communicators_}) {
+        for (auto &comm : *group) {
+            if (comm != nullptr) {
+                infinicclCommDestroy(comm);
+            }
         }
     }
 }

csrc/engine/distributed/communication_group.hpp

@@ -18,15 +18,32 @@ struct RankInfo {
     int tp_size;
     // Tensor parallelism rank number of this rank
     int tp_rank;
-    // Communicator handle
+    int pp_size;
+    int pp_rank;
+    int world_size;
+    int global_rank;
+    // Tensor-, pipeline-, and world-parallel communicator handles.
     infinicclComm_t comm;
+    infinicclComm_t pp_comm;
+    infinicclComm_t world_comm;
 
     RankInfo(infinicore::Device _device = infinicore::context::getDevice())
-        : tp_size(1), tp_rank(0), device(_device), comm(nullptr){};
+        : tp_size(1), tp_rank(0), pp_size(1), pp_rank(0), world_size(1),
+          global_rank(0), device(_device), comm(nullptr), pp_comm(nullptr),
+          world_comm(nullptr){};
+
+    bool is_pipeline_first_stage() const { return pp_rank == 0; }
+    bool is_pipeline_last_stage() const { return pp_rank + 1 == pp_size; }
+    bool is_output_rank() const {
+        return is_pipeline_last_stage() && tp_rank == 0;
+    }
 
     std::string to_string() const {
         std::stringstream ss;
-        ss << "RankInfo: device=" << device.toString() << ", tp_size=" << tp_size << ", tp_rank=" << tp_rank;
+        ss << "RankInfo: device=" << device.toString()
+           << ", global_rank=" << global_rank << "/" << world_size
+           << ", tp_rank=" << tp_rank << "/" << tp_size
+           << ", pp_rank=" << pp_rank << "/" << pp_size;
         return ss.str();
     }
 };
@@ -42,12 +59,16 @@ class CommunicationGroup {
 
     int get_world_size() const;
 
+    int get_output_rank() const;
+
     ~CommunicationGroup();
 
 protected:
     DistConfig dist_config_;
     infinicore::Device::Type device_type_;
-    std::vector<infinicclComm_t> communicators_;
+    std::vector<infinicclComm_t> tp_communicators_;
+    std::vector<infinicclComm_t> pp_communicators_;
+    std::vector<infinicclComm_t> world_communicators_;
 };
 
 } // namespace infinilm::engine::distributed

csrc/engine/distributed/dist_config.cpp

@@ -1,18 +1,30 @@
 #include "dist_config.hpp"
 
+#include <stdexcept>
+
 namespace infinilm::engine::distributed {
 DistConfig::DistConfig()
     : tp_device_ids{0} {}
 
-DistConfig::DistConfig(int tp_size)
-    : tp_device_ids(tp_size, 0) {
-    for (int i = 0; i < tp_size; ++i) {
+DistConfig::DistConfig(int tp_size, int pp_size)
+    : tp_device_ids(tp_size * pp_size, 0),
+      tensor_parallel_size(tp_size),
+      pipeline_parallel_size(pp_size) {
+    if (tp_size < 1 || pp_size < 1) {
+        throw std::invalid_argument("TP and PP sizes must be positive");
+    }
+    for (int i = 0; i < tp_size * pp_size; ++i) {
         tp_device_ids[i] = i;
     }
 }
 
 DistConfig::DistConfig(const std::vector<int> &tp_device_ids_)
-    : tp_device_ids(tp_device_ids_) {}
+    : tp_device_ids(tp_device_ids_),
+      tensor_parallel_size(static_cast<int>(tp_device_ids_.size())) {}
+
+int DistConfig::world_size() const {
+    return tensor_parallel_size * pipeline_parallel_size;
+}
 
 DistConfig::operator std::string() const {
     std::string repr = "DistConfig(tp_device_ids=[";
@@ -22,7 +34,10 @@ DistConfig::operator std::string() const {
             repr += ", ";
         }
     }
-    repr += "], moe_ep_backend=" + moe_ep_backend + ", moe_ep_size=" + std::to_string(moe_ep_size) + ")";
+    repr += "], tp_size=" + std::to_string(tensor_parallel_size)
+          + ", pp_size=" + std::to_string(pipeline_parallel_size)
+          + ", moe_ep_backend=" + moe_ep_backend
+          + ", moe_ep_size=" + std::to_string(moe_ep_size) + ")";
     return repr;
 }
 

csrc/engine/distributed/dist_config.hpp

@@ -7,15 +7,19 @@
 namespace infinilm::engine::distributed {
 
 struct DistConfig {
-    // Device IDs for each rank in tensor parallelism
+    // Device IDs for all ranks, ordered as contiguous TP groups per PP stage.
     std::vector<int> tp_device_ids;
+    int tensor_parallel_size{1};
+    int pipeline_parallel_size{1};
     std::string moe_ep_backend{"disabled"};
     size_t moe_ep_size{1};
 
     DistConfig();
-    explicit DistConfig(int tp_size);
+    explicit DistConfig(int tp_size, int pp_size = 1);
     explicit DistConfig(const std::vector<int> &tp_device_ids_);
 
+    int world_size() const;
+
     explicit operator std::string() const;
 };
 

csrc/engine/infer_engine.cpp

@@ -170,6 +170,7 @@ InferEngine::Input::to_model_input(infinicore::Device device) const {
         to_device(past_sequence_lengths), // @todo: on device in the future
         to_device(total_sequence_lengths),
         to_device(input_offsets),
+        to_device(request_ids),
         to_device(cu_seqlens),
         to_device(block_tables),
         to_device(slot_mapping),
@@ -181,15 +182,19 @@ InferEngine::Input::to_model_input(infinicore::Device device) const {
         to_device_vec(image_grid_thw),
         image_req_ids,
         visual_token_ranges,
-        to_device(target_hidden_states)};
+        to_device(target_hidden_states),
+        sample_all_positions};
 
     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};
+        input.slot_mapping,
+        max_context_len,
+        is_mixed_batch};
 
     infinilm::global_state::get_forward_context().mamba_metadata = {
         input.input_offsets,
@@ -213,7 +218,7 @@ InferEngine::Output InferEngine::forward(const InferEngine::Input &input) {
         worker->wait();
     }
 
-    return workers_[0]->get_output();
+    return workers_[communication_group_.get_output_rank()]->get_output();
 }
 
 void InferEngine::compile() {

csrc/engine/rank_barrier.cpp

@@ -3,17 +3,20 @@
 namespace infinilm::engine {
 RankBarrier::RankBarrier(size_t num_ranks) : thread_count_(num_ranks), generation_(0), arrived_(0) {}
 
-void RankBarrier::wait() {
+bool RankBarrier::wait(bool success) {
     std::unique_lock<std::mutex> lock(mutex_);
-    int gen = generation_;
+    const size_t gen = generation_;
+    generation_success_ = generation_success_ && success;
 
     if (++arrived_ == thread_count_) {
-        // last thread
+        completed_results_.push_back(generation_success_);
         generation_++;
         arrived_ = 0;
+        generation_success_ = true;
         cv_.notify_all();
     } else {
         cv_.wait(lock, [&] { return gen != generation_; });
     }
+    return completed_results_[gen];
 }
 } // namespace infinilm::engine

csrc/engine/rank_barrier.hpp

@@ -2,18 +2,21 @@
 
 #include <condition_variable>
 #include <mutex>
+#include <vector>
 
 namespace infinilm::engine {
 class RankBarrier {
 public:
     explicit RankBarrier(size_t nranks);
 
-    void wait();
+    bool wait(bool success = true);
 
 private:
     const size_t thread_count_;
     size_t arrived_;
     size_t generation_;
+    bool generation_success_{true};
+    std::vector<bool> completed_results_;
     std::mutex mutex_;
     std::condition_variable cv_;
 };