diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp index f267794e2..1c8acdd35 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -1,9 +1,23 @@ #include "paged_compiler.hpp" #include "../../global_state/global_state.hpp" #include "../../utils.hpp" +#include "../workspace/workspace_context.hpp" + +#include 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 &model, RankBarrier *barrier) : GraphCompiler(model, barrier) { for (size_t b = 1; b < 64; ++b) { @@ -25,14 +39,19 @@ void PagedCompiler::compile() { size_t nblocks = dynamic_cast(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) { + auto make_decode_input = [&](size_t b, bool mrope_position_ids) { InfinilmModel::Input input; - input.input_ids = infinicore::Tensor::empty({1, b}, infinicore::DataType::I64, infinicore::context::getDevice()); - input.position_ids = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice()); + input.input_ids = infinicore::Tensor::empty({1, b}, infinicore::DataType::I32, 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.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()); @@ -66,8 +85,11 @@ void PagedCompiler::compile() { { const size_t warmup_batch_size = std::min(max_batch_size, static_cast(64)); - auto input = make_decode_input(warmup_batch_size); - model_->forward(input); + auto input = make_decode_input(warmup_batch_size, false); + { + WorkspaceForwardGuard forward_guard(maybe_current_workspace()); + model_->forward(input); + } infinicore::context::syncStream(); // Warmup runs the eager Marlin path and may leave per-layer lock // workspaces dirty. Reset before CUDA graph capture so capture @@ -76,11 +98,17 @@ void PagedCompiler::compile() { infinicore::context::syncStream(); } - for (size_t b : decode_batch_sizes_) { - auto input = make_decode_input(b); + auto capture_decode_graph = [&](size_t b, + bool mrope_position_ids, + std::unordered_map &compiled_map, + const std::string &scope_name) { + auto input = make_decode_input(b, mrope_position_ids); barrier_->wait(); - (void)model_->forward(input); + { + WorkspaceForwardGuard forward_guard(maybe_current_workspace()); + (void)model_->forward(input); + } infinicore::context::syncStream(); // Capture must not start with stale Marlin locks from previous // warmup/capture attempts. This reset is intentionally outside @@ -88,7 +116,15 @@ void PagedCompiler::compile() { // before every graph replay in get_compiled(). model_->reset_runtime_state(); infinicore::context::syncStream(); + barrier_->wait(); infinicore::context::startGraphRecording(); + WorkspaceCollectiveScopeGuard collective_scope_guard( + maybe_current_workspace(), + scope_name + std::to_string(b)); + WorkspaceForwardGuard forward_guard(maybe_current_workspace()); + // Capture runtime workspace resets so graph replay does not need + // an extra host-launched reset phase before every decode token. + model_->reset_runtime_state(); auto output = model_->forward(input); auto graph = infinicore::context::stopGraphRecording(); barrier_->wait(); @@ -96,7 +132,15 @@ void PagedCompiler::compile() { auto shared_output = std::shared_ptr( new InfinilmModel::Output{infinicore::graph::GraphTensor(output.logits)}); - compiled_map_decode_[b] = CompiledResult{std::move(input), std::make_tuple(graph, shared_output)}; + 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_, "paged_decode."); + if (compile_mrope_decode) { + capture_decode_graph(b, true, compiled_map_decode_mrope_, "paged_decode_mrope."); + } } } } @@ -110,12 +154,23 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & if (batch_size != input.input_ids.value()->size(1)) { return {nullptr, nullptr}; } else { - auto result = compiled_map_decode_.find(batch_size); - if (result == compiled_map_decode_.end()) { + 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()) { 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()); @@ -132,16 +187,12 @@ 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()); - // CUDA graph replay reuses the same per-layer Marlin workspaces. - // The graph itself does not contain a workspace reset, so enqueue - // one on the same stream before launch. This is correct but costs - // decode latency; the intended follow-up is a reusable global - // zero workspace/lock buffer shared by all Marlin layers. - model_->reset_runtime_state(); - auto graph = std::get<0>(result->second.compiled); auto shared_output = std::shared_ptr(new InfinilmModel::Output{std::get<1>(result->second.compiled)->logits->resume_from_blob_()}); diff --git a/csrc/engine/compiler/paged_compiler.hpp b/csrc/engine/compiler/paged_compiler.hpp index a1125864d..fc3baad61 100644 --- a/csrc/engine/compiler/paged_compiler.hpp +++ b/csrc/engine/compiler/paged_compiler.hpp @@ -27,5 +27,10 @@ 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 diff --git a/csrc/engine/compiler/static_batching_compiler.cpp b/csrc/engine/compiler/static_batching_compiler.cpp index af2f4799f..7d77e0c10 100644 --- a/csrc/engine/compiler/static_batching_compiler.cpp +++ b/csrc/engine/compiler/static_batching_compiler.cpp @@ -1,6 +1,7 @@ #include "static_batching_compiler.hpp" #include "../../cache/cache.hpp" #include "../../global_state/global_state.hpp" +#include "../workspace/workspace_context.hpp" namespace infinilm::engine { StaticBatchingCompiler::StaticBatchingCompiler(const std::shared_ptr &model, RankBarrier *barrier) @@ -11,7 +12,7 @@ void StaticBatchingCompiler::compile() { if (model_->get_cache_config() != nullptr && dynamic_cast(model_->get_cache_config())) { size_t b = dynamic_cast(model_->get_cache_config())->max_batch_size(); InfinilmModel::Input input; - input.input_ids = infinicore::Tensor::empty({b, 1}, infinicore::DataType::I64, infinicore::context::getDevice()); + input.input_ids = infinicore::Tensor::empty({b, 1}, infinicore::DataType::I32, infinicore::context::getDevice()); input.position_ids = infinicore::Tensor::empty({b, 1}, infinicore::DataType::I64, infinicore::context::getDevice()); input.past_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice()); input.total_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice()); @@ -29,10 +30,21 @@ void StaticBatchingCompiler::compile() { }; barrier_->wait(); - (void)model_->forward(input); + { + WorkspaceForwardGuard forward_guard(maybe_current_workspace()); + (void)model_->forward(input); + } infinicore::context::syncStream(); + barrier_->wait(); infinicore::context::startGraphRecording(); + WorkspaceCollectiveScopeGuard collective_scope_guard( + maybe_current_workspace(), + "static_batching." + std::to_string(b)); + WorkspaceForwardGuard forward_guard(maybe_current_workspace()); + // Capture runtime workspace resets so graph replay does not need + // an extra host-launched reset phase before every decode token. + model_->reset_runtime_state(); auto output = model_->forward(input); auto graph = infinicore::context::stopGraphRecording(); barrier_->wait(); diff --git a/csrc/engine/distributed/async_collective.cpp b/csrc/engine/distributed/async_collective.cpp new file mode 100644 index 000000000..04a11e3fd --- /dev/null +++ b/csrc/engine/distributed/async_collective.cpp @@ -0,0 +1,263 @@ +#include "async_collective.hpp" + +#include +#include + +#include +#include +#include +#include +#include + +namespace infinilm::engine::distributed { +namespace { + +thread_local AsyncCollectiveContext *current_context = nullptr; + +std::string trim(const std::string &value) { + std::size_t begin = 0; + while (begin < value.size() && std::isspace(static_cast(value[begin]))) { + ++begin; + } + std::size_t end = value.size(); + while (end > begin && std::isspace(static_cast(value[end - 1]))) { + --end; + } + return value.substr(begin, end - begin); +} + +std::string to_lower(std::string value) { + std::transform(value.begin(), value.end(), value.begin(), [](unsigned char ch) { + return static_cast(std::tolower(ch)); + }); + return value; +} + +void check_status(infiniStatus_t status, const char *what) { + if (status != INFINI_STATUS_SUCCESS) { + throw std::runtime_error(std::string(what) + " failed with status " + std::to_string(static_cast(status))); + } +} + +bool env_flag_enabled(const char *name) { + const char *raw = std::getenv(name); + if (raw == nullptr) { + return false; + } + const std::string value = to_lower(trim(raw)); + return value == "1" || value == "true" || value == "on" || value == "yes"; +} + +} // namespace + +struct AsyncCollectiveContext::EventSlot { + EventSlot(const infinicore::Device &device, std::uint64_t initial_generation) + : compute_ready(device, INFINIRT_EVENT_DISABLE_TIMING), + collective_done(device, INFINIRT_EVENT_DISABLE_TIMING), + generation(initial_generation) {} + + infinicore::DeviceEvent compute_ready; + infinicore::DeviceEvent collective_done; + std::uint64_t generation = 0; +}; + +bool async_collectives_enabled_by_env() { + return env_flag_enabled("INFINILM_ENABLE_ASYNC_COLLECTIVES") || + env_flag_enabled("INFINILM_ENABLE_COMM_OVERLAP"); +} + +bool async_collectives_force_enabled_by_env() { + return env_flag_enabled("INFINILM_COMM_OVERLAP_FORCE"); +} + +bool async_collectives_supported_by_policy(const infinicore::Device &device, int tp_size) { + if (device.getType() != infinicore::Device::Type::NVIDIA || tp_size <= 1) { + return false; + } + return tp_size == 2 || async_collectives_force_enabled_by_env(); +} + +bool async_collectives_enabled_for_rank(const infinicore::Device &device, int tp_size, infinicclComm_t communicator) { + return communicator != nullptr && + async_collectives_enabled_by_env() && + async_collectives_supported_by_policy(device, tp_size); +} + +AsyncCollectiveContext::AsyncCollectiveContext(const infinicore::Device &device, bool enabled) + : device_(device), enabled_(enabled) { + if (!enabled_) { + return; + } + + auto previous_device = infinicore::context::getDevice(); + if (previous_device != device_) { + infinicore::context::setDevice(device_); + } + + try { + check_status(infinirtStreamCreate(&stream_), "infinirtStreamCreate"); + const std::size_t event_pool_size = 64; + event_slots_.reserve(event_pool_size); + for (std::size_t i = 0; i < event_pool_size; ++i) { + event_slots_.push_back(std::make_unique(device_, generation_++)); + } + } catch (...) { + event_slots_.clear(); + if (stream_ != nullptr) { + (void)infinirtStreamDestroy(stream_); + stream_ = nullptr; + } + if (previous_device != device_) { + infinicore::context::setDevice(previous_device); + } + throw; + } + + if (previous_device != device_) { + infinicore::context::setDevice(previous_device); + } +} + +AsyncCollectiveContext::~AsyncCollectiveContext() { + if (stream_ == nullptr) { + return; + } + + try { + auto previous_device = infinicore::context::getDevice(); + if (previous_device != device_) { + infinicore::context::setDevice(device_); + } + (void)infinirtStreamSynchronize(stream_); + event_slots_.clear(); + (void)infinirtStreamDestroy(stream_); + stream_ = nullptr; + if (previous_device != device_) { + infinicore::context::setDevice(previous_device); + } + } catch (...) { + stream_ = nullptr; + } +} + +AsyncCollectiveContext::EventSlot &AsyncCollectiveContext::next_event_slot() { + if (event_slots_.empty()) { + throw std::runtime_error("async collective event pool is not initialized"); + } + EventSlot &slot = *event_slots_[next_slot_]; + next_slot_ = (next_slot_ + 1) % event_slots_.size(); + slot.generation = generation_++; + return slot; +} + +void AsyncCollectiveContext::validate_allreduce(infinicore::Tensor output, + const infinicore::Tensor &input, + infinicclComm_t communicator) const { + if (!enabled_ || stream_ == nullptr) { + throw std::runtime_error("async collective context is disabled"); + } + if (communicator == nullptr) { + throw std::runtime_error("async allreduce requires a non-null communicator"); + } + if (!output || !input) { + throw std::runtime_error("async allreduce received an empty tensor"); + } + if (output->device() != input->device() || output->device() != device_) { + throw std::runtime_error("async allreduce tensors must be on the context device"); + } + if (!output->is_contiguous() || !input->is_contiguous()) { + throw std::runtime_error("async allreduce requires contiguous tensors"); + } + if (output->numel() != input->numel()) { + throw std::runtime_error("async allreduce input/output numel mismatch"); + } + if (output->dtype() != input->dtype()) { + throw std::runtime_error("async allreduce input/output dtype mismatch"); + } +} + +PendingCollective AsyncCollectiveContext::allreduce(infinicore::Tensor output, + const infinicore::Tensor &input, + infinicclReduceOp_t op, + infinicclComm_t communicator) { + validate_allreduce(output, input, communicator); + + auto previous_device = infinicore::context::getDevice(); + if (previous_device != device_) { + infinicore::context::setDevice(device_); + } + + EventSlot &slot = next_event_slot(); + slot.compute_ready.record(infinicore::context::getStream()); + slot.compute_ready.wait(stream_); + + check_status(infinicclAllReduce( + const_cast(input->data()), + output->data(), + input->numel(), + static_cast(static_cast(input->dtype())), + op, + communicator, + stream_), + "infinicclAllReduce"); + + slot.collective_done.record(stream_); + + if (previous_device != device_) { + infinicore::context::setDevice(previous_device); + } + + return PendingCollective{(next_slot_ + event_slots_.size() - 1) % event_slots_.size(), slot.generation, true}; +} + +void AsyncCollectiveContext::wait(const PendingCollective &pending) { + if (!pending.valid) { + return; + } + if (pending.slot >= event_slots_.size()) { + throw std::runtime_error("invalid pending collective slot"); + } + EventSlot &slot = *event_slots_[pending.slot]; + if (slot.generation != pending.generation) { + throw std::runtime_error("pending collective event slot was reused before wait"); + } + + auto previous_device = infinicore::context::getDevice(); + if (previous_device != device_) { + infinicore::context::setDevice(device_); + } + slot.collective_done.wait(infinicore::context::getStream()); + if (previous_device != device_) { + infinicore::context::setDevice(previous_device); + } +} + +void AsyncCollectiveContext::allreduce_and_wait(infinicore::Tensor output, + const infinicore::Tensor &input, + infinicclReduceOp_t op, + infinicclComm_t communicator) { + auto pending = allreduce(output, input, op, communicator); + wait(pending); +} + +AsyncCollectiveContextGuard::AsyncCollectiveContextGuard(AsyncCollectiveContext *context) + : previous_(current_context) { + current_context = context; +} + +AsyncCollectiveContextGuard::~AsyncCollectiveContextGuard() { + current_context = previous_; +} + +AsyncCollectiveContext *maybe_current_async_collective_context() { + return current_context; +} + +AsyncCollectiveContext ¤t_async_collective_context() { + if (current_context == nullptr) { + throw std::runtime_error("no active async collective context"); + } + return *current_context; +} + +} // namespace infinilm::engine::distributed diff --git a/csrc/engine/distributed/async_collective.hpp b/csrc/engine/distributed/async_collective.hpp new file mode 100644 index 000000000..72beb39fe --- /dev/null +++ b/csrc/engine/distributed/async_collective.hpp @@ -0,0 +1,82 @@ +#pragma once + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace infinicore { +class DeviceEvent; +} + +namespace infinilm::engine::distributed { + +struct PendingCollective { + std::size_t slot = 0; + std::uint64_t generation = 0; + bool valid = false; +}; + +class AsyncCollectiveContext { +public: + explicit AsyncCollectiveContext(const infinicore::Device &device, bool enabled); + ~AsyncCollectiveContext(); + + AsyncCollectiveContext(const AsyncCollectiveContext &) = delete; + AsyncCollectiveContext &operator=(const AsyncCollectiveContext &) = delete; + + bool enabled() const { return enabled_; } + infinirtStream_t stream() const { return stream_; } + + PendingCollective allreduce(infinicore::Tensor output, + const infinicore::Tensor &input, + infinicclReduceOp_t op, + infinicclComm_t communicator); + void wait(const PendingCollective &pending); + void allreduce_and_wait(infinicore::Tensor output, + const infinicore::Tensor &input, + infinicclReduceOp_t op, + infinicclComm_t communicator); + +private: + struct EventSlot; + + EventSlot &next_event_slot(); + void validate_allreduce(infinicore::Tensor output, + const infinicore::Tensor &input, + infinicclComm_t communicator) const; + + infinicore::Device device_; + bool enabled_ = false; + infinirtStream_t stream_ = nullptr; + std::vector> event_slots_; + std::size_t next_slot_ = 0; + std::uint64_t generation_ = 1; +}; + +class AsyncCollectiveContextGuard { +public: + explicit AsyncCollectiveContextGuard(AsyncCollectiveContext *context); + ~AsyncCollectiveContextGuard(); + + AsyncCollectiveContextGuard(const AsyncCollectiveContextGuard &) = delete; + AsyncCollectiveContextGuard &operator=(const AsyncCollectiveContextGuard &) = delete; + +private: + AsyncCollectiveContext *previous_ = nullptr; +}; + +bool async_collectives_enabled_by_env(); +bool async_collectives_force_enabled_by_env(); +bool async_collectives_supported_by_policy(const infinicore::Device &device, int tp_size); +bool async_collectives_enabled_for_rank(const infinicore::Device &device, int tp_size, infinicclComm_t communicator); +AsyncCollectiveContext *maybe_current_async_collective_context(); +AsyncCollectiveContext ¤t_async_collective_context(); + +} // namespace infinilm::engine::distributed diff --git a/csrc/engine/distributed/communication_group.cpp b/csrc/engine/distributed/communication_group.cpp index 782faa9ec..8eac53c22 100644 --- a/csrc/engine/distributed/communication_group.cpp +++ b/csrc/engine/distributed/communication_group.cpp @@ -1,5 +1,6 @@ #include "communication_group.hpp" #include "../../utils.hpp" +#include namespace infinilm::engine::distributed { @@ -22,6 +23,12 @@ CommunicationGroup::CommunicationGroup(const DistConfig &dist_config, infinicore communicators_.data(), dist_config.tp_device_ids.size(), dist_config.tp_device_ids.data())); + for (auto *comm : communicators_) { + RUN_INFINI(infinicclCommSetAllReduceBackend(comm, dist_config_.allreduce_backend)); + } + if (dist_config_.allreduce_backend != INFINICCL_ALLREDUCE_BACKEND_NCCL) { + spdlog::info("InfiniCCL allreduce backend set to {}", std::string(dist_config_)); + } } } @@ -35,6 +42,7 @@ RankInfo CommunicationGroup::get_rank_info(int rank) const { info.tp_rank = rank; info.device = infinicore::Device(device_type_, dist_config_.tp_device_ids[rank]); info.comm = communicators_[rank]; + info.allreduce_backend = dist_config_.allreduce_backend; return info; } @@ -42,6 +50,14 @@ int CommunicationGroup::get_world_size() const { return dist_config_.tp_device_ids.size(); } +void CommunicationGroup::clear_registered_allreduce_buffers() { + if (communicators_.size() > 1) { + RUN_INFINI(infinicclCommClearAllReduceBuffers( + communicators_.data(), + static_cast(communicators_.size()))); + } +} + CommunicationGroup::~CommunicationGroup() { if (communicators_.size() > 1) { for (auto &comm : communicators_) { diff --git a/csrc/engine/distributed/communication_group.hpp b/csrc/engine/distributed/communication_group.hpp index e4f3c81a8..a7958d2c6 100644 --- a/csrc/engine/distributed/communication_group.hpp +++ b/csrc/engine/distributed/communication_group.hpp @@ -20,9 +20,11 @@ struct RankInfo { int tp_rank; // Communicator handle infinicclComm_t comm; + // User-requested allreduce backend for this TP group. + infinicclAllReduceBackend_t allreduce_backend; RankInfo(infinicore::Device _device = infinicore::context::getDevice()) - : tp_size(1), tp_rank(0), device(_device), comm(nullptr){}; + : tp_size(1), tp_rank(0), device(_device), comm(nullptr), allreduce_backend(INFINICCL_ALLREDUCE_BACKEND_NCCL){}; std::string to_string() const { std::stringstream ss; @@ -42,6 +44,8 @@ class CommunicationGroup { int get_world_size() const; + void clear_registered_allreduce_buffers(); + ~CommunicationGroup(); protected: diff --git a/csrc/engine/distributed/dist_config.cpp b/csrc/engine/distributed/dist_config.cpp index 44a73daf1..ad09f1b04 100644 --- a/csrc/engine/distributed/dist_config.cpp +++ b/csrc/engine/distributed/dist_config.cpp @@ -1,18 +1,35 @@ #include "dist_config.hpp" namespace infinilm::engine::distributed { +namespace { + +const char *allreduce_backend_name(infinicclAllReduceBackend_t backend) { + switch (backend) { + case INFINICCL_ALLREDUCE_BACKEND_AUTO: + return "auto"; + case INFINICCL_ALLREDUCE_BACKEND_NCCL: + return "nccl"; + case INFINICCL_ALLREDUCE_BACKEND_CUSTOM: + return "custom"; + default: + return "unknown"; + } +} + +} // namespace + DistConfig::DistConfig() - : tp_device_ids{0} {} + : tp_device_ids{0}, allreduce_backend(INFINICCL_ALLREDUCE_BACKEND_NCCL) {} -DistConfig::DistConfig(int tp_size) - : tp_device_ids(tp_size, 0) { +DistConfig::DistConfig(int tp_size, infinicclAllReduceBackend_t allreduce_backend_) + : tp_device_ids(tp_size, 0), allreduce_backend(allreduce_backend_) { for (int i = 0; i < tp_size; ++i) { tp_device_ids[i] = i; } } -DistConfig::DistConfig(const std::vector &tp_device_ids_) - : tp_device_ids(tp_device_ids_) {} +DistConfig::DistConfig(const std::vector &tp_device_ids_, infinicclAllReduceBackend_t allreduce_backend_) + : tp_device_ids(tp_device_ids_), allreduce_backend(allreduce_backend_) {} DistConfig::operator std::string() const { std::string repr = "DistConfig(tp_device_ids=["; @@ -22,7 +39,11 @@ DistConfig::operator std::string() const { repr += ", "; } } - repr += "], moe_ep_backend=" + moe_ep_backend + ", moe_ep_size=" + std::to_string(moe_ep_size) + ")"; + repr += "], moe_ep_backend=" + moe_ep_backend; + repr += ", moe_ep_size=" + std::to_string(moe_ep_size); + repr += ", allreduce_backend="; + repr += allreduce_backend_name(allreduce_backend); + repr += ")"; return repr; } diff --git a/csrc/engine/distributed/dist_config.hpp b/csrc/engine/distributed/dist_config.hpp index 4affcc9be..a47b6d077 100644 --- a/csrc/engine/distributed/dist_config.hpp +++ b/csrc/engine/distributed/dist_config.hpp @@ -1,6 +1,8 @@ #pragma once #include +#include + #include #include @@ -11,10 +13,11 @@ struct DistConfig { std::vector tp_device_ids; std::string moe_ep_backend{"disabled"}; size_t moe_ep_size{1}; + infinicclAllReduceBackend_t allreduce_backend; DistConfig(); - explicit DistConfig(int tp_size); - explicit DistConfig(const std::vector &tp_device_ids_); + explicit DistConfig(int tp_size, infinicclAllReduceBackend_t allreduce_backend_ = INFINICCL_ALLREDUCE_BACKEND_NCCL); + explicit DistConfig(const std::vector &tp_device_ids_, infinicclAllReduceBackend_t allreduce_backend_ = INFINICCL_ALLREDUCE_BACKEND_NCCL); explicit operator std::string() const; }; diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp index a5221eb37..1c862625f 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -204,16 +204,27 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { } InferEngine::Output InferEngine::forward(const InferEngine::Input &input) { + Input local_input = input; + if (last_output_ready_event_ && local_input.input_ids.has_value() && + local_input.input_ids.value()->device().getType() != infinicore::Device::Type::CPU) { + // Decode feeds the previous sampled GPU tensor back as input. Keep the + // worker stream ordered without forcing a host synchronization. + local_input.wait_event = last_output_ready_event_; + } + // Trigger each worker to run inference for (auto &worker : workers_) { - worker->run(input); + worker->run(local_input); } // Wait for all workers for (auto &worker : workers_) { worker->wait(); } - return workers_[0]->get_output(); + auto output = workers_[0]->get_output(); + last_output_ids_ = output.output_ids; + last_output_ready_event_ = output.ready_event; + return output; } void InferEngine::compile() { @@ -227,16 +238,39 @@ void InferEngine::compile() { for (auto &worker : workers_) { worker->wait(); } + // Keep graph-registered allreduce buffers alive with the compiled CUDA graph. + // Prefill/decode backend selection should be explicit instead of unregistering here. } //------------------------------------------------------ // Destructor //------------------------------------------------------ InferEngine::~InferEngine() { - // Close all workers + close(); +} + +void InferEngine::close() { + if (closed_) { + return; + } + closed_ = true; + if (last_saved_output_event_) { + last_saved_output_event_->synchronize(); + } + + for (auto &worker : workers_) { + worker->request_close(); + } for (auto &worker : workers_) { - worker->close(); + worker->join(); } + + last_output_ready_event_.reset(); + last_saved_output_event_.reset(); + last_output_ids_.reset(); + request_output_refs_.clear(); + workers_.clear(); + barrier_.reset(); } const distributed::DistConfig &InferEngine::get_dist_config() const { @@ -254,9 +288,62 @@ void InferEngine::reset_cache(const cache::CacheConfig *new_config) { worker->wait(); } cache_config_ = new_config->unique_copy(); + reset_request_state(); this->compile(); } +void InferEngine::reset_request_state() { + sync_last_output(); + last_output_ready_event_.reset(); + last_output_ids_.reset(); + last_saved_output_event_.reset(); + request_output_refs_.clear(); +} + +void InferEngine::sync_last_output() { + if (last_output_ready_event_) { + last_output_ready_event_->synchronize(); + } + if (last_saved_output_event_) { + last_saved_output_event_->synchronize(); + } + request_output_refs_.clear(); +} + +void InferEngine::copy_last_output_to(infinicore::Tensor dst) { + if (!last_output_ids_) { + throw std::runtime_error("No previous output tensor is available to copy"); + } + if (!dst) { + throw std::runtime_error("Destination output tensor is empty"); + } + if (dst->shape() != last_output_ids_->shape()) { + throw std::runtime_error( + "Cannot copy output with different shape. Src: " + last_output_ids_->info() + + " Dst: " + dst->info()); + } + if (!(dst->device() == last_output_ids_->device())) { + throw std::runtime_error( + "Destination output tensor must be on the same device as the sampled token. Src: " + + last_output_ids_->info() + " Dst: " + dst->info()); + } + + infinicore::context::setDevice(dst->device()); + if (last_saved_output_event_ && last_saved_output_event_->is_recorded() && last_saved_output_event_->query()) { + request_output_refs_.clear(); + } + if (last_output_ready_event_) { + infinicore::context::streamWaitEvent( + infinicore::context::getStream(), last_output_ready_event_->get()); + } + dst->copy_from(last_output_ids_); + request_output_refs_.push_back(last_output_ids_); + if (!last_saved_output_event_ || !(last_saved_output_event_->device() == dst->device())) { + last_saved_output_event_ = std::make_shared(dst->device()); + } + last_saved_output_event_->record(infinicore::context::getStream()); +} + std::vector> InferEngine::get_kv_cache() { std::vector> kv_cache_list; if (workers_.empty()) { diff --git a/csrc/engine/infer_engine.hpp b/csrc/engine/infer_engine.hpp index 4c0c0345c..b4bdf2195 100644 --- a/csrc/engine/infer_engine.hpp +++ b/csrc/engine/infer_engine.hpp @@ -55,6 +55,14 @@ class InferEngine { void reset_cache(const cache::CacheConfig *new_config); + void reset_request_state(); + + void sync_last_output(); + + void copy_last_output_to(infinicore::Tensor dst); + + void close(); + std::vector> get_kv_cache(); ~InferEngine(); @@ -74,6 +82,11 @@ class InferEngine { std::string weight_load_mode_ = "async"; bool weights_finalized_ = false; bool use_mla_{false}; + bool closed_{false}; + std::shared_ptr last_output_ready_event_; + infinicore::Tensor last_output_ids_; + std::shared_ptr last_saved_output_event_; + std::vector request_output_refs_; }; } // namespace infinilm::engine diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp index 11696e1eb..738aff014 100644 --- a/csrc/engine/rank_worker.cpp +++ b/csrc/engine/rank_worker.cpp @@ -1,8 +1,12 @@ #include "rank_worker.hpp" +#include "topology/device_topology.hpp" +#include "../utils.hpp" #include "../models/model_factory.hpp" +#include "workspace/workspace_context.hpp" #include "infinicore/ops.hpp" #include #include +#include namespace infinilm::engine { @@ -246,9 +250,9 @@ std::vector RankWorker::get_kv_cache() { } //------------------------------------------------------ -// close -- request shutdown and join thread +// request_close -- request shutdown without waiting for the thread //------------------------------------------------------ -void RankWorker::close() { +void RankWorker::request_close() { { std::lock_guard lock(mutex_); should_exit_ = true; @@ -256,12 +260,25 @@ void RankWorker::close() { job_cmd_ = Command::STOP; } cv_.notify_all(); +} +//------------------------------------------------------ +// join -- wait for worker thread shutdown +//------------------------------------------------------ +void RankWorker::join() { if (thread_.joinable()) { thread_.join(); } } +//------------------------------------------------------ +// close -- request shutdown and join thread +//------------------------------------------------------ +void RankWorker::close() { + request_close(); + join(); +} + //------------------------------------------------------ // get_output (thread safe) //------------------------------------------------------ @@ -280,6 +297,33 @@ void RankWorker::thread_loop() { // Initialize device & model outside of holding the main mutex to avoid blocking callers. infinicore::context::setDevice(rank_info_.device); + auto affinity_binding = topology::bind_current_thread_to_device_numa(rank_info_.device); + if (affinity_binding.applied) { + spdlog::info( + "{} bound worker thread to NUMA node {} CPUs {} via {}{}", + rank_info_.to_string(), + affinity_binding.numa_node, + affinity_binding.cpu_list, + affinity_binding.provider, + affinity_binding.pci_bus_id.empty() ? "" : " pci=" + affinity_binding.pci_bus_id); + } else if (affinity_binding.attempted) { + spdlog::debug( + "{} skipped worker CPU affinity binding via {}: {}", + rank_info_.to_string(), + affinity_binding.provider, + affinity_binding.reason); + } + workspace_manager_ = std::make_unique(rank_info_.device); + const bool enable_async_collectives = distributed::async_collectives_enabled_for_rank( + rank_info_.device, + rank_info_.tp_size, + rank_info_.comm); + async_collective_context_ = std::make_unique( + rank_info_.device, + enable_async_collectives); + if (async_collective_context_->enabled()) { + spdlog::info("{} enabled async collective context", rank_info_.to_string()); + } // Initialize global enviromnet. infinilm::global_state::initialize_model_parallel(rank_info_); @@ -413,22 +457,41 @@ void RankWorker::thread_loop() { } else if (local_cmd == Command::RUN) { try { { + WorkspaceForwardGuard forward_guard(workspace_manager_.get()); + WorkspaceContextGuard workspace_guard(workspace_manager_.get()); + distributed::AsyncCollectiveContextGuard async_collective_guard(async_collective_context_.get()); std::lock_guard lk(mutex_); infinicore::Tensor logits; infinicore::Tensor hidden_states; - // All-position speculative/MTP runs need eager mode because - // hidden states are not part of compiled graph outputs. - if (!local_args.sample_all_positions && compiler_ != nullptr) { - auto [graph, output] = compiler_->get_compiled(local_args.to_model_input(infinicore::Device::cpu())); - if (graph != nullptr && output != nullptr) { - graph->run(); - logits = output->logits; - } + if (local_args.wait_event) { + infinicore::context::setDevice(rank_info_.device); + infinicore::context::streamWaitEvent( + infinicore::context::getStream(), local_args.wait_event->get()); + } + auto model_args = local_args.to_model_input(rank_info_.device); + std::shared_ptr graph; + std::shared_ptr graph_output; + if (compiler_ != nullptr && !local_args.sample_all_positions) { + std::tie(graph, graph_output) = compiler_->get_compiled(model_args); + } + + const bool use_compiled_decode_graph = graph != nullptr && graph_output != nullptr; + if (rank_info_.comm != nullptr && + rank_info_.allreduce_backend == INFINICCL_ALLREDUCE_BACKEND_CUSTOM) { + RUN_INFINI(infinicclCommSetAllReduceBackend( + rank_info_.comm, + use_compiled_decode_graph ? INFINICCL_ALLREDUCE_BACKEND_CUSTOM + : INFINICCL_ALLREDUCE_BACKEND_NCCL)); + } + + if (use_compiled_decode_graph) { + graph->run(); + logits = graph_output->logits; } - // Fall back to eager mode + // Fall back to eager mode. This covers prefill and unsupported decode shapes; + // when custom was requested, it is forced to NCCL above for this eager path. if (!logits) { - auto model_args = local_args.to_model_input(rank_info_.device); auto model_output = model_->forward(model_args); logits = model_output.logits; hidden_states = model_output.hidden_states; @@ -449,9 +512,13 @@ void RankWorker::thread_loop() { int32_t *input_offsets = (int32_t *)local_args.input_offsets.value()->data(); const bool sample_all_positions = local_args.sample_all_positions; - const size_t n_out = sample_all_positions ? static_cast(input_offsets[n_req]) : n_req; - auto output_ids{infinicore::Tensor::empty({n_out}, infinicore::DataType::I64, rank_info_.device)}; - + const size_t n_out = sample_all_positions + ? static_cast(input_offsets[n_req]) + : n_req; + const auto output_dtype = sample_all_positions + ? infinicore::DataType::I64 + : infinicore::DataType::I32; + auto output_ids{infinicore::Tensor::empty({n_out}, output_dtype, rank_info_.device)}; for (size_t i{0}; i < n_out; ++i) { size_t score_idx = i; if (!sample_all_positions) { @@ -464,12 +531,9 @@ void RankWorker::thread_loop() { out, score, random_val, top_p, top_k, temperature); } - output_ids = output_ids->to(infinicore::Device::cpu()); - - infinicore::context::syncStream(); - - auto out{Output{output_ids, logits, hidden_states}}; - + auto ready_event = std::make_shared(rank_info_.device); + ready_event->record(infinicore::context::getStream()); + auto out{Output{output_ids, logits, hidden_states, ready_event}}; output_ = std::move(out); } @@ -509,6 +573,8 @@ void RankWorker::thread_loop() { } else if (local_cmd == Command::COMPILE) { try { if (compiler_ != nullptr) { + WorkspaceContextGuard workspace_guard(workspace_manager_.get()); + distributed::AsyncCollectiveContextGuard async_collective_guard(async_collective_context_.get()); compiler_->compile(); } { @@ -532,8 +598,20 @@ void RankWorker::thread_loop() { // Shouldn't reach here (no-op) } } // while - // Some clean up should be done before exiting the thread + // Release graph/model-owned GPU resources on the worker's CUDA context + // instead of leaving them to Python interpreter shutdown. + infinicore::context::setDevice(rank_info_.device); + infinicore::context::syncStream(); + { + std::lock_guard lk(mutex_); + output_ = Output{}; + pending_args_ = Input{}; + } compiler_.reset(); + model_.reset(); + workspace_manager_.reset(); + async_collective_context_.reset(); + infinicore::context::syncStream(); } catch (const std::exception &e) { // Top-level exception: ensure any waiters are woken and the thread exits cleanly. { diff --git a/csrc/engine/rank_worker.hpp b/csrc/engine/rank_worker.hpp index d396ef6f1..de1ad12e8 100644 --- a/csrc/engine/rank_worker.hpp +++ b/csrc/engine/rank_worker.hpp @@ -6,8 +6,11 @@ #include "../global_state/global_state.hpp" #include "../models/model_factory.hpp" #include "compiler/general_compiler.hpp" +#include "distributed/async_collective.hpp" #include "distributed/distributed.hpp" #include "rank_barrier.hpp" +#include "workspace/inference_workspace_manager.hpp" +#include "infinicore/device_event.hpp" #include #include @@ -79,6 +82,9 @@ class RankWorker { float top_p{1}; + // GPU relay dependency from the previous sampled token. This is not exposed to Python. + std::shared_ptr wait_event; + infinilm::InfinilmModel::Input to_model_input(infinicore::Device device) const; }; @@ -86,6 +92,7 @@ class RankWorker { infinicore::Tensor output_ids; infinicore::Tensor logits; infinicore::Tensor hidden_states; + std::shared_ptr ready_event; }; RankWorker(std::shared_ptr infinilm_config, @@ -126,6 +133,12 @@ class RankWorker { // Wait until run job completes. The result can be retrieved with get_output(). void wait(); + // Request worker shutdown. This only signals the worker and returns immediately. + void request_close(); + + // Join the worker thread after shutdown has been requested. + void join(); + // Request worker shutdown and join the thread. void close(); @@ -145,6 +158,8 @@ class RankWorker { ForwardContext forward_context_; std::shared_ptr model_; std::shared_ptr cache_; + std::unique_ptr workspace_manager_; + std::unique_ptr async_collective_context_; // Backends backends::AttentionBackend attention_backend_; diff --git a/csrc/engine/topology/device_topology.cpp b/csrc/engine/topology/device_topology.cpp new file mode 100644 index 000000000..19ce03e15 --- /dev/null +++ b/csrc/engine/topology/device_topology.cpp @@ -0,0 +1,435 @@ +#include "device_topology.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef __linux__ +#include +#include +#include +#endif + +namespace infinilm::engine::topology { +namespace { + +std::string trim(const std::string &value) { + std::size_t begin = 0; + while (begin < value.size() && std::isspace(static_cast(value[begin]))) { + ++begin; + } + std::size_t end = value.size(); + while (end > begin && std::isspace(static_cast(value[end - 1]))) { + --end; + } + return value.substr(begin, end - begin); +} + +std::string to_lower(std::string value) { + std::transform(value.begin(), value.end(), value.begin(), [](unsigned char ch) { + return static_cast(std::tolower(ch)); + }); + return value; +} + +bool env_truthy(const char *name) { + const char *raw = std::getenv(name); + if (raw == nullptr) { + return false; + } + const std::string value = to_lower(trim(raw)); + return value == "1" || value == "true" || value == "on" || value == "yes"; +} + +std::optional getenv_nonempty(const char *name) { + const char *raw = std::getenv(name); + if (raw == nullptr) { + return std::nullopt; + } + std::string value = trim(raw); + if (value.empty()) { + return std::nullopt; + } + return value; +} + +std::vector split_mapping_entries(const std::string &value) { + std::vector entries; + std::string current; + for (char ch : value) { + if (ch == ',' || ch == ';') { + current = trim(current); + if (!current.empty()) { + entries.push_back(current); + } + current.clear(); + } else { + current.push_back(ch); + } + } + current = trim(current); + if (!current.empty()) { + entries.push_back(current); + } + return entries; +} + +std::optional lookup_device_mapping(const char *env_name, const infinicore::Device &device) { + auto env_value = getenv_nonempty(env_name); + if (!env_value.has_value()) { + return std::nullopt; + } + + const std::string type_name = to_lower(infinicore::Device::toString(device.getType())); + const std::string index = std::to_string(device.getIndex()); + const std::vector expected_keys = { + type_name + ":" + index, + type_name + "." + index, + type_name + index, + index, + "*", + }; + + for (const auto &entry : split_mapping_entries(*env_value)) { + const auto pos = entry.find('='); + if (pos == std::string::npos) { + continue; + } + const std::string key = to_lower(trim(entry.substr(0, pos))); + const std::string value = trim(entry.substr(pos + 1)); + if (value.empty()) { + continue; + } + if (std::find(expected_keys.begin(), expected_keys.end(), key) != expected_keys.end()) { + return value; + } + } + return std::nullopt; +} + +std::optional parse_int(const std::string &value) { + try { + std::size_t consumed = 0; + int parsed = std::stoi(value, &consumed, 10); + if (consumed == value.size()) { + return parsed; + } + } catch (...) { + } + return std::nullopt; +} + +#ifdef __linux__ + +struct PciBusIdQuery { + std::string bus_id; + std::string provider; +}; + +std::optional read_first_line(const std::string &path) { + std::ifstream file(path); + if (!file.good()) { + return std::nullopt; + } + std::string line; + std::getline(file, line); + line = trim(line); + if (line.empty()) { + return std::nullopt; + } + return line; +} + +std::vector pci_bus_id_candidates(std::string bus_id) { + bus_id = to_lower(trim(bus_id)); + std::vector candidates; + if (bus_id.empty()) { + return candidates; + } + + candidates.push_back(bus_id); + const auto colon = bus_id.find(':'); + if (colon == std::string::npos) { + candidates.push_back("0000:" + bus_id); + } else if (colon == 8) { + candidates.push_back(bus_id.substr(4)); + } else if (colon != 4) { + candidates.push_back("0000:" + bus_id); + } + + std::sort(candidates.begin(), candidates.end()); + candidates.erase(std::unique(candidates.begin(), candidates.end()), candidates.end()); + return candidates; +} + +std::optional query_nvidia_pci_bus_id(const infinicore::Device &device) { + const char *library_names[] = {"libcudart.so", "libcudart.so.12", "libcudart.so.11.0"}; + void *handle = nullptr; + for (const char *name : library_names) { + handle = dlopen(name, RTLD_LAZY | RTLD_LOCAL); + if (handle != nullptr) { + break; + } + } + if (handle == nullptr) { + return std::nullopt; + } + + using GetPciBusIdFn = int (*)(char *, int, int); + auto get_pci_bus_id = reinterpret_cast(dlsym(handle, "cudaDeviceGetPCIBusId")); + if (get_pci_bus_id == nullptr) { + dlclose(handle); + return std::nullopt; + } + + char bus_id[64] = {}; + int status = get_pci_bus_id(bus_id, static_cast(sizeof(bus_id)), static_cast(device.getIndex())); + dlclose(handle); + if (status != 0 || bus_id[0] == '\0') { + return std::nullopt; + } + return PciBusIdQuery{bus_id, "cuda-runtime"}; +} + +std::optional query_device_pci_bus_id(const infinicore::Device &device) { + if (auto override_value = lookup_device_mapping("INFINILM_DEVICE_PCI_BUS_IDS", device)) { + return PciBusIdQuery{*override_value, "env-pci"}; + } + + switch (device.getType()) { + case infinicore::Device::Type::NVIDIA: + return query_nvidia_pci_bus_id(device); + case infinicore::Device::Type::HYGON: + case infinicore::Device::Type::ILUVATAR: + case infinicore::Device::Type::METAX: + case infinicore::Device::Type::MOORE: + case infinicore::Device::Type::CAMBRICON: + case infinicore::Device::Type::ASCEND: + case infinicore::Device::Type::KUNLUN: + case infinicore::Device::Type::QY: + case infinicore::Device::Type::ALI: + case infinicore::Device::Type::CPU: + case infinicore::Device::Type::COUNT: + return std::nullopt; + } + return std::nullopt; +} + +std::optional query_numa_node_from_pci_bus_id(const std::string &bus_id) { + for (const auto &candidate : pci_bus_id_candidates(bus_id)) { + auto numa_node = read_first_line("/sys/bus/pci/devices/" + candidate + "/numa_node"); + if (!numa_node.has_value()) { + continue; + } + auto parsed = parse_int(*numa_node); + if (parsed.has_value()) { + return parsed; + } + } + return std::nullopt; +} + +std::optional read_numa_cpu_list(int numa_node) { + if (numa_node < 0) { + return std::nullopt; + } + return read_first_line("/sys/devices/system/node/node" + std::to_string(numa_node) + "/cpulist"); +} + +bool add_cpu_range(cpu_set_t &set, int begin, int end) { + if (begin < 0 || end < begin) { + return false; + } + for (int cpu = begin; cpu <= end; ++cpu) { + if (cpu >= CPU_SETSIZE) { + return false; + } + CPU_SET(cpu, &set); + } + return true; +} + +bool parse_cpu_list(const std::string &cpu_list, cpu_set_t &set) { + CPU_ZERO(&set); + bool any = false; + std::stringstream ss(cpu_list); + std::string token; + while (std::getline(ss, token, ',')) { + token = trim(token); + if (token.empty()) { + continue; + } + const auto dash = token.find('-'); + if (dash == std::string::npos) { + auto cpu = parse_int(token); + if (!cpu.has_value() || !add_cpu_range(set, *cpu, *cpu)) { + return false; + } + } else { + auto begin = parse_int(trim(token.substr(0, dash))); + auto end = parse_int(trim(token.substr(dash + 1))); + if (!begin.has_value() || !end.has_value() || !add_cpu_range(set, *begin, *end)) { + return false; + } + } + any = true; + } + return any; +} + +std::string format_cpu_set(const cpu_set_t &set) { + std::ostringstream out; + bool first = true; + int cpu = 0; + while (cpu < CPU_SETSIZE) { + if (!CPU_ISSET(cpu, &set)) { + ++cpu; + continue; + } + const int begin = cpu; + while (cpu + 1 < CPU_SETSIZE && CPU_ISSET(cpu + 1, &set)) { + ++cpu; + } + const int end = cpu; + if (!first) { + out << ","; + } + out << begin; + if (end != begin) { + out << "-" << end; + } + first = false; + ++cpu; + } + return out.str(); +} + +bool intersect_with_current_affinity(cpu_set_t &target) { + cpu_set_t current; + CPU_ZERO(¤t); + if (sched_getaffinity(0, sizeof(current), ¤t) != 0) { + return false; + } + + bool any = false; + for (int cpu = 0; cpu < CPU_SETSIZE; ++cpu) { + const bool keep = CPU_ISSET(cpu, &target) && CPU_ISSET(cpu, ¤t); + if (!keep) { + CPU_CLR(cpu, &target); + } else { + any = true; + } + } + return any; +} + +CpuAffinityBinding apply_cpu_list(std::string cpu_list, int numa_node, std::string provider, std::string pci_bus_id) { + CpuAffinityBinding result; + result.attempted = true; + result.numa_node = numa_node; + result.provider = std::move(provider); + result.pci_bus_id = std::move(pci_bus_id); + + cpu_set_t target; + if (!parse_cpu_list(cpu_list, target)) { + result.reason = "invalid CPU list: " + cpu_list; + return result; + } + if (!intersect_with_current_affinity(target)) { + result.reason = "target CPU list is outside the process affinity mask: " + cpu_list; + return result; + } + if (sched_setaffinity(0, sizeof(target), &target) != 0) { + result.reason = std::string("sched_setaffinity failed: ") + std::strerror(errno); + return result; + } + + result.applied = true; + result.cpu_list = format_cpu_set(target); + return result; +} + +CpuAffinityBinding bind_current_thread_to_device_numa_linux(const infinicore::Device &device) { + CpuAffinityBinding result; + if (env_truthy("INFINILM_DISABLE_CPU_AFFINITY")) { + result.reason = "disabled by INFINILM_DISABLE_CPU_AFFINITY"; + return result; + } + + if (auto cpu_list = getenv_nonempty("INFINILM_CPU_AFFINITY_CPUS")) { + return apply_cpu_list(*cpu_list, -1, "env-cpu-list", ""); + } + + if (auto numa_override = lookup_device_mapping("INFINILM_DEVICE_NUMA_NODES", device)) { + auto parsed_node = parse_int(*numa_override); + if (!parsed_node.has_value()) { + result.attempted = true; + result.provider = "env-numa"; + result.reason = "invalid NUMA node: " + *numa_override; + return result; + } + auto cpu_list = read_numa_cpu_list(*parsed_node); + if (!cpu_list.has_value()) { + result.attempted = true; + result.provider = "env-numa"; + result.numa_node = *parsed_node; + result.reason = "cannot read CPU list for NUMA node " + std::to_string(*parsed_node); + return result; + } + return apply_cpu_list(*cpu_list, *parsed_node, "env-numa", ""); + } + + auto pci_bus_id = query_device_pci_bus_id(device); + if (!pci_bus_id.has_value()) { + result.attempted = true; + result.provider = "topology-provider"; + result.reason = "no PCI bus id provider for " + device.toString(); + return result; + } + + auto numa_node = query_numa_node_from_pci_bus_id(pci_bus_id->bus_id); + if (!numa_node.has_value() || *numa_node < 0) { + result.attempted = true; + result.provider = pci_bus_id->provider; + result.pci_bus_id = pci_bus_id->bus_id; + result.reason = "cannot resolve NUMA node for PCI bus id " + pci_bus_id->bus_id; + return result; + } + + auto cpu_list = read_numa_cpu_list(*numa_node); + if (!cpu_list.has_value()) { + result.attempted = true; + result.provider = pci_bus_id->provider; + result.pci_bus_id = pci_bus_id->bus_id; + result.numa_node = *numa_node; + result.reason = "cannot read CPU list for NUMA node " + std::to_string(*numa_node); + return result; + } + + return apply_cpu_list(*cpu_list, *numa_node, pci_bus_id->provider, pci_bus_id->bus_id); +} + +#else + +CpuAffinityBinding bind_current_thread_to_device_numa_linux(const infinicore::Device &) { + CpuAffinityBinding result; + result.reason = "CPU affinity binding is only supported on Linux"; + return result; +} + +#endif + +} // namespace + +CpuAffinityBinding bind_current_thread_to_device_numa(const infinicore::Device &device) { + return bind_current_thread_to_device_numa_linux(device); +} + +} // namespace infinilm::engine::topology diff --git a/csrc/engine/topology/device_topology.hpp b/csrc/engine/topology/device_topology.hpp new file mode 100644 index 000000000..3be1d8dea --- /dev/null +++ b/csrc/engine/topology/device_topology.hpp @@ -0,0 +1,21 @@ +#pragma once + +#include + +#include + +namespace infinilm::engine::topology { + +struct CpuAffinityBinding { + bool attempted = false; + bool applied = false; + int numa_node = -1; + std::string cpu_list; + std::string pci_bus_id; + std::string provider; + std::string reason; +}; + +CpuAffinityBinding bind_current_thread_to_device_numa(const infinicore::Device &device); + +} // namespace infinilm::engine::topology diff --git a/csrc/engine/workspace/inference_workspace_manager.cpp b/csrc/engine/workspace/inference_workspace_manager.cpp new file mode 100644 index 000000000..0e11ed85b --- /dev/null +++ b/csrc/engine/workspace/inference_workspace_manager.cpp @@ -0,0 +1,444 @@ +#include "inference_workspace_manager.hpp" + +#include "../../utils.hpp" + +#include "infinicore/context/context.hpp" + +#include +#include + +namespace infinilm::engine { + +namespace { + +size_t numel_of(const infinicore::Shape &shape) { + size_t numel = 1; + for (auto dim : shape) { + numel *= dim; + } + return numel; +} + +class ZeroTensorGraphOperator final : public infinicore::graph::GraphOperator { +public: + explicit ZeroTensorGraphOperator(const infinicore::Tensor &tensor) + : tensor_(tensor), bytes_(tensor->nbytes()) {} + + void run() const override { + infinicore::context::setDeviceMemoryAsync( + const_cast(tensor_->data()), 0, bytes_, infinicore::context::getStream()); + } + +private: + infinicore::graph::GraphTensor tensor_; + size_t bytes_ = 0; +}; + +} // namespace + +infinicore::Tensor WorkspaceBuffer::as_tensor(const infinicore::Shape &shape, infinicore::DataType dtype) const { + const auto requested_bytes = numel_of(shape) * infinicore::dsize(dtype); + if (requested_bytes > bytes) { + throw std::runtime_error("workspace buffer is too small for requested tensor view"); + } + if (owner && dtype == infinicore::DataType::U8 && shape.size() == 1 && shape[0] == requested_bytes) { + return owner->narrow({{0, 0, requested_bytes}}); + } + return infinicore::Tensor::from_blob(ptr, shape, dtype, device); +} + +InferenceWorkspaceManager::InferenceWorkspaceManager(infinicore::Device device) : device_(device) {} + +void InferenceWorkspaceManager::begin_forward() { + if (in_forward_) { + throw std::runtime_error("InferenceWorkspaceManager::begin_forward called while a forward is active"); + } + reclaim_retired_typed_blocks(); + in_forward_ = true; + transient_cursor_ = 0; + collective_cursor_ = 0; +} + +void InferenceWorkspaceManager::end_forward() { + for (auto it = transient_active_.rbegin(); it != transient_active_.rend(); ++it) { + auto &block = *it; + block.active = false; + if (block.typed) { + if (block.tensor.use_count() == 1) { + typed_transient_free_[block.key].push_back(std::move(block)); + } else { + transient_retired_.push_back(std::move(block)); + } + } else { + transient_free_[block.capacity].push_back(std::move(block)); + } + } + transient_active_.clear(); + in_forward_ = false; +} + +void InferenceWorkspaceManager::begin_warmup() { + if (in_forward_) { + throw std::runtime_error("cannot begin workspace warmup while a forward is active"); + } + transient_plan_.clear(); + transient_cursor_ = 0; + mode_ = WorkspaceMode::Warmup; +} + +void InferenceWorkspaceManager::end_warmup() { + if (in_forward_) { + throw std::runtime_error("cannot end workspace warmup while a forward is active"); + } + mode_ = WorkspaceMode::Dynamic; + transient_cursor_ = 0; +} + +void InferenceWorkspaceManager::freeze_transient_pool() { + if (in_forward_) { + throw std::runtime_error("cannot freeze workspace pool while a forward is active"); + } + mode_ = WorkspaceMode::Frozen; + transient_cursor_ = 0; +} + +void InferenceWorkspaceManager::unfreeze_transient_pool() { + mode_ = WorkspaceMode::Dynamic; + transient_cursor_ = 0; +} + +WorkspaceBuffer InferenceWorkspaceManager::acquire_temp_buffer( + std::string_view tag, + size_t bytes, + size_t alignment, + WorkspaceZeroPolicy zero_policy) { + if (!in_forward_) { + throw std::runtime_error("transient inference workspace can only be acquired during forward: tag=" + std::string(tag) + ", scope=" + collective_scope_ + ", graph_recording=" + (infinicore::context::isGraphRecording() ? std::string("true") : std::string("false"))); + } + const auto capacity = normalize_bytes(bytes, alignment); + auto &free_list = transient_free_[capacity]; + + if (mode_ == WorkspaceMode::Warmup) { + transient_plan_.push_back(TransientRequest{std::string(tag), bytes, capacity}); + } else if (mode_ == WorkspaceMode::Frozen) { + if (transient_cursor_ >= transient_plan_.size()) { + throw std::runtime_error("frozen workspace pool received more transient requests than warmup recorded"); + } + const auto &expected = transient_plan_[transient_cursor_]; + if (expected.capacity != capacity) { + throw std::runtime_error( + "frozen workspace pool request size mismatch for " + std::string(tag) + + ": expected capacity " + std::to_string(expected.capacity) + + ", got " + std::to_string(capacity)); + } + } + ++transient_cursor_; + + Block block; + if (!free_list.empty()) { + block = std::move(free_list.back()); + free_list.pop_back(); + } else { + if (mode_ == WorkspaceMode::Frozen) { + throw std::runtime_error("frozen workspace pool has no free block for " + std::string(tag)); + } + block.tensor = infinicore::Tensor::empty({capacity}, infinicore::DataType::U8, device_); + block.capacity = capacity; + block.lifetime = WorkspaceLifetime::Transient; + } + + block.owner = "transient"; + block.key = std::string(tag); + block.active = true; + if (zero_policy == WorkspaceZeroPolicy::OnAcquire || zero_policy == WorkspaceZeroPolicy::OnCreate) { + zero_tensor(block.tensor); + } + + transient_active_.push_back(std::move(block)); + return make_buffer(transient_active_.back(), bytes); +} + +infinicore::Tensor InferenceWorkspaceManager::acquire_temp_tensor( + std::string_view tag, + const infinicore::Shape &shape, + infinicore::DataType dtype, + size_t /*alignment*/, + WorkspaceZeroPolicy zero_policy) { + if (!in_forward_) { + throw std::runtime_error("transient inference tensor can only be acquired during forward: tag=" + std::string(tag) + ", scope=" + collective_scope_ + ", graph_recording=" + (infinicore::context::isGraphRecording() ? std::string("true") : std::string("false"))); + } + + const auto bytes = numel_of(shape) * infinicore::dsize(dtype); + const auto key = typed_key(shape, dtype); + auto &free_list = typed_transient_free_[key]; + + if (mode_ == WorkspaceMode::Warmup) { + transient_plan_.push_back(TransientRequest{std::string(tag), bytes, bytes}); + } else if (mode_ == WorkspaceMode::Frozen) { + if (transient_cursor_ >= transient_plan_.size()) { + throw std::runtime_error("frozen workspace pool received more transient tensor requests than warmup recorded"); + } + const auto &expected = transient_plan_[transient_cursor_]; + if (expected.capacity != bytes) { + throw std::runtime_error( + "frozen workspace pool tensor size mismatch for " + std::string(tag) + + ": expected " + std::to_string(expected.capacity) + + ", got " + std::to_string(bytes)); + } + } + ++transient_cursor_; + + Block block; + if (!free_list.empty()) { + block = std::move(free_list.back()); + free_list.pop_back(); + } else { + if (mode_ == WorkspaceMode::Frozen) { + throw std::runtime_error("frozen workspace pool has no free typed tensor for " + std::string(tag)); + } + block.tensor = infinicore::Tensor::empty(shape, dtype, device_); + block.capacity = bytes; + block.lifetime = WorkspaceLifetime::Transient; + block.typed = true; + } + + block.owner = "transient"; + block.key = key; + block.active = true; + if (zero_policy == WorkspaceZeroPolicy::OnAcquire || zero_policy == WorkspaceZeroPolicy::OnCreate) { + zero_tensor(block.tensor); + } + + transient_active_.push_back(std::move(block)); + return transient_active_.back().tensor; +} + +void InferenceWorkspaceManager::set_collective_scope(std::string_view scope) { + collective_scope_ = std::string(scope); +} + +void InferenceWorkspaceManager::clear_collective_scope() { + collective_scope_.clear(); +} + +std::string InferenceWorkspaceManager::next_collective_key(std::string_view tag) { + if (!in_forward_) { + throw std::runtime_error("collective workspace keys can only be acquired during forward"); + } + std::string key; + if (!collective_scope_.empty()) { + key += collective_scope_; + key += ":"; + } + key += std::string(tag); + key += "."; + key += std::to_string(collective_cursor_++); + return key; +} + +WorkspaceBuffer InferenceWorkspaceManager::get_persistent_buffer( + std::string_view owner, + std::string_view key, + size_t bytes, + WorkspaceLifetime lifetime, + WorkspaceZeroPolicy zero_policy) { + if (lifetime == WorkspaceLifetime::Transient) { + throw std::runtime_error("persistent workspace request cannot use transient lifetime"); + } + + const auto map_key = full_key(owner, key); + const auto capacity = normalize_bytes(bytes, 256); + auto it = persistent_.find(map_key); + if (it == persistent_.end()) { + Block block; + block.tensor = infinicore::Tensor::empty({capacity}, infinicore::DataType::U8, device_); + block.capacity = capacity; + block.owner = std::string(owner); + block.key = map_key; + block.lifetime = lifetime; + if (zero_policy == WorkspaceZeroPolicy::OnCreate || zero_policy == WorkspaceZeroPolicy::OnAcquire) { + zero_tensor(block.tensor); + } + it = persistent_.emplace(map_key, std::move(block)).first; + } else { + auto &block = it->second; + if (block.owner != std::string(owner)) { + throw std::runtime_error("persistent workspace owner mismatch for key " + map_key); + } + if (block.lifetime != lifetime) { + throw std::runtime_error("persistent workspace lifetime mismatch for key " + map_key); + } + if (block.capacity < capacity) { + if (graph_frozen_) { + throw std::runtime_error("cannot grow workspace after graph freeze: " + map_key); + } + block.tensor = infinicore::Tensor::empty({capacity}, infinicore::DataType::U8, device_); + block.capacity = capacity; + if (zero_policy == WorkspaceZeroPolicy::OnCreate || zero_policy == WorkspaceZeroPolicy::OnAcquire) { + zero_tensor(block.tensor); + } + } else if (zero_policy == WorkspaceZeroPolicy::OnAcquire) { + zero_tensor(block.tensor); + } + } + return make_buffer(it->second, bytes); +} + +infinicore::Tensor InferenceWorkspaceManager::get_persistent_tensor( + std::string_view owner, + std::string_view key, + const infinicore::Shape &shape, + infinicore::DataType dtype, + WorkspaceLifetime lifetime, + WorkspaceZeroPolicy zero_policy) { + const auto bytes = numel_of(shape) * infinicore::dsize(dtype); + if (dtype == infinicore::DataType::U8 && shape.size() == 1) { + if (lifetime == WorkspaceLifetime::Transient) { + throw std::runtime_error("persistent workspace request cannot use transient lifetime"); + } + const auto map_key = full_key(owner, key); + auto it = persistent_.find(map_key); + if (it == persistent_.end()) { + Block block; + block.tensor = infinicore::Tensor::empty(shape, dtype, device_); + block.capacity = bytes; + block.owner = std::string(owner); + block.key = map_key; + block.lifetime = lifetime; + if (zero_policy == WorkspaceZeroPolicy::OnCreate || zero_policy == WorkspaceZeroPolicy::OnAcquire) { + zero_tensor(block.tensor); + } + it = persistent_.emplace(map_key, std::move(block)).first; + } else { + auto &block = it->second; + if (block.owner != std::string(owner)) { + throw std::runtime_error("persistent workspace owner mismatch for key " + map_key); + } + if (block.lifetime != lifetime) { + throw std::runtime_error("persistent workspace lifetime mismatch for key " + map_key); + } + if (block.capacity < bytes) { + if (graph_frozen_) { + throw std::runtime_error("cannot grow workspace after graph freeze: " + map_key); + } + block.tensor = infinicore::Tensor::empty(shape, dtype, device_); + block.capacity = bytes; + if (zero_policy == WorkspaceZeroPolicy::OnCreate || zero_policy == WorkspaceZeroPolicy::OnAcquire) { + zero_tensor(block.tensor); + } + } else if (zero_policy == WorkspaceZeroPolicy::OnAcquire) { + zero_tensor(block.tensor); + } + } + return it->second.tensor; + } + + return get_persistent_buffer(owner, key, bytes, lifetime, zero_policy).as_tensor(shape, dtype); +} + +void InferenceWorkspaceManager::reserve_persistent_buffer( + std::string_view owner, + std::string_view key, + size_t bytes, + WorkspaceLifetime lifetime, + WorkspaceZeroPolicy zero_policy) { + (void)get_persistent_buffer(owner, key, bytes, lifetime, zero_policy); +} + +void InferenceWorkspaceManager::freeze_for_graph() { + graph_frozen_ = true; +} + +void InferenceWorkspaceManager::unfreeze_graph() { + graph_frozen_ = false; +} + +void InferenceWorkspaceManager::zero(std::string_view key) { + auto it = persistent_.find(std::string(key)); + if (it != persistent_.end()) { + zero_tensor(it->second.tensor); + } +} + +void InferenceWorkspaceManager::zero_by_prefix(std::string_view prefix) { + for (auto &[key, block] : persistent_) { + if (key.compare(0, prefix.size(), prefix.data(), prefix.size()) == 0) { + zero_tensor(block.tensor); + } + } +} + +WorkspaceStats InferenceWorkspaceManager::stats() const { + WorkspaceStats stats; + for (const auto &[capacity, blocks] : transient_free_) { + stats.transient_reserved_bytes += capacity * blocks.size(); + stats.transient_cached_blocks += blocks.size(); + } + for (const auto &[_, blocks] : typed_transient_free_) { + for (const auto &block : blocks) { + stats.transient_reserved_bytes += block.capacity; + ++stats.transient_cached_blocks; + } + } + for (const auto &block : transient_retired_) { + stats.transient_reserved_bytes += block.capacity; + } + for (const auto &block : transient_active_) { + stats.transient_reserved_bytes += block.capacity; + stats.transient_active_bytes += block.capacity; + ++stats.transient_active_blocks; + } + for (const auto &[_, block] : persistent_) { + stats.persistent_reserved_bytes += block.capacity; + ++stats.persistent_blocks; + } + return stats; +} + +size_t InferenceWorkspaceManager::normalize_bytes(size_t bytes, size_t alignment) const { + if (bytes == 0) { + return 0; + } + if (alignment == 0) { + alignment = 1; + } + return ((bytes + alignment - 1) / alignment) * alignment; +} + +WorkspaceBuffer InferenceWorkspaceManager::make_buffer(Block &block, size_t requested_bytes) { + return WorkspaceBuffer{block.tensor->data(), requested_bytes, device_, block.tensor}; +} + +void InferenceWorkspaceManager::zero_tensor(infinicore::Tensor &tensor) const { + if (infinicore::context::isGraphRecording()) { + infinicore::context::addGraphOperator(std::make_shared(tensor)); + } else { + set_zeros_device_async(tensor); + } +} + +std::string InferenceWorkspaceManager::typed_key(const infinicore::Shape &shape, infinicore::DataType dtype) const { + std::string key = std::to_string(static_cast(dtype)); + for (auto dim : shape) { + key += ":" + std::to_string(dim); + } + return key; +} + +void InferenceWorkspaceManager::reclaim_retired_typed_blocks() { + std::vector still_retired; + for (auto &block : transient_retired_) { + if (block.tensor.use_count() == 1) { + typed_transient_free_[block.key].push_back(std::move(block)); + } else { + still_retired.push_back(std::move(block)); + } + } + transient_retired_ = std::move(still_retired); +} + +std::string InferenceWorkspaceManager::full_key(std::string_view owner, std::string_view key) const { + return std::string(owner) + "." + std::string(key); +} + +} // namespace infinilm::engine diff --git a/csrc/engine/workspace/inference_workspace_manager.hpp b/csrc/engine/workspace/inference_workspace_manager.hpp new file mode 100644 index 000000000..cab3f1fbf --- /dev/null +++ b/csrc/engine/workspace/inference_workspace_manager.hpp @@ -0,0 +1,155 @@ +#pragma once + +#include "infinicore/device.hpp" +#include "infinicore/dtype.hpp" +#include "infinicore/tensor.hpp" + +#include +#include +#include +#include +#include + +namespace infinilm::engine { + +enum class WorkspaceLifetime { + Transient, + Persistent, + GraphPersistent, +}; + +enum class WorkspaceZeroPolicy { + None, + OnCreate, + OnAcquire, +}; + +enum class WorkspaceMode { + Dynamic, + Warmup, + Frozen, +}; + +struct WorkspaceStats { + size_t transient_reserved_bytes = 0; + size_t persistent_reserved_bytes = 0; + size_t transient_active_bytes = 0; + size_t transient_active_blocks = 0; + size_t transient_cached_blocks = 0; + size_t persistent_blocks = 0; +}; + +struct WorkspaceBuffer { + void *ptr = nullptr; + size_t bytes = 0; + infinicore::Device device; + infinicore::Tensor owner; + + template + T *as() const { + return reinterpret_cast(ptr); + } + + infinicore::Tensor as_tensor(const infinicore::Shape &shape, infinicore::DataType dtype) const; +}; + +class InferenceWorkspaceManager { +public: + explicit InferenceWorkspaceManager(infinicore::Device device); + + void begin_forward(); + void end_forward(); + + void begin_warmup(); + void end_warmup(); + void freeze_transient_pool(); + void unfreeze_transient_pool(); + + WorkspaceBuffer acquire_temp_buffer( + std::string_view tag, + size_t bytes, + size_t alignment = 256, + WorkspaceZeroPolicy zero_policy = WorkspaceZeroPolicy::None); + + infinicore::Tensor acquire_temp_tensor( + std::string_view tag, + const infinicore::Shape &shape, + infinicore::DataType dtype, + size_t alignment = 256, + WorkspaceZeroPolicy zero_policy = WorkspaceZeroPolicy::None); + + void set_collective_scope(std::string_view scope); + void clear_collective_scope(); + std::string next_collective_key(std::string_view tag); + + WorkspaceBuffer get_persistent_buffer( + std::string_view owner, + std::string_view key, + size_t bytes, + WorkspaceLifetime lifetime = WorkspaceLifetime::Persistent, + WorkspaceZeroPolicy zero_policy = WorkspaceZeroPolicy::None); + + infinicore::Tensor get_persistent_tensor( + std::string_view owner, + std::string_view key, + const infinicore::Shape &shape, + infinicore::DataType dtype, + WorkspaceLifetime lifetime = WorkspaceLifetime::Persistent, + WorkspaceZeroPolicy zero_policy = WorkspaceZeroPolicy::None); + + void reserve_persistent_buffer( + std::string_view owner, + std::string_view key, + size_t bytes, + WorkspaceLifetime lifetime = WorkspaceLifetime::Persistent, + WorkspaceZeroPolicy zero_policy = WorkspaceZeroPolicy::None); + + void freeze_for_graph(); + void unfreeze_graph(); + + void zero(std::string_view key); + void zero_by_prefix(std::string_view prefix); + + WorkspaceStats stats() const; + +private: + struct Block { + infinicore::Tensor tensor; + size_t capacity = 0; + std::string owner; + std::string key; + WorkspaceLifetime lifetime = WorkspaceLifetime::Transient; + bool active = false; + bool typed = false; + }; + + struct TransientRequest { + std::string tag; + size_t bytes = 0; + size_t capacity = 0; + }; + + size_t normalize_bytes(size_t bytes, size_t alignment) const; + WorkspaceBuffer make_buffer(Block &block, size_t requested_bytes); + std::string typed_key(const infinicore::Shape &shape, infinicore::DataType dtype) const; + void reclaim_retired_typed_blocks(); + void zero_tensor(infinicore::Tensor &tensor) const; + std::string full_key(std::string_view owner, std::string_view key) const; + + infinicore::Device device_; + bool in_forward_ = false; + bool graph_frozen_ = false; + WorkspaceMode mode_ = WorkspaceMode::Dynamic; + size_t transient_cursor_ = 0; + size_t collective_cursor_ = 0; + std::string collective_scope_; + + std::unordered_map> transient_free_; + std::unordered_map> typed_transient_free_; + std::vector transient_active_; + std::vector transient_retired_; + std::unordered_map persistent_; + std::vector transient_plan_; +}; + +} // namespace infinilm::engine diff --git a/csrc/engine/workspace/tensor_allocator.hpp b/csrc/engine/workspace/tensor_allocator.hpp new file mode 100644 index 000000000..b2f78f387 --- /dev/null +++ b/csrc/engine/workspace/tensor_allocator.hpp @@ -0,0 +1,35 @@ +#pragma once + +#include "workspace_context.hpp" + +#include "infinicore/context/context.hpp" +#include "infinicore/tensor.hpp" + +#include + +namespace infinilm::engine { + +inline infinicore::Tensor allocate_inference_tensor( + std::string_view tag, + const infinicore::Shape &shape, + infinicore::DataType dtype, + const infinicore::Device &device, + WorkspaceZeroPolicy zero_policy = WorkspaceZeroPolicy::None) { + if (infinicore::context::isGraphRecording()) { + if (zero_policy == WorkspaceZeroPolicy::OnAcquire || zero_policy == WorkspaceZeroPolicy::OnCreate) { + return infinicore::Tensor::zeros(shape, dtype, device); + } + return infinicore::Tensor::empty(shape, dtype, device); + } + + (void)tag; + // Do not pool returned activation tensors here. InfiniCore view/as_strided creates + // a new TensorImpl that shares storage, but InfiniLM cannot observe that storage + // lifetime yet, so reusing these buffers can race with surviving views. + if (zero_policy == WorkspaceZeroPolicy::OnAcquire || zero_policy == WorkspaceZeroPolicy::OnCreate) { + return infinicore::Tensor::zeros(shape, dtype, device); + } + return infinicore::Tensor::empty(shape, dtype, device); +} + +} // namespace infinilm::engine diff --git a/csrc/engine/workspace/workspace_context.cpp b/csrc/engine/workspace/workspace_context.cpp new file mode 100644 index 000000000..78faade85 --- /dev/null +++ b/csrc/engine/workspace/workspace_context.cpp @@ -0,0 +1,57 @@ +#include "workspace_context.hpp" + +#include + +namespace infinilm::engine { + +namespace { +thread_local InferenceWorkspaceManager *current_manager = nullptr; +} + +WorkspaceContextGuard::WorkspaceContextGuard(InferenceWorkspaceManager *manager) + : previous_(current_manager) { + current_manager = manager; +} + +WorkspaceContextGuard::~WorkspaceContextGuard() { + current_manager = previous_; +} + +WorkspaceForwardGuard::WorkspaceForwardGuard(InferenceWorkspaceManager *manager) + : manager_(manager), active_(manager != nullptr) { + if (active_) { + manager_->begin_forward(); + } +} + +WorkspaceForwardGuard::~WorkspaceForwardGuard() { + if (active_) { + manager_->end_forward(); + } +} + +WorkspaceCollectiveScopeGuard::WorkspaceCollectiveScopeGuard(InferenceWorkspaceManager *manager, std::string_view scope) + : manager_(manager), active_(manager != nullptr) { + if (active_) { + manager_->set_collective_scope(scope); + } +} + +WorkspaceCollectiveScopeGuard::~WorkspaceCollectiveScopeGuard() { + if (active_) { + manager_->clear_collective_scope(); + } +} + +InferenceWorkspaceManager *maybe_current_workspace() { + return current_manager; +} + +InferenceWorkspaceManager ¤t_workspace() { + if (current_manager == nullptr) { + throw std::runtime_error("no active inference workspace context"); + } + return *current_manager; +} + +} // namespace infinilm::engine diff --git a/csrc/engine/workspace/workspace_context.hpp b/csrc/engine/workspace/workspace_context.hpp new file mode 100644 index 000000000..fa08e31e1 --- /dev/null +++ b/csrc/engine/workspace/workspace_context.hpp @@ -0,0 +1,50 @@ +#pragma once + +#include "inference_workspace_manager.hpp" + +#include + +namespace infinilm::engine { + +class WorkspaceContextGuard { +public: + explicit WorkspaceContextGuard(InferenceWorkspaceManager *manager); + ~WorkspaceContextGuard(); + + WorkspaceContextGuard(const WorkspaceContextGuard &) = delete; + WorkspaceContextGuard &operator=(const WorkspaceContextGuard &) = delete; + +private: + InferenceWorkspaceManager *previous_ = nullptr; +}; + +class WorkspaceForwardGuard { +public: + explicit WorkspaceForwardGuard(InferenceWorkspaceManager *manager); + ~WorkspaceForwardGuard(); + + WorkspaceForwardGuard(const WorkspaceForwardGuard &) = delete; + WorkspaceForwardGuard &operator=(const WorkspaceForwardGuard &) = delete; + +private: + InferenceWorkspaceManager *manager_ = nullptr; + bool active_ = false; +}; + +class WorkspaceCollectiveScopeGuard { +public: + WorkspaceCollectiveScopeGuard(InferenceWorkspaceManager *manager, std::string_view scope); + ~WorkspaceCollectiveScopeGuard(); + + WorkspaceCollectiveScopeGuard(const WorkspaceCollectiveScopeGuard &) = delete; + WorkspaceCollectiveScopeGuard &operator=(const WorkspaceCollectiveScopeGuard &) = delete; + +private: + InferenceWorkspaceManager *manager_ = nullptr; + bool active_ = false; +}; + +InferenceWorkspaceManager *maybe_current_workspace(); +InferenceWorkspaceManager ¤t_workspace(); + +} // namespace infinilm::engine diff --git a/csrc/layers/linear/base_linear.cpp b/csrc/layers/linear/base_linear.cpp index f7d1c80de..ce2eacb47 100644 --- a/csrc/layers/linear/base_linear.cpp +++ b/csrc/layers/linear/base_linear.cpp @@ -32,7 +32,7 @@ BaseLinear::BaseLinear(size_t in_features, size_t out_features, } } -infinicore::Tensor BaseLinear::compute_linear(infinicore::Tensor &input) const { +infinicore::Tensor BaseLinear::compute_linear(infinicore::Tensor &input, const infinicore::Tensor *output) const { // Build params map from direct parameters only (not state_dict which uses a // static local and is not thread-safe across RankWorker threads). infinilm::quantization::ParamsMap params; @@ -40,7 +40,7 @@ infinicore::Tensor BaseLinear::compute_linear(infinicore::Tensor &input) const { params[name] = static_cast(param); } - return quantization_->forward(params, input, has_bias_, alpha_); + return quantization_->forward(params, input, has_bias_, alpha_, output); } infinicore::Tensor BaseLinear::forward(infinicore::Tensor &input) const { diff --git a/csrc/layers/linear/base_linear.hpp b/csrc/layers/linear/base_linear.hpp index a36954836..81244d17f 100644 --- a/csrc/layers/linear/base_linear.hpp +++ b/csrc/layers/linear/base_linear.hpp @@ -57,7 +57,7 @@ class BaseLinear : public infinicore::nn::Module { const infinicore::nn::Parameter &get_parameter_ref(const std::string &name) const; protected: - infinicore::Tensor compute_linear(infinicore::Tensor &input) const; + infinicore::Tensor compute_linear(infinicore::Tensor &input, const infinicore::Tensor *output = nullptr) const; size_t in_features_; size_t out_features_; diff --git a/csrc/layers/linear/fused_linear.cpp b/csrc/layers/linear/fused_linear.cpp index 1bcb96c94..07dc21874 100644 --- a/csrc/layers/linear/fused_linear.cpp +++ b/csrc/layers/linear/fused_linear.cpp @@ -151,9 +151,10 @@ GateUpParallelLinear::GateUpParallelLinear(size_t hidden_size, size_t intermedia std::tuple GateUpParallelLinear::forward_split(infinicore::Tensor &input) { auto output = this->forward(input); - auto cols = output->shape()[2]; - auto gate_output = output->narrow({{2, 0, cols / 2}}); - auto up_output = output->narrow({{2, cols / 2, cols / 2}}); + 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}}); return std::make_tuple(gate_output, up_output); } diff --git a/csrc/layers/linear/linear.cpp b/csrc/layers/linear/linear.cpp index 84982409f..d18a28e8d 100644 --- a/csrc/layers/linear/linear.cpp +++ b/csrc/layers/linear/linear.cpp @@ -1,9 +1,54 @@ #include "linear.hpp" +#include "../../engine/distributed/async_collective.hpp" +#include "../../engine/workspace/workspace_context.hpp" #include "infinicore/ops.hpp" #include "infinicore/ops/distributed/allreduce.hpp" +#include +#include +#include +#include #include +#include +#include namespace infinilm::nn { +namespace { + +std::optional parse_chunk_rows_env() { + const char *raw = std::getenv("INFINILM_COMM_OVERLAP_CHUNK_ROWS"); + if (raw == nullptr || raw[0] == '\0') { + return std::nullopt; + } + char *end = nullptr; + const auto value = std::strtoull(raw, &end, 10); + if (end == raw) { + return std::nullopt; + } + return static_cast(value); +} + +size_t default_comm_overlap_chunk_rows(infinicore::Size tp_size) { + if (tp_size > 2 && infinilm::engine::distributed::async_collectives_force_enabled_by_env()) { + return 2048; + } + return 8192; +} + +size_t comm_overlap_chunk_rows(size_t rows, infinicore::Size tp_size) { + constexpr size_t kMaxPendingChunks = 32; + const size_t requested = parse_chunk_rows_env().value_or(default_comm_overlap_chunk_rows(tp_size)); + if (requested == 0) { + return 0; + } + const size_t min_chunk = (rows + kMaxPendingChunks - 1) / kMaxPendingChunks; + return std::max(requested, min_chunk); +} + +bool can_use_async_comm_overlap(const infinicore::Tensor &input, infinicore::Size tp_size) { + return infinilm::engine::distributed::async_collectives_supported_by_policy(input->device(), static_cast(tp_size)); +} + +} // namespace // ---- Linear ---- @@ -84,13 +129,125 @@ RowParallelLinear::RowParallelLinear(size_t in_features, size_t out_features, tp_size_(tp_size), communicator_(communicator) { } +bool RowParallelLinearOutput::has_pending() const { + return async_context != nullptr && !pending_collectives.empty(); +} + +void RowParallelLinearOutput::wait() const { + if (async_context == nullptr) { + return; + } + for (const auto &pending : pending_collectives) { + async_context->wait(pending); + } +} + infinicore::Tensor RowParallelLinear::forward(infinicore::Tensor &input) const { - auto output = BaseLinear::forward(input); + auto result = forward_async(input); + result.wait(); + return result.output; +} + +RowParallelLinearOutput RowParallelLinear::forward_async(infinicore::Tensor &input) const { + const bool needs_allreduce = (tp_size_ > 1) && (communicator_ != nullptr); + auto *workspace = infinilm::engine::maybe_current_workspace(); + auto *async_context = infinilm::engine::distributed::maybe_current_async_collective_context(); + const bool can_async_allreduce = needs_allreduce && async_context != nullptr && async_context->enabled() && !infinicore::context::isGraphRecording() && can_use_async_comm_overlap(input, tp_size_); + const bool supports_preallocated_output = quantization_->supports_preallocated_output(); + + auto output_shape = input->shape(); + output_shape.back() = out_features_; + const auto input_shape = input->shape(); + const size_t input_features = input_shape.back(); + const size_t rows = input->numel() / input_features; + const size_t chunk_rows = comm_overlap_chunk_rows(rows, tp_size_); + const bool use_chunked_overlap = can_async_allreduce && supports_preallocated_output && chunk_rows > 0 && rows > chunk_rows; + + if (use_chunked_overlap) { + infinicore::Tensor partial_output; + infinicore::Tensor output; + if (workspace != nullptr) { + partial_output = workspace->acquire_temp_tensor( + "row_parallel.chunked_allreduce_input", + output_shape, + input->dtype(), + 256); + output = workspace->acquire_temp_tensor( + "row_parallel.chunked_allreduce_output", + output_shape, + input->dtype(), + 256); + } else { + partial_output = infinicore::Tensor::empty(output_shape, input->dtype(), input->device()); + output = infinicore::Tensor::empty(output_shape, input->dtype(), input->device()); + } + + auto input_flat = input->view({rows, input_features}); + auto partial_flat = partial_output->view({rows, out_features_}); + auto output_flat = output->view({rows, out_features_}); + + std::vector pending_collectives; + pending_collectives.reserve((rows + chunk_rows - 1) / chunk_rows); + for (size_t begin = 0; begin < rows; begin += chunk_rows) { + const size_t len = std::min(chunk_rows, rows - begin); + auto input_chunk = input_flat->narrow({{0, begin, len}}); + auto partial_chunk = partial_flat->narrow({{0, begin, len}}); + auto output_chunk = output_flat->narrow({{0, begin, len}}); + compute_linear(input_chunk, &partial_chunk); + pending_collectives.push_back(async_context->allreduce( + output_chunk, + partial_chunk, + INFINICCL_SUM, + communicator_)); + } + return RowParallelLinearOutput{output, partial_output, async_context, std::move(pending_collectives)}; + } + + const bool use_workspace_output = needs_allreduce && workspace != nullptr && supports_preallocated_output; + + infinicore::Tensor output; + infinicore::Tensor allreduce_input; + std::optional preallocated_partial_output; + std::optional preallocated_reduced_output; + if (use_workspace_output) { + preallocated_partial_output.emplace(workspace->acquire_temp_tensor( + "row_parallel.allreduce_input", + output_shape, + input->dtype(), + 256)); + allreduce_input = compute_linear(input, &preallocated_partial_output.value()); + preallocated_reduced_output.emplace(workspace->acquire_temp_tensor( + "row_parallel.allreduce_output", + output_shape, + input->dtype(), + 256)); + output = preallocated_reduced_output.value(); + + if (infinicore::context::isGraphRecording()) { + const auto key = workspace->next_collective_key("row_parallel.allreduce"); + const auto status = infinicclCommRegisterAllReduceBuffer( + communicator_, + key.c_str(), + const_cast(allreduce_input->data()), + allreduce_input->nbytes()); + if (status != INFINI_STATUS_SUCCESS) { + throw std::runtime_error("failed to register row-parallel allreduce workspace buffer: " + std::to_string(static_cast(status))); + } + } + } else { + output = BaseLinear::forward(input); + allreduce_input = output; + } - if ((tp_size_ > 1) && (communicator_ != nullptr)) { - infinicore::op::distributed::allreduce_(output, output, INFINICCL_SUM, communicator_); + std::vector pending_collectives; + if (needs_allreduce) { + if (can_async_allreduce) { + pending_collectives.push_back(async_context->allreduce(output, allreduce_input, INFINICCL_SUM, communicator_)); + } else { + infinicore::op::distributed::allreduce_(output, allreduce_input, INFINICCL_SUM, communicator_); + } } - return output; + return RowParallelLinearOutput{output, allreduce_input, can_async_allreduce ? async_context : nullptr, std::move(pending_collectives)}; } std::string RowParallelLinear::extra_repr() const { diff --git a/csrc/layers/linear/linear.hpp b/csrc/layers/linear/linear.hpp index 566cee77c..f4e99eb9a 100644 --- a/csrc/layers/linear/linear.hpp +++ b/csrc/layers/linear/linear.hpp @@ -1,11 +1,13 @@ #pragma once +#include "../../engine/distributed/async_collective.hpp" #include "../quantization/quantization.hpp" #include "base_linear.hpp" #include "infinicore/nn/module.hpp" #include "infinicore/ops.hpp" #include #include +#include namespace infinilm::nn { @@ -52,6 +54,16 @@ class ColumnParallelLinear : public BaseLinear { infinicore::Size tp_size_ = 1; }; +struct RowParallelLinearOutput { + infinicore::Tensor output; + infinicore::Tensor allreduce_input; + engine::distributed::AsyncCollectiveContext *async_context = nullptr; + std::vector pending_collectives; + + bool has_pending() const; + void wait() const; +}; + class RowParallelLinear : public BaseLinear { public: // Without quantization (backward compat) @@ -70,6 +82,7 @@ class RowParallelLinear : public BaseLinear { infinicclComm_t communicator = nullptr); infinicore::Tensor forward(infinicore::Tensor &input) const; + RowParallelLinearOutput forward_async(infinicore::Tensor &input) const; std::string extra_repr() const; protected: @@ -87,6 +100,7 @@ namespace infinilm::layers::linear { using ReplicatedLinear = infinilm::nn::Linear; using ColumnParallelLinear = infinilm::nn::ColumnParallelLinear; using RowParallelLinear = infinilm::nn::RowParallelLinear; +using RowParallelLinearOutput = infinilm::nn::RowParallelLinearOutput; using BaseLinear = infinilm::nn::BaseLinear; } // namespace infinilm::layers::linear diff --git a/csrc/layers/mlp/mlp.cpp b/csrc/layers/mlp/mlp.cpp index f7604c505..dd999b6f2 100644 --- a/csrc/layers/mlp/mlp.cpp +++ b/csrc/layers/mlp/mlp.cpp @@ -26,14 +26,19 @@ MLP::MLP(std::shared_ptr model_config, use_bias_, dtype, device, tp_rank, tp_size, rank_info.comm); } -infinicore::Tensor MLP::forward(const infinicore::Tensor &hidden_states) const { +infinilm::layers::linear::RowParallelLinearOutput MLP::forward_async(const infinicore::Tensor &hidden_states) const { // 1. Project to gate and up auto hidden_states_mutable = hidden_states; auto [gate, up] = gate_up_proj_->forward_split(hidden_states_mutable); // 2. Apply SwiGLU: silu(gate) * up auto intermediate = infinicore::op::swiglu(up, gate); - // 3. Project down - auto output = down_proj_->forward(intermediate); - return output; + // 3. Project down and optionally leave row-parallel allreduce pending. + return down_proj_->forward_async(intermediate); +} + +infinicore::Tensor MLP::forward(const infinicore::Tensor &hidden_states) const { + auto output = forward_async(hidden_states); + output.wait(); + return output.output; } } // namespace infinilm::layers::mlp diff --git a/csrc/layers/mlp/mlp.hpp b/csrc/layers/mlp/mlp.hpp index abd81bd88..45b557721 100644 --- a/csrc/layers/mlp/mlp.hpp +++ b/csrc/layers/mlp/mlp.hpp @@ -35,6 +35,7 @@ class MLP : public infinicore::nn::Module { * @return Output tensor of shape [batch, seq_len, hidden_size] */ infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + infinilm::layers::linear::RowParallelLinearOutput forward_async(const infinicore::Tensor &hidden_states) const; void process_weights_after_loading() override { gate_up_proj_->process_weights_after_loading(); diff --git a/csrc/layers/moe/router/topk_router.cpp b/csrc/layers/moe/router/topk_router.cpp index e70e55cf8..5cac964ec 100644 --- a/csrc/layers/moe/router/topk_router.cpp +++ b/csrc/layers/moe/router/topk_router.cpp @@ -1,5 +1,6 @@ #include "topk_router.hpp" +#include "../../../utils.hpp" #include "infinicore/ops.hpp" #include @@ -12,6 +13,9 @@ 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; } @@ -53,9 +57,13 @@ TopKRouter::TopKRouter(std::shared_ptr model_conf use_correction_bias_ = model_config->get_or({"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("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("hidden_size")}, model_config->get_dtype(), device)); + ({num_experts_, model_config->get("hidden_size")}, router_dtype, device)); if (use_correction_bias_) { INFINICORE_NN_PARAMETER_INIT(e_score_correction_bias, ({num_experts_}, infinicore::DataType::F32, device)); @@ -79,9 +87,17 @@ TopKRouter::TopKRouter(std::shared_ptr model_conf std::tuple TopKRouter::forward(const infinicore::Tensor &hidden_states) const { ASSERT(hidden_states->ndim() == 2); - size_t ntoken = hidden_states->shape()[0]; - auto router_logits = infinicore::op::linear(hidden_states, weight_, std::nullopt, 1.0f); + auto router_input = hidden_states->is_contiguous() ? hidden_states : hidden_states->contiguous(); + auto router_weight = static_cast(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_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()); @@ -98,6 +114,14 @@ std::tuple 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, diff --git a/csrc/layers/moe/router/topk_router.hpp b/csrc/layers/moe/router/topk_router.hpp index ce1d468b7..27b8c1947 100644 --- a/csrc/layers/moe/router/topk_router.hpp +++ b/csrc/layers/moe/router/topk_router.hpp @@ -11,6 +11,7 @@ namespace infinilm::layers::moe { enum class TopKRouterBackend { Softmax, + LegacySoftmax, Sigmoid, FusedGate, }; diff --git a/csrc/layers/quantization/awq_marlin.cpp b/csrc/layers/quantization/awq_marlin.cpp index 47a4c269e..ca861e518 100644 --- a/csrc/layers/quantization/awq_marlin.cpp +++ b/csrc/layers/quantization/awq_marlin.cpp @@ -45,13 +45,31 @@ infinicore::Tensor AWQMarlin::forward( const ParamsMap ¶ms, const infinicore::Tensor &input, bool has_bias, - float /*alpha*/) const { + float alpha) const { + return forward(params, input, has_bias, alpha, nullptr); +} + +infinicore::Tensor AWQMarlin::forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float /*alpha*/, + const infinicore::Tensor *preallocated_output) const { auto input_contiguous = input->is_contiguous() ? input : input->contiguous(); const auto &shape = input_contiguous->shape(); const size_t k = shape.back(); const size_t m = input_contiguous->numel() / k; auto flat_input = input_contiguous->view({m, k}); - auto output = infinicore::Tensor::empty({m, output_size_per_partition_}, input->dtype(), input->device()); + + auto out_shape = shape; + out_shape.back() = output_size_per_partition_; + auto output = preallocated_output != nullptr + ? (*preallocated_output)->view({m, output_size_per_partition_}) + : infinicore::Tensor::empty({m, output_size_per_partition_}, input->dtype(), input->device()); + if (preallocated_output != nullptr && + ((*preallocated_output)->shape() != out_shape || (*preallocated_output)->dtype() != input->dtype() || (*preallocated_output)->device() != input->device())) { + throw std::runtime_error("preallocated AWQ Marlin output does not match linear output shape, dtype, or device"); + } auto qweight = params.at("qweight"); auto scales = params.at("scales"); @@ -83,8 +101,6 @@ infinicore::Tensor AWQMarlin::forward( infinicore::op::add_(output, output, bias->as_strided(output->shape(), {0, 1})); } - auto out_shape = shape; - out_shape.back() = output_size_per_partition_; return output->view(out_shape); } @@ -117,6 +133,15 @@ infinicore::Tensor AWQMarlin::forward( throw std::runtime_error("AWQ Marlin is not available because InfiniCore was built without Marlin GEMM headers."); } +infinicore::Tensor AWQMarlin::forward( + const ParamsMap &, + const infinicore::Tensor &, + bool, + float, + const infinicore::Tensor *) const { + throw std::runtime_error("AWQ Marlin is not available because InfiniCore was built without Marlin GEMM headers."); +} + std::vector AWQMarlin::split_params( const std::unordered_map &, const std::vector &, diff --git a/csrc/layers/quantization/awq_marlin.hpp b/csrc/layers/quantization/awq_marlin.hpp index a22fb43b3..b31a9aa6e 100644 --- a/csrc/layers/quantization/awq_marlin.hpp +++ b/csrc/layers/quantization/awq_marlin.hpp @@ -22,6 +22,15 @@ class AWQMarlin : public BaseQuantization { bool has_bias, float alpha = 1.0f) const override; + infinicore::Tensor forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float alpha, + const infinicore::Tensor *output) const override; + + bool supports_preallocated_output() const override { return true; } + std::vector split_params( const std::unordered_map ¶ms, const std::vector &splits, @@ -41,15 +50,11 @@ class AWQMarlin : public BaseQuantization { infinicore::Tensor &g_idx, infinicore::Tensor &perm) const; + // Per-layer Marlin workspace, matching vLLM/SGLang's graph-stable layer state. + mutable infinicore::Tensor workspace_; + size_t input_size_per_partition_; size_t output_size_per_partition_; - // Per-layer Marlin workspace. It must be all-zero before each launch - // because the current InfiniCore Marlin kernels use it as lock state. - // TODO: replace per-layer memset with a shared global zero workspace, or - // update the kernels so the lock region is self-reset at completion. The - // remaining gap to vLLM is mainly from two sources: TP communication cost - // and this workspace memset/reset path. - mutable infinicore::Tensor workspace_; }; } // namespace infinilm::quantization diff --git a/csrc/layers/quantization/base_quantization.hpp b/csrc/layers/quantization/base_quantization.hpp index 6c37f4121..fcf7b9155 100644 --- a/csrc/layers/quantization/base_quantization.hpp +++ b/csrc/layers/quantization/base_quantization.hpp @@ -59,6 +59,18 @@ class BaseQuantization : public std::enable_shared_from_this { bool has_bias, float alpha = 1.0f) const = 0; + virtual infinicore::Tensor forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float alpha, + const infinicore::Tensor *output) const { + (void)output; + return forward(params, input, has_bias, alpha); + } + + virtual bool supports_preallocated_output() const { return false; } + // Dimension for fused-split (gate/up, q/k/v) of a column-parallel weight. // For NoneQuantization weight [out, in], split is on dim0. // For AWQ qweight [in, out/pack], split is on dim1. diff --git a/csrc/layers/quantization/gptq_marlin.cpp b/csrc/layers/quantization/gptq_marlin.cpp index e95d85b5a..ceb06e640 100644 --- a/csrc/layers/quantization/gptq_marlin.cpp +++ b/csrc/layers/quantization/gptq_marlin.cpp @@ -45,13 +45,31 @@ infinicore::Tensor GPTQMarlin::forward( const ParamsMap ¶ms, const infinicore::Tensor &input, bool has_bias, - float /*alpha*/) const { + float alpha) const { + return forward(params, input, has_bias, alpha, nullptr); +} + +infinicore::Tensor GPTQMarlin::forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float /*alpha*/, + const infinicore::Tensor *preallocated_output) const { auto input_contiguous = input->is_contiguous() ? input : input->contiguous(); const auto &shape = input_contiguous->shape(); const size_t k = shape.back(); const size_t m = input_contiguous->numel() / k; auto flat_input = input_contiguous->view({m, k}); - auto output = infinicore::Tensor::empty({m, output_size_per_partition_}, input->dtype(), input->device()); + + auto out_shape = shape; + out_shape.back() = output_size_per_partition_; + auto output = preallocated_output != nullptr + ? (*preallocated_output)->view({m, output_size_per_partition_}) + : infinicore::Tensor::empty({m, output_size_per_partition_}, input->dtype(), input->device()); + if (preallocated_output != nullptr && + ((*preallocated_output)->shape() != out_shape || (*preallocated_output)->dtype() != input->dtype() || (*preallocated_output)->device() != input->device())) { + throw std::runtime_error("preallocated GPTQ Marlin output does not match linear output shape, dtype, or device"); + } auto qweight = params.at("qweight"); auto scales = params.at("scales"); @@ -83,8 +101,6 @@ infinicore::Tensor GPTQMarlin::forward( infinicore::op::add_(output, output, bias->as_strided(output->shape(), {0, 1})); } - auto out_shape = shape; - out_shape.back() = output_size_per_partition_; return output->view(out_shape); } @@ -117,6 +133,15 @@ infinicore::Tensor GPTQMarlin::forward( throw std::runtime_error("GPTQ Marlin is not available because InfiniCore was built without Marlin GEMM headers."); } +infinicore::Tensor GPTQMarlin::forward( + const ParamsMap &, + const infinicore::Tensor &, + bool, + float, + const infinicore::Tensor *) const { + throw std::runtime_error("GPTQ Marlin is not available because InfiniCore was built without Marlin GEMM headers."); +} + std::vector GPTQMarlin::split_params( const std::unordered_map &, const std::vector &, diff --git a/csrc/layers/quantization/gptq_marlin.hpp b/csrc/layers/quantization/gptq_marlin.hpp index c16c79438..10509892d 100644 --- a/csrc/layers/quantization/gptq_marlin.hpp +++ b/csrc/layers/quantization/gptq_marlin.hpp @@ -24,6 +24,15 @@ class GPTQMarlin : public BaseQuantization { bool has_bias, float alpha = 1.0f) const override; + infinicore::Tensor forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float alpha, + const infinicore::Tensor *output) const override; + + bool supports_preallocated_output() const override { return true; } + std::vector split_params( const std::unordered_map ¶ms, const std::vector &splits, @@ -43,16 +52,12 @@ class GPTQMarlin : public BaseQuantization { infinicore::Tensor &g_idx, infinicore::Tensor &perm) const; + // Per-layer Marlin workspace, matching vLLM/SGLang's graph-stable layer state. + mutable infinicore::Tensor workspace_; + size_t input_size_per_partition_; size_t output_size_per_partition_; bool is_k_full_; - // Per-layer Marlin workspace. It must be all-zero before each launch - // because the current InfiniCore Marlin kernels use it as lock state. - // TODO: replace per-layer memset with a shared global zero workspace, or - // update the kernels so the lock region is self-reset at completion. The - // remaining gap to vLLM is mainly from two sources: TP communication cost - // and this workspace memset/reset path. - mutable infinicore::Tensor workspace_; }; } // namespace infinilm::quantization diff --git a/csrc/layers/quantization/gptq_qy.cpp b/csrc/layers/quantization/gptq_qy.cpp index 4098e452d..ee3221d0d 100644 --- a/csrc/layers/quantization/gptq_qy.cpp +++ b/csrc/layers/quantization/gptq_qy.cpp @@ -1,7 +1,10 @@ #include "gptq_qy.hpp" +#include "../../engine/workspace/tensor_allocator.hpp" +#include "../../utils.hpp" #include "infinicore/ops.hpp" #include "infinicore/ops/linear_w4a16_gptq_qy.hpp" #include +#include namespace infinilm::quantization { @@ -33,13 +36,41 @@ infinicore::Tensor GPTQ_QY::forward( const ParamsMap ¶ms, const infinicore::Tensor &input, bool has_bias, - float /*alpha*/) const { + float alpha) const { + return forward(params, input, has_bias, alpha, nullptr); +} + +infinicore::Tensor GPTQ_QY::forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float /*alpha*/, + const infinicore::Tensor *preallocated_output) const { auto input_contiguous = input->is_contiguous() ? input : input->contiguous(); auto qweight = params.at("qweight"); auto qzeros = params.at("qzeros"); auto scales = params.at("scales"); - auto output = infinicore::op::linear_w4a16_gptq_qy(input_contiguous->contiguous(), qweight, qzeros, scales, 0, 4); + auto input_for_linear = input_contiguous->contiguous(); + auto output_shape = input_for_linear->shape(); + output_shape.back() = qweight->shape()[1]; + + infinicore::Tensor output; + if (preallocated_output != nullptr) { + output = *preallocated_output; + if (output->shape() != output_shape || output->dtype() != input_for_linear->dtype() || output->device() != input_for_linear->device()) { + throw std::runtime_error("preallocated GPTQ_QY output does not match linear output shape, dtype, or device"); + } + set_zeros_device_async(output); + } else { + output = infinilm::engine::allocate_inference_tensor( + "linear.gptq_qy.output", + output_shape, + input_for_linear->dtype(), + input_for_linear->device(), + infinilm::engine::WorkspaceZeroPolicy::OnAcquire); + } + infinicore::op::linear_w4a16_gptq_qy_(output, input_for_linear, qweight, scales, qzeros, 0, 4); if (has_bias) { auto bias = params.at("bias"); diff --git a/csrc/layers/quantization/gptq_qy.hpp b/csrc/layers/quantization/gptq_qy.hpp index 634b4aaf7..034fa1cf5 100644 --- a/csrc/layers/quantization/gptq_qy.hpp +++ b/csrc/layers/quantization/gptq_qy.hpp @@ -112,7 +112,16 @@ class GPTQ_QY : public BaseQuantization { bool has_bias, float alpha = 1.0f) const override; + infinicore::Tensor forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float alpha, + const infinicore::Tensor *output) const override; + // Split fused linear parameters into named sub-parameters + bool supports_preallocated_output() const override { return true; } + std::vector split_params( const std::unordered_map ¶ms, const std::vector &splits, diff --git a/csrc/layers/quantization/none_quantization.cpp b/csrc/layers/quantization/none_quantization.cpp index e1f67a7d1..2a98c0afe 100644 --- a/csrc/layers/quantization/none_quantization.cpp +++ b/csrc/layers/quantization/none_quantization.cpp @@ -1,6 +1,8 @@ #include "none_quantization.hpp" +#include "../../engine/workspace/tensor_allocator.hpp" #include "infinicore/ops/linear.hpp" #include +#include namespace infinilm::quantization { @@ -26,6 +28,15 @@ infinicore::Tensor NoneQuantization::forward( const infinicore::Tensor &input, bool has_bias, float alpha) const { + return forward(params, input, has_bias, alpha, nullptr); +} + +infinicore::Tensor NoneQuantization::forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float alpha, + const infinicore::Tensor *preallocated_output) const { auto input_contiguous = input->is_contiguous() ? input : input->contiguous(); auto weight = params.at("weight"); @@ -35,7 +46,26 @@ infinicore::Tensor NoneQuantization::forward( bias_opt = params.at("bias"); } - return infinicore::op::linear(input_contiguous->contiguous(), weight->contiguous(), bias_opt, alpha); + auto input_for_linear = input_contiguous->contiguous(); + auto weight_for_linear = weight->contiguous(); + auto output_shape = input_for_linear->shape(); + output_shape.back() = weight_for_linear->shape()[0]; + + infinicore::Tensor output; + if (preallocated_output != nullptr) { + output = *preallocated_output; + if (output->shape() != output_shape || output->dtype() != input_for_linear->dtype() || output->device() != input_for_linear->device()) { + throw std::runtime_error("preallocated NoneQuantization output does not match linear output shape, dtype, or device"); + } + } else { + output = infinilm::engine::allocate_inference_tensor( + "linear.none.output", + output_shape, + input_for_linear->dtype(), + input_for_linear->device()); + } + infinicore::op::linear_(output, input_for_linear, weight_for_linear, bias_opt, alpha); + return output; } std::vector NoneQuantization::split_params( diff --git a/csrc/layers/quantization/none_quantization.hpp b/csrc/layers/quantization/none_quantization.hpp index 18fe4bf17..fb9bfed0f 100644 --- a/csrc/layers/quantization/none_quantization.hpp +++ b/csrc/layers/quantization/none_quantization.hpp @@ -27,6 +27,15 @@ class NoneQuantization : public BaseQuantization { bool has_bias, float alpha = 1.0f) const override; + infinicore::Tensor forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float alpha, + const infinicore::Tensor *output) const override; + + bool supports_preallocated_output() const override { return true; } + std::vector split_params( const std::unordered_map ¶ms, const std::vector &splits, diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_attention.cpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_attention.cpp new file mode 100644 index 000000000..44941deb2 --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_attention.cpp @@ -0,0 +1,163 @@ +#include "ernie4_5_attention.hpp" + +#include "../../global_state/global_state.hpp" +#include "../../layers/attention/attention.hpp" +#include "../../utils.hpp" +#include "infinicore/ops.hpp" + +#include + +namespace infinilm::models::ernie4_5_moe_vl { + +Ernie4_5Attention::Ernie4_5Attention(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) { + layer_idx_ = layer_idx; + hidden_size_ = model_config->get("hidden_size"); + head_dim_ = model_config->get("head_dim"); + + const auto &dtype = model_config->get_dtype(); + const size_t total_num_heads = model_config->get("num_attention_heads"); + const size_t total_num_kv_heads = model_config->get("num_key_value_heads"); + const bool use_bias = model_config->get_or("attention_bias", false); + const bool use_output_bias = model_config->get_or("attention_output_bias", false); + const auto &config_json = model_config->get_config_json(); + + attention_backend_ = infinilm::global_state::get_infinilm_config().attention_backend; + const engine::distributed::RankInfo &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); + const int tp_rank = infinilm::global_state::get_tensor_model_parallel_rank(); + const int tp_size = infinilm::global_state::get_tensor_model_parallel_world_size(); + + num_attention_heads_ = total_num_heads / tp_size; + num_key_value_heads_ = total_num_kv_heads < static_cast(tp_size) + ? 1 + : total_num_kv_heads / tp_size; + + auto quantization_method = model_config->get_quantization_method(); + auto register_fn = [this](const std::string &name, infinicore::nn::Parameter parameter) { + this->register_parameter(name, std::move(parameter)); + }; + qkv_proj_ = std::make_shared( + hidden_size_, head_dim_, total_num_heads, total_num_kv_heads, + "q_proj", "k_proj", "v_proj", register_fn, + quantization_method, use_bias, dtype, device, rank_info); + o_proj_ = this->register_module( + "o_proj", total_num_heads * head_dim_, hidden_size_, quantization_method, + use_output_bias, dtype, device, tp_rank, tp_size, rank_info.comm); + + rotary_emb_ = infinilm::layers::rotary_embedding::get_rope(model_config, device); + rope_theta_ = model_config->get_or("rope_theta", 500000.0); + if (config_json.contains("rope_scaling") && config_json["rope_scaling"].contains("mrope_section")) { + const auto §ions = config_json["rope_scaling"]["mrope_section"]; + if (sections.is_array() && sections.size() == 3) { + mrope_section_h_ = sections.at(0).get(); + mrope_section_w_ = sections.at(1).get(); + mrope_section_t_ = sections.at(2).get(); + } + } + + const float scaling = 1.0f / std::sqrt(static_cast(head_dim_)); + attn_ = std::make_shared( + num_attention_heads_, head_dim_, scaling, num_key_value_heads_, layer_idx_, + kv_cache_k_scale_, kv_cache_v_scale_, attention_backend_); + + infinilm::layers::attention::init_kv_cache_quant_params(register_fn, device, kv_cache_k_scale_, kv_cache_v_scale_); +} + +infinicore::Tensor Ernie4_5Attention::forward(const infinicore::Tensor &positions, + const infinicore::Tensor &hidden_states) const { + if (::infinilm::backends::AttentionBackend::STATIC_ATTN == attention_backend_) { + return forward_static_(positions, hidden_states); + } + return forward_paged_(positions, hidden_states); +} + +infinicore::Tensor Ernie4_5Attention::forward_static_(const infinicore::Tensor &position_ids, + const infinicore::Tensor &hidden_states) const { + auto hidden_states_mutable = hidden_states; + const auto shape = hidden_states->shape(); + const size_t batch_size = shape[0]; + const size_t seq_len = shape[1]; + + auto [q, k, v] = qkv_proj_->forward_split(hidden_states_mutable); + auto q_reshaped = q->view({batch_size, seq_len, num_attention_heads_, head_dim_}); + auto k_reshaped = k->view({batch_size, seq_len, num_key_value_heads_, head_dim_}); + auto v_reshaped = v->view({batch_size, seq_len, num_key_value_heads_, head_dim_}); + + auto pos_shape = position_ids->shape(); + if ((pos_shape.size() == 3 && pos_shape[0] == 1 && pos_shape[1] == seq_len && pos_shape[2] == 3) || (pos_shape.size() == 2 && pos_shape[0] == seq_len && pos_shape[1] == 3)) { + ASSERT_EQ(batch_size, 1); + auto q_mrope = q_reshaped->squeeze(0); + auto k_mrope = k_reshaped->squeeze(0); + infinicore::Tensor pos_ids_for_mrope = position_ids; + if (pos_shape.size() == 3) { + pos_ids_for_mrope = position_ids->squeeze(0); + } + infinicore::op::ernie45_mrope_( + q_mrope, k_mrope, pos_ids_for_mrope, + rope_theta_, mrope_section_h_, mrope_section_w_, mrope_section_t_); + auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped); + return o_proj_->forward(attn_output); + } + + infinicore::Tensor pos_ids_for_rope = position_ids; + if (pos_shape.size() == 2) { + auto pos_narrowed = position_ids->narrow({{0, 0, 1}}); + pos_ids_for_rope = pos_narrowed->contiguous()->view({pos_shape[1]}); + } else if (pos_shape.size() == 1) { + pos_ids_for_rope = position_ids->contiguous(); + } else { + throw std::runtime_error("Ernie4_5Attention: Unexpected position_ids shape"); + } + + auto q_rope = infinicore::Tensor::empty({batch_size, num_attention_heads_, seq_len, head_dim_}, q_reshaped->dtype(), q_reshaped->device())->permute({0, 2, 1, 3}); + rotary_emb_->forward(q_rope, q_reshaped, pos_ids_for_rope); + rotary_emb_->forward(k_reshaped, pos_ids_for_rope, true); + + auto attn_output = attn_->forward(q_rope, k_reshaped, v_reshaped); + return o_proj_->forward(attn_output); +} + +infinicore::Tensor Ernie4_5Attention::forward_paged_(const infinicore::Tensor &position_ids, + const infinicore::Tensor &hidden_states) const { + auto hidden_states_mutable = hidden_states; + const auto shape = hidden_states->shape(); + const size_t batch_size = shape[0]; + const size_t seq_len = shape[1]; + + ASSERT_EQ(batch_size, 1); + + auto [q, k, v] = qkv_proj_->forward_split(hidden_states_mutable); + auto q_reshaped = q->view({seq_len, num_attention_heads_, head_dim_}); + auto k_reshaped = k->view({seq_len, num_key_value_heads_, head_dim_}); + auto v_reshaped = v->view({seq_len, num_key_value_heads_, head_dim_}); + + auto pos_shape = position_ids->shape(); + if ((pos_shape.size() == 3 && pos_shape[0] == 1 && pos_shape[1] == seq_len && pos_shape[2] == 3) || (pos_shape.size() == 2 && pos_shape[0] == seq_len && pos_shape[1] == 3)) { + infinicore::Tensor pos_ids_for_mrope = position_ids; + if (pos_shape.size() == 3) { + pos_ids_for_mrope = position_ids->squeeze(0); + } + infinicore::op::ernie45_mrope_( + q_reshaped, k_reshaped, pos_ids_for_mrope, + rope_theta_, mrope_section_h_, mrope_section_w_, mrope_section_t_); + auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped); + return o_proj_->forward(attn_output); + } + + infinicore::Tensor pos_ids_for_rope = position_ids; + if (pos_shape.size() == 2) { + auto pos_narrowed = position_ids->narrow({{0, 0, 1}}); + pos_ids_for_rope = pos_narrowed->view({pos_shape[1]}); + } else if (pos_shape.size() != 1) { + throw std::runtime_error("Ernie4_5Attention: Unexpected position_ids shape"); + } + + rotary_emb_->forward(q_reshaped, pos_ids_for_rope, true); + rotary_emb_->forward(k_reshaped, pos_ids_for_rope, true); + + auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped); + return o_proj_->forward(attn_output); +} + +} // namespace infinilm::models::ernie4_5_moe_vl diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_attention.hpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_attention.hpp new file mode 100644 index 000000000..1a3c6fbaf --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_attention.hpp @@ -0,0 +1,51 @@ +#pragma once + +#include "../../layers/common_modules.hpp" + +namespace infinilm::models::ernie4_5_moe_vl { + +class Ernie4_5Attention : public infinicore::nn::Module { +public: + Ernie4_5Attention(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &positions, + const infinicore::Tensor &hidden_states) const; + + void process_weights_after_loading() override { + qkv_proj_->process_weights_after_loading(); + } + + void reset_runtime_state() const override { + qkv_proj_->reset_runtime_state(); + } + +private: + infinicore::Tensor forward_static_(const infinicore::Tensor &positions, + const infinicore::Tensor &hidden_states) const; + + infinicore::Tensor forward_paged_(const infinicore::Tensor &positions, + const infinicore::Tensor &hidden_states) const; + + std::shared_ptr qkv_proj_; + std::shared_ptr o_proj_; + std::shared_ptr rotary_emb_; + std::shared_ptr attn_; + + ::infinilm::backends::AttentionBackend attention_backend_; + size_t layer_idx_{0}; + size_t num_attention_heads_{0}; + size_t num_key_value_heads_{0}; + size_t hidden_size_{0}; + size_t head_dim_{0}; + double rope_theta_{500000.0}; + size_t mrope_section_h_{22}; + size_t mrope_section_w_{22}; + size_t mrope_section_t_{20}; + + INFINICORE_NN_PARAMETER(kv_cache_k_scale); + INFINICORE_NN_PARAMETER(kv_cache_v_scale); +}; + +} // namespace infinilm::models::ernie4_5_moe_vl diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_decoder_layer.cpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_decoder_layer.cpp new file mode 100644 index 000000000..84f53fc1c --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_decoder_layer.cpp @@ -0,0 +1,111 @@ +#include "ernie4_5_decoder_layer.hpp" + +#include "infinicore/ops.hpp" + +#include + +namespace infinilm::models::ernie4_5_moe_vl { + +namespace { + +size_t min_config_index(const nlohmann::json &value, size_t fallback) { + if (value.is_array()) { + size_t result = value.empty() ? fallback : value.at(0).get(); + for (const auto &item : value) { + result = std::min(result, item.get()); + } + return result; + } + if (value.is_number_unsigned() || value.is_number_integer()) { + return value.get(); + } + return fallback; +} + +size_t max_config_index(const nlohmann::json &value, size_t fallback) { + if (value.is_array()) { + size_t result = value.empty() ? fallback : value.at(0).get(); + for (const auto &item : value) { + result = std::max(result, item.get()); + } + return result; + } + if (value.is_number_unsigned() || value.is_number_integer()) { + return value.get(); + } + return fallback; +} + +} // namespace + +Ernie4_5DecoderLayer::Ernie4_5DecoderLayer(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) + : layer_idx_(layer_idx), + use_moe_(is_moe_layer(model_config, layer_idx)) { + const auto &dtype = model_config->get_dtype(); + const size_t hidden_size = model_config->get("hidden_size"); + const double rms_norm_eps = model_config->get("rms_norm_eps"); + + INFINICORE_NN_MODULE_INIT(input_layernorm, hidden_size, rms_norm_eps, dtype, device); + INFINICORE_NN_MODULE_INIT(post_attention_layernorm, hidden_size, rms_norm_eps, dtype, device); + INFINICORE_NN_MODULE_INIT(self_attn, model_config, layer_idx, device); + + if (use_moe_) { + mlp_ = this->register_module("mlp", model_config, device, layer_idx); + } else { + mlp_ = this->register_module("mlp", model_config, device); + } +} + +std::tuple Ernie4_5DecoderLayer::forward(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) { + input_layernorm_->forward_inplace(hidden_states, residual); + hidden_states = self_attn_->forward(positions, hidden_states); + post_attention_layernorm_->forward_inplace(hidden_states, residual); + hidden_states = forward_mlp_(hidden_states); + return std::make_tuple(hidden_states, residual); +} + +infinicore::Tensor Ernie4_5DecoderLayer::forward(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states) { + auto residual = hidden_states; + hidden_states = input_layernorm_->forward(hidden_states); + hidden_states = self_attn_->forward(positions, hidden_states); + hidden_states = infinicore::op::add(residual, hidden_states); + + residual = hidden_states; + hidden_states = post_attention_layernorm_->forward(hidden_states); + hidden_states = forward_mlp_(hidden_states); + hidden_states = infinicore::op::add(residual, hidden_states); + return hidden_states; +} + +infinicore::Tensor Ernie4_5DecoderLayer::forward_mlp_(const infinicore::Tensor &hidden_states) const { + if (use_moe_) { + return std::static_pointer_cast(mlp_)->forward(hidden_states); + } + return std::static_pointer_cast(mlp_)->forward(hidden_states); +} + +bool Ernie4_5DecoderLayer::is_moe_layer(const std::shared_ptr &model_config, + size_t layer_idx) { + const auto &config_json = model_config->get_config_json(); + const bool use_moe = model_config->get_or("use_moe", true); + if (!use_moe) { + return false; + } + + const size_t interval = model_config->get_or("moe_layer_interval", 1); + const size_t start = config_json.contains("moe_layer_start_index") + ? min_config_index(config_json.at("moe_layer_start_index"), 0) + : 0; + const size_t end = config_json.contains("moe_layer_end_index") + ? max_config_index(config_json.at("moe_layer_end_index"), model_config->get("num_hidden_layers") - 1) + : model_config->get("num_hidden_layers") - 1; + + return interval != 0 && ((layer_idx + 1) % interval == 0) && layer_idx >= start && layer_idx <= end; +} + +} // namespace infinilm::models::ernie4_5_moe_vl diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_decoder_layer.hpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_decoder_layer.hpp new file mode 100644 index 000000000..68b22915c --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_decoder_layer.hpp @@ -0,0 +1,41 @@ +#pragma once + +#include "../../layers/common_modules.hpp" +#include "ernie4_5_attention.hpp" +#include "ernie4_5_sparse_moe_block.hpp" + +#include +#include + +namespace infinilm::models::ernie4_5_moe_vl { + +class Ernie4_5DecoderLayer : public infinicore::nn::Module { +public: + Ernie4_5DecoderLayer(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + + std::tuple forward(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual); + + infinicore::Tensor forward(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states); + + size_t layer_idx() const { return layer_idx_; } + +private: + infinicore::Tensor forward_mlp_(const infinicore::Tensor &hidden_states) const; + static bool is_moe_layer(const std::shared_ptr &model_config, + size_t layer_idx); + + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, input_layernorm); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_attention_layernorm); + INFINICORE_NN_MODULE(Ernie4_5Attention, self_attn); + std::shared_ptr mlp_; + + size_t layer_idx_{0}; + bool use_moe_{false}; +}; + +} // namespace infinilm::models::ernie4_5_moe_vl diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_for_conditional_generation.cpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_for_conditional_generation.cpp new file mode 100644 index 000000000..015f7f86f --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_for_conditional_generation.cpp @@ -0,0 +1,341 @@ +#include "ernie4_5_for_conditional_generation.hpp" +#include "../../global_state/global_state.hpp" +#include "../models_registry.hpp" + +#include "infinicore/nn/rope.hpp" +#include "infinicore/ops.hpp" + +#include +#include +#include +#include + +namespace infinilm::models::ernie4_5_moe_vl { + +namespace { + +size_t first_size_t_or(const nlohmann::json &config_json, + const std::string &key, + size_t default_value) { + if (!config_json.contains(key) || config_json.at(key).is_null()) { + return default_value; + } + const auto &value = config_json.at(key); + if (value.is_array()) { + if (value.empty()) { + return default_value; + } + return value.at(0).get(); + } + return value.get(); +} + +size_t second_size_t_or(const nlohmann::json &config_json, + const std::string &key, + size_t default_value) { + if (!config_json.contains(key) || config_json.at(key).is_null()) { + return default_value; + } + const auto &value = config_json.at(key); + if (value.is_array()) { + if (value.size() < 2) { + return first_size_t_or(config_json, key, default_value); + } + return value.at(1).get(); + } + return value.get(); +} + +void set_if_missing(nlohmann::json &config_json, + const std::string &key, + const nlohmann::json &value) { + if (!config_json.contains(key) || config_json.at(key).is_null()) { + config_json[key] = value; + } +} + +void append_visual_ranges(std::vector &ranges, + const infinicore::Tensor &image_bound, + const infinicore::Tensor &input_ids, + const nlohmann::json &config_json, + size_t token_offset) { + auto bounds_cpu = image_bound->to(infinicore::Device::cpu()); + auto bound_slice = bounds_cpu; + if (bound_slice->shape().size() == 3) { + ASSERT_EQ(bound_slice->size(0), 1); + bound_slice = bound_slice->squeeze(0); + } + + auto ids_cpu = input_ids->to(infinicore::Device::cpu())->contiguous(); + auto ids_shape = ids_cpu->shape(); + const size_t seq_len = ids_shape.empty() ? 0 : ids_shape.back(); + std::vector boundary_token_ids; + for (const auto *key : {"image_start_token_id", "image_end_token_id", "video_start_token_id", "video_end_token_id"}) { + if (config_json.contains(key) && !config_json.at(key).is_null()) { + boundary_token_ids.push_back(config_json.at(key).get()); + } + } + + auto is_boundary_token = [&boundary_token_ids](int64_t token_id) { + return std::find(boundary_token_ids.begin(), boundary_token_ids.end(), token_id) != boundary_token_ids.end(); + }; + + auto token_at = [&ids_cpu](size_t index) -> int64_t { + if (ids_cpu->dtype() == infinicore::DataType::I64) { + const auto *ids = reinterpret_cast(ids_cpu->data()); + return ids[index]; + } + if (ids_cpu->dtype() == infinicore::DataType::I32) { + const auto *ids = reinterpret_cast(ids_cpu->data()); + return ids[index]; + } + throw std::runtime_error("Ernie4_5ForConditionalGeneration: input_ids must be int32 or int64"); + }; + + for (size_t i = 0; i < bound_slice->size(0); ++i) { + auto bound = bound_slice->narrow({{0, i, 1}})->squeeze(0); + const auto *bound_ptr = reinterpret_cast(bound->data()); + const int64_t start = bound_ptr[0]; + const int64_t end = bound_ptr[1]; + if (end > start) { + size_t range_start = static_cast(start); + size_t range_end = static_cast(end); + if (range_start > 0 && range_start - 1 < seq_len && is_boundary_token(token_at(range_start - 1))) { + --range_start; + } + if (range_end < seq_len && is_boundary_token(token_at(range_end))) { + ++range_end; + } + ranges.push_back(token_offset + range_start); + ranges.push_back(token_offset + range_end); + } + } +} + +infinicore::Tensor expand_grid_for_vision_tower(const infinicore::Tensor &grid_thw) { + auto grid_cpu = grid_thw->to(infinicore::Device::cpu())->contiguous(); + auto shape = grid_cpu->shape(); + if (shape.size() != 2 || shape[1] != 3) { + throw std::runtime_error("Ernie4_5ForConditionalGeneration: grid_thw must have shape [num_items, 3]"); + } + + std::vector expanded; + expanded.reserve(shape[0] * 3); + const auto append_grid = [&expanded](int64_t t, int64_t h, int64_t w) { + for (int64_t ti = 0; ti < t; ++ti) { + (void)ti; + expanded.push_back(1); + expanded.push_back(h); + expanded.push_back(w); + } + }; + + if (grid_cpu->dtype() == infinicore::DataType::I64) { + const auto *ptr = reinterpret_cast(grid_cpu->data()); + for (size_t i = 0; i < shape[0]; ++i) { + append_grid(ptr[i * 3 + 0], ptr[i * 3 + 1], ptr[i * 3 + 2]); + } + } else if (grid_cpu->dtype() == infinicore::DataType::I32) { + const auto *ptr = reinterpret_cast(grid_cpu->data()); + for (size_t i = 0; i < shape[0]; ++i) { + append_grid(ptr[i * 3 + 0], ptr[i * 3 + 1], ptr[i * 3 + 2]); + } + } else { + throw std::runtime_error("Ernie4_5ForConditionalGeneration: grid_thw must be int32 or int64"); + } + + auto expanded_cpu = infinicore::Tensor::from_blob( + expanded.data(), + {expanded.size() / 3, 3}, + infinicore::DataType::I64, + infinicore::Device::cpu()); + return expanded_cpu->to(grid_thw->device()); +} + +} // namespace + +Ernie4_5ForConditionalGeneration::Ernie4_5ForConditionalGeneration( + std::shared_ptr model_config, + const infinicore::Device &device) { + model_config_ = model_config; + const auto &dtype = model_config->get_dtype(); + const size_t hidden_size = model_config->get("hidden_size"); + const size_t vocab_size = model_config->get("vocab_size"); + const auto &config_json = model_config->get_config_json(); + + INFINICORE_NN_MODULE_INIT(model, model_config, device); + if (config_json.contains("vision_config")) { + INFINICORE_NN_MODULE_INIT(vision_model, + config_json.at("vision_config"), + model_config->get_or("rms_norm_eps", 1e-6), + dtype, + device); + INFINICORE_NN_MODULE_INIT(resampler_model, config_json, dtype, device); + } + INFINICORE_NN_MODULE_INIT(lm_head, hidden_size, vocab_size, false, dtype, device); +} + +infinicore::Tensor Ernie4_5ForConditionalGeneration::replace_embeddings( + const infinicore::Tensor &inputs_embeds, + const infinicore::Tensor &vision_hidden, + const infinicore::Tensor &image_bound) const { + auto out = infinicore::Tensor::empty(inputs_embeds->shape(), inputs_embeds->dtype(), inputs_embeds->device()); + out->copy_from(inputs_embeds); + + auto bounds_cpu = image_bound->to(infinicore::Device::cpu()); + const size_t batch_size = inputs_embeds->size(0); + ASSERT_EQ(batch_size, 1); + + auto out_slice = out->squeeze(0); + auto bound_slice = bounds_cpu; + if (bound_slice->shape().size() == 3) { + ASSERT_EQ(bound_slice->size(0), 1); + bound_slice = bound_slice->squeeze(0); + } + + const size_t num_ranges = bound_slice->size(0); + const size_t vision_len = vision_hidden->size(0); + size_t vision_offset = 0; + for (size_t i = 0; i < num_ranges; ++i) { + auto bound = bound_slice->narrow({{0, i, 1}})->squeeze(0); + auto *bound_ptr = reinterpret_cast(bound->data()); + const int64_t start = bound_ptr[0]; + const int64_t end = bound_ptr[1]; + if (end <= start) { + continue; + } + const size_t len = static_cast(end - start); + if (vision_offset + len > vision_len) { + throw std::runtime_error("Ernie4_5ForConditionalGeneration: image_bound exceeds vision hidden length"); + } + auto patch_embed = vision_hidden->narrow({{0, vision_offset, len}}); + out_slice->narrow({{0, static_cast(start), len}})->copy_from(patch_embed); + vision_offset += len; + } + if (vision_offset != vision_len) { + throw std::runtime_error("Ernie4_5ForConditionalGeneration: unused vision hidden tokens after image_bound replacement"); + } + return out; +} + +infinilm::InfinilmModel::Output Ernie4_5ForConditionalGeneration::forward( + const infinilm::InfinilmModel::Input &input) const { + auto &mm_metadata = global_state::get_forward_context().mm_metadata; + mm_metadata.visual_token_ranges.reset(); + if (input.pixel_values.has_value() && !input.pixel_values->empty()) { + mm_metadata.visual_token_ranges = std::vector{}; + if (!input.input_ids.has_value()) { + throw std::runtime_error("Ernie4_5ForConditionalGeneration: input_ids is required"); + } + if (!input.position_ids.has_value()) { + throw std::runtime_error("Ernie4_5ForConditionalGeneration: position_ids is required"); + } + if (!input.tgt_sizes.has_value()) { + throw std::runtime_error("Ernie4_5ForConditionalGeneration: tgt_sizes/image_grid_thw is required for vision input"); + } + if (!input.image_bound.has_value()) { + throw std::runtime_error("Ernie4_5ForConditionalGeneration: image_bound is required for vision input"); + } + if (input.pixel_values->size() != input.image_bound->size() || input.pixel_values->size() != input.tgt_sizes->size()) { + throw std::runtime_error("Ernie4_5ForConditionalGeneration: pixel_values, image_bound and tgt_sizes must have the same number of elements"); + } + + const auto &config_json = model_config_->get_config_json(); + auto input_ids = input.input_ids.value(); + auto inputs_embeds = model_->embed_tokens(input_ids); + + const auto &pixel_values = input.pixel_values.value(); + const auto &image_bound = input.image_bound.value(); + const auto &tgt_sizes = input.tgt_sizes.value(); + + if (input.input_offsets.has_value() && mm_metadata.image_req_ids.has_value()) { + auto input_offsets_cpu = input.input_offsets.value()->to(infinicore::Device::cpu()); + const auto *offsets = reinterpret_cast(input_offsets_cpu->data()); + const auto &image_req_ids = mm_metadata.image_req_ids.value(); + if (image_req_ids.size() != pixel_values.size()) { + throw std::runtime_error("Ernie4_5ForConditionalGeneration: image_req_ids size does not match pixel_values"); + } + for (size_t image_idx = 0; image_idx < pixel_values.size(); ++image_idx) { + const size_t req_id = image_req_ids.at(image_idx); + auto vision_grid = expand_grid_for_vision_tower(tgt_sizes.at(image_idx)); + auto vision_hidden = vision_model_->forward(pixel_values.at(image_idx), vision_grid); + vision_hidden = resampler_model_->forward(vision_hidden, tgt_sizes.at(image_idx)); + auto request_embeds = inputs_embeds->narrow( + {{1, static_cast(offsets[req_id]), static_cast(offsets[req_id + 1] - offsets[req_id])}}); + auto replaced = replace_embeddings(request_embeds, vision_hidden, image_bound.at(image_idx)); + request_embeds->copy_from(replaced); + auto request_input_ids = input_ids->narrow( + {{1, static_cast(offsets[req_id]), static_cast(offsets[req_id + 1] - offsets[req_id])}}); + append_visual_ranges( + mm_metadata.visual_token_ranges.value(), + image_bound.at(image_idx), + request_input_ids, + config_json, + static_cast(offsets[req_id])); + } + } else { + if (pixel_values.size() != 1) { + throw std::runtime_error("Ernie4_5ForConditionalGeneration: batched vision input requires image_req_ids"); + } + auto vision_grid = expand_grid_for_vision_tower(tgt_sizes.at(0)); + auto vision_hidden = vision_model_->forward(pixel_values.at(0), vision_grid); + vision_hidden = resampler_model_->forward(vision_hidden, tgt_sizes.at(0)); + inputs_embeds = replace_embeddings(inputs_embeds, vision_hidden, image_bound.at(0)); + append_visual_ranges(mm_metadata.visual_token_ranges.value(), image_bound.at(0), input_ids, config_json, 0); + } + + auto hidden_states = model_->forward_embeds(inputs_embeds, input.position_ids.value()); + mm_metadata.visual_token_ranges.reset(); + return {lm_head_->forward(hidden_states)}; + } + auto hidden_states = model_->forward(input); + return {lm_head_->forward(hidden_states)}; +} + +std::shared_ptr create_ernie4_5_moe_vl_model_config(std::shared_ptr model_config) { + const std::string &model_type = model_config->get("model_type"); + if ("ernie4_5_moe_vl" != model_type && "ernie4_5_vl_moe" != model_type) { + throw std::runtime_error("create_ernie4_5_moe_vl_model_config: unsupported model_type: " + model_type); + } + + nlohmann::json &config_json = model_config->get_config_json(); + const size_t hidden_size = model_config->get("hidden_size"); + const size_t num_attention_heads = model_config->get("num_attention_heads"); + + set_if_missing(config_json, "head_dim", hidden_size / num_attention_heads); + set_if_missing(config_json, "attention_bias", false); + set_if_missing(config_json, "attention_output_bias", false); + set_if_missing(config_json, "mlp_bias", false); + set_if_missing(config_json, "use_moe", true); + + set_if_missing(config_json, "vision_num_experts", second_size_t_or(config_json, "moe_num_experts", 0)); + set_if_missing( + config_json, + "vision_moe_intermediate_size", + second_size_t_or(config_json, "moe_intermediate_size", model_config->get("intermediate_size"))); + set_if_missing(config_json, "num_experts", first_size_t_or(config_json, "moe_num_experts", 0)); + set_if_missing(config_json, "num_experts_per_tok", first_size_t_or(config_json, "moe_k", 1)); + config_json["moe_intermediate_size"] = first_size_t_or(config_json, "moe_intermediate_size", model_config->get("intermediate_size")); + set_if_missing(config_json, "norm_topk_prob", true); + set_if_missing(config_json, "moe_router_backend", "softmax"); + set_if_missing(config_json, "moe_router_dtype", "float32"); + set_if_missing(config_json, "moe_router_use_correction_bias", true); + set_if_missing(config_json, "enable_ernie_vision_experts", true); + + model_config->set_rope_algo(infinicore::nn::RoPE::Algo::GPT_J); + + return model_config; +} + +} // namespace infinilm::models::ernie4_5_moe_vl + +namespace { +INFINILM_REGISTER_CAUSAL_LM_MODEL( + ernie4_5_moe_vl, + infinilm::models::ernie4_5_moe_vl::Ernie4_5ForConditionalGeneration, + infinilm::models::ernie4_5_moe_vl::create_ernie4_5_moe_vl_model_config); +INFINILM_REGISTER_CAUSAL_LM_MODEL( + ernie4_5_vl_moe, + infinilm::models::ernie4_5_moe_vl::Ernie4_5ForConditionalGeneration, + infinilm::models::ernie4_5_moe_vl::create_ernie4_5_moe_vl_model_config); +} // namespace diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_for_conditional_generation.hpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_for_conditional_generation.hpp new file mode 100644 index 000000000..0d1e6a3fa --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_for_conditional_generation.hpp @@ -0,0 +1,33 @@ +#pragma once + +#include "ernie4_5_decoder_layer.hpp" +#include "ernie4_5_vision.hpp" + +namespace infinilm::models::ernie4_5_moe_vl { + +using Ernie4_5Model = infinilm::layers::causal_lm_templates::TextModel; + +class Ernie4_5ForConditionalGeneration : public infinilm::InfinilmModel { +public: + Ernie4_5ForConditionalGeneration(std::shared_ptr model_config, + const infinicore::Device &device); + + infinilm::InfinilmModel::Output forward(const infinilm::InfinilmModel::Input &input) const override; + bool supports_mrope_position_ids() const override { + return true; + } + +private: + infinicore::Tensor replace_embeddings(const infinicore::Tensor &inputs_embeds, + const infinicore::Tensor &vision_hidden, + const infinicore::Tensor &image_bound) const; + + INFINICORE_NN_MODULE(Ernie4_5Model, model); + INFINICORE_NN_MODULE(Ernie4_5VisionTransformer, vision_model); + INFINICORE_NN_MODULE(Ernie4_5VariableResolutionResampler, resampler_model); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, lm_head); +}; + +std::shared_ptr create_ernie4_5_moe_vl_model_config(std::shared_ptr model_config); + +} // namespace infinilm::models::ernie4_5_moe_vl diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_sparse_moe_block.cpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_sparse_moe_block.cpp new file mode 100644 index 000000000..a00bd243b --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_sparse_moe_block.cpp @@ -0,0 +1,159 @@ +#include "ernie4_5_sparse_moe_block.hpp" + +#include "../../global_state/global_state.hpp" +#include "infinicore/ops.hpp" + +#include + +namespace infinilm::models::ernie4_5_moe_vl { + +namespace { + +std::shared_ptr make_shared_experts_config( + const std::shared_ptr &model_config) { + auto shared_config_json = model_config->get_config_json(); + const size_t num_shared_experts = model_config->get_or("moe_num_shared_experts", 0); + const size_t moe_intermediate_size = model_config->get("moe_intermediate_size"); + shared_config_json["intermediate_size"] = moe_intermediate_size * num_shared_experts; + return std::make_shared(shared_config_json); +} + +size_t first_size_t_or(const nlohmann::json &config_json, + const std::string &key, + size_t default_value) { + if (!config_json.contains(key) || config_json.at(key).is_null()) { + return default_value; + } + const auto &value = config_json.at(key); + if (value.is_array()) { + if (value.empty()) { + return default_value; + } + return value.at(0).get(); + } + return value.get(); +} + +size_t second_size_t_or(const nlohmann::json &config_json, + const std::string &key, + size_t default_value) { + if (!config_json.contains(key) || config_json.at(key).is_null()) { + return default_value; + } + const auto &value = config_json.at(key); + if (value.is_array()) { + if (value.size() < 2) { + return first_size_t_or(config_json, key, default_value); + } + return value.at(1).get(); + } + return value.get(); +} + +std::shared_ptr make_vision_experts_config( + const std::shared_ptr &model_config) { + auto vision_config_json = model_config->get_config_json(); + vision_config_json["num_experts"] = vision_config_json.value( + "vision_num_experts", + second_size_t_or(vision_config_json, "moe_num_experts", model_config->get("num_experts"))); + vision_config_json["moe_intermediate_size"] = vision_config_json.value( + "vision_moe_intermediate_size", + second_size_t_or( + vision_config_json, + "moe_intermediate_size", + model_config->get("moe_intermediate_size"))); + return std::make_shared(vision_config_json); +} + +} // namespace + +Ernie4_5SparseMoeBlock::Ernie4_5SparseMoeBlock(std::shared_ptr model_config, + const infinicore::Device &device, + size_t layer_id) + : infinilm::layers::moe::SparseMoeBlock(model_config, device, layer_id) { + has_shared_experts_ = model_config->get_or("moe_num_shared_experts", 0) > 0; + if (has_shared_experts_) { + auto shared_config = make_shared_experts_config(model_config); + INFINICORE_NN_MODULE_INIT(shared_experts, shared_config, device); + } + + const auto &config_json = model_config->get_config_json(); + has_vision_experts_ = model_config->get_or("enable_ernie_vision_experts", false) && config_json.contains("moe_num_experts") && config_json.at("moe_num_experts").is_array() && config_json.at("moe_num_experts").size() > 1; + if (has_vision_experts_) { + auto vision_config = make_vision_experts_config(model_config); + INFINICORE_NN_MODULE_INIT(vision_gate, vision_config, device); + INFINICORE_NN_MODULE_INIT(vision_experts, vision_config, device); + INFINICORE_NN_MODULE_INIT(vision_fused_moe, vision_config, device, layer_id); + } +} + +infinicore::Tensor Ernie4_5SparseMoeBlock::forward_routed_( + const infinicore::Tensor &hidden_states, + const infinilm::layers::moe::TopKRouter &gate, + const infinilm::layers::moe::FusedMoeExperts &experts, + const infinilm::layers::moe::FusedMoE &fused_moe) const { + auto [routing_weights, selected_experts] = gate.forward(hidden_states); + infinilm::layers::moe::TopKOutput topk_output{ + routing_weights, + selected_experts, + infinicore::Tensor(), + }; + + auto output = fused_moe.forward( + hidden_states, + topk_output, + experts.moe_weights()); + if (has_shared_experts_) { + auto shared_output = shared_experts_->forward(hidden_states); + output = infinicore::op::add(output, shared_output); + } + return output; +} + +infinicore::Tensor Ernie4_5SparseMoeBlock::forward_text_(const infinicore::Tensor &hidden_states) const { + return forward_routed_(hidden_states, *gate_, *experts_, *fused_moe_); +} + +infinicore::Tensor Ernie4_5SparseMoeBlock::forward_vision_(const infinicore::Tensor &hidden_states) const { + return forward_routed_(hidden_states, *vision_gate_, *vision_experts_, *vision_fused_moe_); +} + +infinicore::Tensor Ernie4_5SparseMoeBlock::forward(const infinicore::Tensor &hidden_states) const { + ASSERT(hidden_states->ndim() == 3); + const auto &visual_ranges = infinilm::global_state::get_forward_context().mm_metadata.visual_token_ranges; + auto flat_hidden = hidden_states->view({hidden_states->size(0) * hidden_states->size(1), hidden_states->size(2)})->contiguous(); + if (!has_vision_experts_ || !visual_ranges.has_value() || visual_ranges->empty()) { + return forward_text_(flat_hidden)->view(hidden_states->shape()); + } + + auto flat_output = infinicore::Tensor::empty(flat_hidden->shape(), flat_hidden->dtype(), flat_hidden->device()); + + size_t current = 0; + const size_t ntokens = flat_hidden->size(0); + const auto &ranges = visual_ranges.value(); + for (size_t i = 0; i + 1 < ranges.size(); i += 2) { + const size_t range_start = ranges[i]; + const size_t range_end = ranges[i + 1]; + const size_t start = std::min(range_start, ntokens); + const size_t end = std::min(range_end, ntokens); + if (start > current) { + const size_t len = start - current; + auto text_output = forward_text_(flat_hidden->narrow({{0, current, len}})->contiguous()); + flat_output->narrow({{0, current, len}})->copy_from(text_output); + } + if (end > start) { + const size_t len = end - start; + auto vision_output = forward_vision_(flat_hidden->narrow({{0, start, len}})->contiguous()); + flat_output->narrow({{0, start, len}})->copy_from(vision_output); + } + current = std::max(current, end); + } + if (current < ntokens) { + const size_t len = ntokens - current; + auto text_output = forward_text_(flat_hidden->narrow({{0, current, len}})->contiguous()); + flat_output->narrow({{0, current, len}})->copy_from(text_output); + } + return flat_output->view(hidden_states->shape()); +} + +} // namespace infinilm::models::ernie4_5_moe_vl diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_sparse_moe_block.hpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_sparse_moe_block.hpp new file mode 100644 index 000000000..2a34377df --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_sparse_moe_block.hpp @@ -0,0 +1,40 @@ +#pragma once + +#include "../../config/model_config.hpp" +#include "../../layers/mlp/mlp.hpp" +#include "../../layers/moe/experts/fused_moe_experts.hpp" +#include "../../layers/moe/fused_moe.hpp" +#include "../../layers/moe/router/topk_router.hpp" +#include "../../layers/moe/sparse_moe_block.hpp" + +#include + +namespace infinilm::models::ernie4_5_moe_vl { + +class Ernie4_5SparseMoeBlock final : public infinilm::layers::moe::SparseMoeBlock { +public: + Ernie4_5SparseMoeBlock(std::shared_ptr model_config, + const infinicore::Device &device, + size_t layer_id = 0); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + +private: + infinicore::Tensor forward_routed_(const infinicore::Tensor &hidden_states, + const infinilm::layers::moe::TopKRouter &gate, + const infinilm::layers::moe::FusedMoeExperts &experts, + const infinilm::layers::moe::FusedMoE &fused_moe) const; + + infinicore::Tensor forward_text_(const infinicore::Tensor &hidden_states) const; + infinicore::Tensor forward_vision_(const infinicore::Tensor &hidden_states) const; + + INFINICORE_NN_MODULE(infinilm::layers::mlp::MLP, shared_experts); + INFINICORE_NN_MODULE(infinilm::layers::moe::TopKRouter, vision_gate); + INFINICORE_NN_MODULE(infinilm::layers::moe::FusedMoeExperts, vision_experts); + INFINICORE_NN_MODULE(infinilm::layers::moe::FusedMoE, vision_fused_moe); + + bool has_shared_experts_{false}; + bool has_vision_experts_{false}; +}; + +} // namespace infinilm::models::ernie4_5_moe_vl diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_vision.cpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_vision.cpp new file mode 100644 index 000000000..67db763ce --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_vision.cpp @@ -0,0 +1,331 @@ +#include "ernie4_5_vision.hpp" + +#include "infinicore/ops.hpp" +#include "infinicore/ops/mha_varlen.hpp" + +#include +#include +#include +#include +#include + +namespace infinilm::models::ernie4_5_moe_vl { +namespace { + +struct GridTHW { + int64_t t; + int64_t h; + int64_t w; +}; + +std::vector read_grid_thw_cpu(const infinicore::Tensor &grid_thw) { + auto grid_cpu = grid_thw->to(infinicore::Device::cpu()); + auto shape = grid_cpu->shape(); + if (shape.size() != 2 || shape[1] != 3) { + throw std::runtime_error("Ernie4_5 VL: grid_thw must have shape [num_images, 3]"); + } + + std::vector grids; + grids.reserve(shape[0]); + for (size_t i = 0; i < shape[0]; ++i) { + GridTHW grid{0, 0, 0}; + if (grid_cpu->dtype() == infinicore::DataType::I64) { + const auto *grid_ptr = reinterpret_cast(grid_cpu->data()); + grid = {grid_ptr[i * 3 + 0], grid_ptr[i * 3 + 1], grid_ptr[i * 3 + 2]}; + } else if (grid_cpu->dtype() == infinicore::DataType::I32) { + const auto *grid_ptr = reinterpret_cast(grid_cpu->data()); + grid = {grid_ptr[i * 3 + 0], grid_ptr[i * 3 + 1], grid_ptr[i * 3 + 2]}; + } else { + throw std::runtime_error("Ernie4_5 VL: grid_thw must be int32 or int64"); + } + if (grid.t <= 0 || grid.h <= 0 || grid.w <= 0) { + throw std::runtime_error("Ernie4_5 VL: invalid grid_thw"); + } + grids.push_back(grid); + } + return grids; +} + +std::pair build_vision_cu_seqlens(const infinicore::Tensor &grid_thw, + size_t seq_len) { + std::vector cu_seqlens; + cu_seqlens.push_back(0); + int max_seqlen = 0; + size_t total = 0; + for (const auto &grid : read_grid_thw_cpu(grid_thw)) { + const int segment_len = static_cast(grid.h * grid.w); + max_seqlen = std::max(max_seqlen, segment_len); + for (int64_t ti = 0; ti < grid.t; ++ti) { + (void)ti; + total += static_cast(segment_len); + cu_seqlens.push_back(static_cast(total)); + } + } + if (total != seq_len) { + throw std::runtime_error("Ernie4_5VisionAttention: grid_thw does not match sequence length"); + } + auto cu_cpu = infinicore::Tensor::from_blob( + cu_seqlens.data(), + {cu_seqlens.size()}, + infinicore::DataType::I32, + infinicore::Device::cpu()); + return {cu_cpu->to(grid_thw->device()), max_seqlen}; +} + +} // namespace + +Ernie4_5VisionPatchEmbed::Ernie4_5VisionPatchEmbed(const nlohmann::json &vision_config, + const infinicore::DataType &dtype, + const infinicore::Device &device) { + const size_t in_chans = vision_config.value("in_chans", vision_config.value("in_channels", 3)); + const size_t patch_size = vision_config.value("patch_size", vision_config.value("spatial_patch_size", 14)); + const size_t embed_dim = vision_config.value("embed_dim", vision_config.value("hidden_size", 1280)); + INFINICORE_NN_MODULE_INIT(proj, in_chans * patch_size * patch_size, embed_dim, false, dtype, device); +} + +infinicore::Tensor Ernie4_5VisionPatchEmbed::forward(const infinicore::Tensor &pixel_values) const { + return proj_->forward(const_cast(pixel_values)); +} + +Ernie4_5VisionMLP::Ernie4_5VisionMLP(size_t hidden_size, + double mlp_ratio, + const infinicore::DataType &dtype, + const infinicore::Device &device) { + const size_t intermediate_size = static_cast(static_cast(hidden_size) * mlp_ratio); + INFINICORE_NN_MODULE_INIT(fc1, hidden_size, intermediate_size, true, dtype, device); + INFINICORE_NN_MODULE_INIT(fc2, intermediate_size, hidden_size, true, dtype, device); +} + +infinicore::Tensor Ernie4_5VisionMLP::forward(const infinicore::Tensor &hidden_states) const { + auto x = fc1_->forward(const_cast(hidden_states)); + x = infinicore::op::quick_gelu(x); + return fc2_->forward(x); +} + +Ernie4_5VisionAttention::Ernie4_5VisionAttention(size_t hidden_size, + size_t num_heads, + const infinicore::DataType &dtype, + const infinicore::Device &device) + : hidden_size_(hidden_size), + num_heads_(num_heads), + head_dim_(hidden_size / num_heads), + scale_(1.0f / std::sqrt(static_cast(head_dim_))) { + if (hidden_size_ % num_heads_ != 0) { + throw std::runtime_error("Ernie4_5VisionAttention: hidden_size must be divisible by num_heads"); + } + INFINICORE_NN_MODULE_INIT(qkv, hidden_size_, 3 * hidden_size_, true, dtype, device); + INFINICORE_NN_MODULE_INIT(proj, hidden_size_, hidden_size_, true, dtype, device); +} + +infinicore::Tensor Ernie4_5VisionAttention::forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &rotary_pos_ids, + const infinicore::Tensor &grid_thw) const { + const size_t seq_len = hidden_states->size(0); + auto qkv = qkv_->forward(const_cast(hidden_states)); + auto q = qkv->narrow({{1, 0, hidden_size_}}) + ->view({seq_len, num_heads_, head_dim_}) + ->contiguous(); + auto k = qkv->narrow({{1, hidden_size_, hidden_size_}}) + ->view({seq_len, num_heads_, head_dim_}) + ->contiguous(); + auto v = qkv->narrow({{1, 2 * hidden_size_, hidden_size_}}) + ->view({seq_len, num_heads_, head_dim_}) + ->contiguous(); + + infinicore::op::ernie45_vision_rope_(q, k, rotary_pos_ids, 10000.0); + auto [cu_seqlens, max_seqlen] = build_vision_cu_seqlens(grid_thw, seq_len); + auto out_tokens = infinicore::op::mha_varlen( + q, + k, + v, + cu_seqlens, + cu_seqlens, + std::nullopt, + max_seqlen, + max_seqlen, + std::nullopt, + scale_); + auto out = out_tokens->view({seq_len, hidden_size_})->contiguous(); + return proj_->forward(out); +} + +Ernie4_5VisionBlock::Ernie4_5VisionBlock(const nlohmann::json &vision_config, + const infinicore::DataType &dtype, + const infinicore::Device &device) { + const size_t hidden_size = vision_config.value("hidden_size", vision_config.value("embed_dim", 1280)); + const size_t num_heads = vision_config.value("num_heads", 16); + const double mlp_ratio = vision_config.value("mlp_ratio", 4.0); + INFINICORE_NN_MODULE_INIT(norm1, hidden_size, 1e-6, dtype, device); + INFINICORE_NN_MODULE_INIT(attn, hidden_size, num_heads, dtype, device); + INFINICORE_NN_MODULE_INIT(norm2, hidden_size, 1e-6, dtype, device); + INFINICORE_NN_MODULE_INIT(mlp, hidden_size, mlp_ratio, dtype, device); +} + +infinicore::Tensor Ernie4_5VisionBlock::forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &rotary_pos_ids, + const infinicore::Tensor &grid_thw) const { + auto residual = hidden_states; + auto x = norm1_->forward(hidden_states); + x = attn_->forward(x, rotary_pos_ids, grid_thw); + x = infinicore::op::add(x, residual); + + residual = x; + x = norm2_->forward(x); + x = mlp_->forward(x); + return infinicore::op::add(x, residual); +} + +Ernie4_5VisionTransformer::Ernie4_5VisionTransformer(const nlohmann::json &vision_config, + double norm_eps, + const infinicore::DataType &dtype, + const infinicore::Device &device) + : spatial_merge_size_(vision_config.value("spatial_merge_size", 2)) { + const size_t hidden_size = vision_config.value("hidden_size", vision_config.value("embed_dim", 1280)); + const size_t depth = vision_config.value("depth", 32); + INFINICORE_NN_MODULE_INIT(patch_embed, vision_config, dtype, device); + INFINICORE_NN_MODULE_VEC_INIT(blocks, depth, Ernie4_5VisionBlock, vision_config, dtype, device); + INFINICORE_NN_MODULE_INIT(ln, hidden_size, norm_eps, dtype, device); +} + +infinicore::Tensor Ernie4_5VisionTransformer::forward(const infinicore::Tensor &pixel_values, + const infinicore::Tensor &grid_thw) const { + auto hidden_states = patch_embed_->forward(pixel_values); + auto rotary_pos_ids = build_rotary_pos_ids(grid_thw, hidden_states->size(0)); + for (const auto &block : blocks_) { + hidden_states = block->forward(hidden_states, rotary_pos_ids, grid_thw); + } + return ln_->forward(hidden_states); +} + +infinicore::Tensor Ernie4_5VisionTransformer::build_rotary_pos_ids(const infinicore::Tensor &grid_thw, + size_t seq_len) const { + std::vector pos; + pos.reserve(seq_len * 2); + for (const auto &grid : read_grid_thw_cpu(grid_thw)) { + const int64_t t = grid.t; + const int64_t h = grid.h; + const int64_t w = grid.w; + if (t <= 0 || h <= 0 || w <= 0 || h % static_cast(spatial_merge_size_) != 0 || w % static_cast(spatial_merge_size_) != 0) { + throw std::runtime_error("Ernie4_5VisionTransformer: invalid grid_thw"); + } + for (int64_t ti = 0; ti < t; ++ti) { + (void)ti; + for (int64_t hb = 0; hb < h; hb += static_cast(spatial_merge_size_)) { + for (int64_t wb = 0; wb < w; wb += static_cast(spatial_merge_size_)) { + for (size_t hs = 0; hs < spatial_merge_size_; ++hs) { + for (size_t ws = 0; ws < spatial_merge_size_; ++ws) { + pos.push_back(static_cast(hb + static_cast(hs))); + pos.push_back(static_cast(wb + static_cast(ws))); + } + } + } + } + } + } + if (pos.size() != seq_len * 2) { + throw std::runtime_error("Ernie4_5VisionTransformer: rotary position count does not match vision sequence length"); + } + auto pos_cpu = infinicore::Tensor::from_blob( + pos.data(), {seq_len, 2}, infinicore::DataType::I32, infinicore::Device::cpu()); + return pos_cpu->to(grid_thw->device()); +} + +Ernie4_5VariableResolutionResampler::Ernie4_5VariableResolutionResampler( + const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device) + : in_dim_(config.value("pixel_hidden_size", 1280)), + out_dim_(config.value("hidden_size", 2560)), + spatial_conv_size_(config.value("spatial_conv_size", 2)), + temporal_conv_size_(config.value("temporal_conv_size", 2)), + use_temporal_conv_(config.value("use_temporal_conv", true)) { + const size_t spatial_dim = in_dim_ * spatial_conv_size_ * spatial_conv_size_; + const size_t temporal_dim = spatial_dim * temporal_conv_size_; + spatial_linear0_ = this->register_module("spatial_linear.0", spatial_dim, spatial_dim, true, dtype, device); + spatial_linear2_ = this->register_module("spatial_linear.2", spatial_dim, spatial_dim, true, dtype, device); + spatial_linear3_ = this->register_module("spatial_linear.3", spatial_dim, 1e-6, dtype, device); + if (use_temporal_conv_) { + temporal_linear0_ = this->register_module("temporal_linear.0", temporal_dim, spatial_dim, true, dtype, device); + temporal_linear2_ = this->register_module("temporal_linear.2", spatial_dim, spatial_dim, true, dtype, device); + temporal_linear3_ = this->register_module("temporal_linear.3", spatial_dim, 1e-6, dtype, device); + } + mlp_ = this->register_module("mlp", spatial_dim, out_dim_, true, dtype, device); + after_norm_ = this->register_module("after_norm", out_dim_, config.value("rms_norm_eps", 1e-6), dtype, device); +} + +infinicore::Tensor Ernie4_5VariableResolutionResampler::spatial_forward(const infinicore::Tensor &hidden_states) const { + const size_t seq_len = hidden_states->size(0); + const size_t pack = spatial_conv_size_ * spatial_conv_size_; + if (seq_len % pack != 0) { + throw std::runtime_error("Ernie4_5VariableResolutionResampler: vision seq_len is not divisible by spatial_conv_size^2"); + } + auto x = hidden_states->contiguous()->view({seq_len / pack, in_dim_ * pack}); + x = spatial_linear0_->forward(x); + x = infinicore::op::gelu(x); + x = spatial_linear2_->forward(x); + return spatial_linear3_->forward(x); +} + +infinicore::Tensor Ernie4_5VariableResolutionResampler::temporal_forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &grid_thw) const { + if (temporal_conv_size_ != 2) { + throw std::runtime_error("Ernie4_5VariableResolutionResampler: only temporal_conv_size=2 is supported"); + } + + const auto grids = read_grid_thw_cpu(grid_thw); + const size_t spatial_pack = spatial_conv_size_ * spatial_conv_size_; + size_t spatial_base = 0; + size_t out_rows = 0; + for (const auto &grid : grids) { + if (grid.h % static_cast(spatial_conv_size_) != 0 || grid.w % static_cast(spatial_conv_size_) != 0) { + throw std::runtime_error("Ernie4_5VariableResolutionResampler: invalid spatial grid"); + } + if (grid.t > 1 && (grid.t % 2) != 0) { + throw std::runtime_error("Ernie4_5VariableResolutionResampler: odd video temporal size is not supported"); + } + const size_t spatial_size = static_cast(grid.h * grid.w) / spatial_pack; + spatial_base += static_cast(grid.t) * spatial_size; + out_rows += ((static_cast(grid.t) + 1) / 2) * spatial_size; + } + if (spatial_base != hidden_states->size(0)) { + throw std::runtime_error("Ernie4_5VariableResolutionResampler: grid_thw does not match spatial hidden length"); + } + + auto x = infinicore::Tensor::empty( + {out_rows, hidden_states->size(1) * 2}, + hidden_states->dtype(), + hidden_states->device()); + spatial_base = 0; + size_t out_offset = 0; + for (const auto &grid : grids) { + const size_t spatial_size = static_cast(grid.h * grid.w) / spatial_pack; + for (int64_t temp_offset = 0; temp_offset < grid.t; temp_offset += 2) { + const size_t first = spatial_base + static_cast(temp_offset) * spatial_size; + const size_t second_temp = static_cast((grid.t > 1) ? temp_offset + 1 : 0); + const size_t second = spatial_base + second_temp * spatial_size; + auto out_slice = x->narrow({{0, out_offset, spatial_size}}); + out_slice->narrow({{1, 0, hidden_states->size(1)}}) + ->copy_from(hidden_states->narrow({{0, first, spatial_size}})); + out_slice->narrow({{1, hidden_states->size(1), hidden_states->size(1)}}) + ->copy_from(hidden_states->narrow({{0, second, spatial_size}})); + out_offset += spatial_size; + } + spatial_base += static_cast(grid.t) * spatial_size; + } + x = temporal_linear0_->forward(x); + x = infinicore::op::gelu(x); + x = temporal_linear2_->forward(x); + return temporal_linear3_->forward(x); +} + +infinicore::Tensor Ernie4_5VariableResolutionResampler::forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &grid_thw) const { + auto x = spatial_forward(hidden_states); + if (use_temporal_conv_) { + x = temporal_forward(x, grid_thw); + } + x = mlp_->forward(x); + return after_norm_->forward(x); +} + +} // namespace infinilm::models::ernie4_5_moe_vl diff --git a/csrc/models/ernie4_5_moe_vl/ernie4_5_vision.hpp b/csrc/models/ernie4_5_moe_vl/ernie4_5_vision.hpp new file mode 100644 index 000000000..39c774a19 --- /dev/null +++ b/csrc/models/ernie4_5_moe_vl/ernie4_5_vision.hpp @@ -0,0 +1,127 @@ +#pragma once + +#include "infinicore/nn/layer_norm.hpp" +#include "infinicore/nn/linear.hpp" +#include "infinicore/nn/module.hpp" +#include "infinicore/nn/rmsnorm.hpp" +#include "infinicore/tensor.hpp" +#include + +namespace infinilm::models::ernie4_5_moe_vl { + +class Ernie4_5VisionPatchEmbed : public infinicore::nn::Module { +public: + Ernie4_5VisionPatchEmbed(const nlohmann::json &vision_config, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &pixel_values) const; + +private: + INFINICORE_NN_MODULE(infinicore::nn::Linear, proj); +}; + +class Ernie4_5VisionMLP : public infinicore::nn::Module { +public: + Ernie4_5VisionMLP(size_t hidden_size, + double mlp_ratio, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + +private: + INFINICORE_NN_MODULE(infinicore::nn::Linear, fc1); + INFINICORE_NN_MODULE(infinicore::nn::Linear, fc2); +}; + +class Ernie4_5VisionAttention : public infinicore::nn::Module { +public: + Ernie4_5VisionAttention(size_t hidden_size, + size_t num_heads, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &rotary_pos_ids, + const infinicore::Tensor &grid_thw) const; + +private: + size_t hidden_size_; + size_t num_heads_; + size_t head_dim_; + float scale_; + + INFINICORE_NN_MODULE(infinicore::nn::Linear, qkv); + INFINICORE_NN_MODULE(infinicore::nn::Linear, proj); +}; + +class Ernie4_5VisionBlock : public infinicore::nn::Module { +public: + Ernie4_5VisionBlock(const nlohmann::json &vision_config, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &rotary_pos_ids, + const infinicore::Tensor &grid_thw) const; + +private: + INFINICORE_NN_MODULE(infinicore::nn::LayerNorm, norm1); + INFINICORE_NN_MODULE(Ernie4_5VisionAttention, attn); + INFINICORE_NN_MODULE(infinicore::nn::LayerNorm, norm2); + INFINICORE_NN_MODULE(Ernie4_5VisionMLP, mlp); +}; + +class Ernie4_5VisionTransformer : public infinicore::nn::Module { +public: + Ernie4_5VisionTransformer(const nlohmann::json &vision_config, + double norm_eps, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &pixel_values, + const infinicore::Tensor &grid_thw) const; + +private: + size_t spatial_merge_size_{2}; + + infinicore::Tensor build_rotary_pos_ids(const infinicore::Tensor &grid_thw, + size_t seq_len) const; + + INFINICORE_NN_MODULE(Ernie4_5VisionPatchEmbed, patch_embed); + INFINICORE_NN_MODULE_VEC(Ernie4_5VisionBlock, blocks); + INFINICORE_NN_MODULE(infinicore::nn::LayerNorm, ln); +}; + +class Ernie4_5VariableResolutionResampler : public infinicore::nn::Module { +public: + Ernie4_5VariableResolutionResampler(const nlohmann::json &config, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &grid_thw) const; + +private: + infinicore::Tensor spatial_forward(const infinicore::Tensor &hidden_states) const; + infinicore::Tensor temporal_forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &grid_thw) const; + + size_t in_dim_; + size_t out_dim_; + size_t spatial_conv_size_; + size_t temporal_conv_size_; + bool use_temporal_conv_; + + std::shared_ptr spatial_linear0_; + std::shared_ptr spatial_linear2_; + std::shared_ptr spatial_linear3_; + std::shared_ptr temporal_linear0_; + std::shared_ptr temporal_linear2_; + std::shared_ptr temporal_linear3_; + std::shared_ptr mlp_; + std::shared_ptr after_norm_; +}; + +} // namespace infinilm::models::ernie4_5_moe_vl diff --git a/csrc/models/infinilm_model.hpp b/csrc/models/infinilm_model.hpp index ac994fd6d..c306ea4cb 100644 --- a/csrc/models/infinilm_model.hpp +++ b/csrc/models/infinilm_model.hpp @@ -20,7 +20,8 @@ class InfinilmModel : public infinicore::nn::Module { struct Input { /// Token IDs tensor of shape `[batch, seq_len]`. std::optional input_ids; - /// Position IDs tensor of shape `[batch, seq_len]` or `[seq_len]`. + /// Position IDs tensor of shape `[batch, seq_len]`, `[seq_len]`, + /// or `[batch, seq_len, 3]` / `[seq_len, 3]` for mrope models. std::optional position_ids; /// Past Lengths of cached sequence for each request, of shape `[num_requests]`. std::optional past_sequence_lengths; @@ -70,6 +71,9 @@ class InfinilmModel : public infinicore::nn::Module { virtual const cache::CacheConfig *get_cache_config() const { return cache_config_.get(); } + virtual bool supports_mrope_position_ids() const { + return false; + } void process_weights_after_loading(); void reset_runtime_state() const; diff --git a/csrc/pybind11/engine/engine.hpp b/csrc/pybind11/engine/engine.hpp index b73ce06ae..2be3ee2d0 100644 --- a/csrc/pybind11/engine/engine.hpp +++ b/csrc/pybind11/engine/engine.hpp @@ -9,11 +9,20 @@ namespace py = pybind11; namespace infinilm::engine::distributed { inline void bind_dist_config(py::module &m) { + py::enum_(m, "AllReduceBackend") + .value("AUTO", INFINICCL_ALLREDUCE_BACKEND_AUTO) + .value("NCCL", INFINICCL_ALLREDUCE_BACKEND_NCCL) + .value("CUSTOM", INFINICCL_ALLREDUCE_BACKEND_CUSTOM); + py::class_(m, "DistConfig") .def(py::init<>(), "Default constructor, empty device list") - .def(py::init(), py::arg("tp_size"), + .def(py::init(), + py::arg("tp_size"), + py::arg("allreduce_backend") = INFINICCL_ALLREDUCE_BACKEND_NCCL, "Constructor with tensor parallel size, auto-assigns device IDs 0..tp_size-1") - .def(py::init &>(), py::arg("tp_device_ids"), + .def(py::init &, infinicclAllReduceBackend_t>(), + py::arg("tp_device_ids"), + py::arg("allreduce_backend") = INFINICCL_ALLREDUCE_BACKEND_NCCL, "Constructor with explicit device IDs") .def_readwrite("tp_device_ids", &DistConfig::tp_device_ids, "List of device IDs used in tensor parallelism") @@ -21,6 +30,8 @@ inline void bind_dist_config(py::module &m) { "MoE expert-parallel backend") .def_readwrite("moe_ep_size", &DistConfig::moe_ep_size, "MoE expert-parallel size") + .def_readwrite("allreduce_backend", &DistConfig::allreduce_backend, + "AllReduce backend selector") .def("__repr__", [](const DistConfig &cfg) { return std::string(cfg); }) @@ -123,6 +134,10 @@ inline void bind_infer_engine(py::module &m) { "Run inference on all ranks with arbitrary arguments") .def( "reset_cache", [](InferEngine &self, std::shared_ptr cfg) { self.reset_cache(cfg ? cfg.get() : nullptr); }, py::arg("cache_config") = py::none()) + .def("reset_request_state", &InferEngine::reset_request_state, "Clear per-request relay state without resetting KV cache or recompiling graphs") + .def("sync_last_output", &InferEngine::sync_last_output, "Synchronize with the worker stream that produced the latest output") + .def("copy_last_output_to", &InferEngine::copy_last_output_to, py::arg("dst"), "Copy the latest output into a caller-owned tensor on the correct stream") + .def("close", &InferEngine::close, "Close worker threads and release engine-owned GPU resources") .def("get_kv_cache", &InferEngine::get_kv_cache, "Get per-rank kv cache list") .def("get_cache_config", [](const InferEngine &self) -> std::shared_ptr { auto cfg = self.get_cache_config(); @@ -244,9 +259,24 @@ inline void bind_infer_engine(py::module &m) { .def_readwrite("top_p", &InferEngine::Input::top_p); py::class_(infer_engine, "Output") - .def_readwrite("output_ids", &InferEngine::Output::output_ids, "Sampled token IDs") - .def_readwrite("logits", &InferEngine::Output::logits, "Raw logits tensor") - .def_readwrite("hidden_states", &InferEngine::Output::hidden_states, "Raw hidden states tensor"); + .def_property_readonly( + "output_ids", + [](const InferEngine::Output &output) { + return output.output_ids; + }, + "Sampled token IDs") + .def_property_readonly( + "logits", + [](const InferEngine::Output &output) { + return output.logits; + }, + "Raw logits tensor") + .def_property_readonly( + "hidden_states", + [](const InferEngine::Output &output) { + return output.hidden_states; + }, + "Raw hidden states tensor"); } } // namespace infinilm::engine diff --git a/docs/inference_optimization_summary.md b/docs/inference_optimization_summary.md new file mode 100644 index 000000000..1d6f478cf --- /dev/null +++ b/docs/inference_optimization_summary.md @@ -0,0 +1,519 @@ +# InfiniLM / InfiniCore ???????? + +## 1. ????? + +?????? InfiniLM ? NVIDIA ????????????????????????? + +- ?? decode ???? token ??????? +- ????????????????? workspace ?? +- ?? TP ??????? custom allreduce ??????? +- ?? CUDA Graph ??????? +- ????? chunked prefill + decode priority ???? +- ?????? benchmark / service benchmark ???? + +???????? + +```bash +ssh qinyiqun@172.22.162.61 +docker exec -it qyq_glm bash +cd /root/Infer/InfiniLM +``` + +????? + +```text +/data-aisoft/mechdancer/models/9g_8b_thinking +/data-aisoft/mechdancer/models/Qwen3-30B-A3B +``` + +## 2. Decode ???? + +### ???? + +- ?? SGLang / vLLM ? decode token ???? +- ????????? GPU ? +- ?? GPU event ????? sampled token ???? decode ?? +- ?? decode token ??? host sync / CPU ???? +- CUDA Graph replay ?????????runtime workspace ???? + +### ???? + +- ?? decode ?? sampled token ?????????? +- ?? `DeviceEvent` ??????? decode ????? host +- ?? replay ?? workspace reset ??????? memset ??? host launch +- ?? graph ? sampled output ????????? tensor ????? + +### ???? + +- decode token ?????? CPU ?? +- ??? CUDA Graph replay +- ??? overlap / custom allreduce / service scheduler ???? + +### ?? + +? `9g_8b_thinking` TP=2?????????decode ITL ????? + +```text +Avg ITL: 9.19 ms +Throughput: 108.77 tok/s +``` + +??????????????? + +```text +Avg ITL: 7.9 ms +Throughput: 126 tok/s +``` + +???????????????????? GPU ????? + +### ???? + +```bash +CUDA_VISIBLE_DEVICES=0,1 python examples/bench.py \ + --device=nvidia \ + --model=/data-aisoft/mechdancer/models/9g_8b_thinking \ + --enable-paged-attn \ + --attn=flash-attn \ + --tp=2 \ + --input-len=16,12800 \ + --output-len=1280 \ + --batch-size=1 \ + --enable-graph +``` + +## 3. Workspace Manager ?? + +### ???? + +- ? InfiniLM ??????? workspace +- ? InfiniCore ?? operator workspace ?? +- ? size ???? buffer +- ???????????? workspace +- ???????? kernel / custom op + +### ???? + +- ?? InfiniLM ? `InferenceWorkspaceManager` +- ?? workspace scope / forward guard +- ???? workspace??? workspace??? workspace +- ????????????????? manager ?? +- ? graph capture ??? buffer ?????????? + +### ???? + +- ? size ?? buffer ??? +- ????????? workspace +- ???? warmup ?? workspace pool ??????? +- ????????? malloc/free ??? +- ??? CUDA Graph capture + +### ?? + +- ????????? CPU ?? +- ?? graph capture/replay ????????????? +- ?????? buffer ??? workspace ?????? + +### ???? + +???? graph benchmark ????????? + +```bash +CUDA_VISIBLE_DEVICES=0,1 python examples/bench.py \ + --device=nvidia \ + --model=/data-aisoft/mechdancer/models/9g_8b_thinking \ + --enable-paged-attn \ + --attn=flash-attn \ + --tp=2 \ + --input-len=16,1280 \ + --output-len=128 \ + --batch-size=1 \ + --enable-graph +``` + +## 4. Custom AllReduce ????? + +### ???? + +- ?? vLLM / SGLang custom allreduce ?? +- NVIDIA ???? +- decode ?????? custom allreduce +- prefill ?????? custom allreduce +- ????????????????????? + +### ???? + +- ?? `--allreduce-backend` +- ?? `nccl` / `custom` / `auto` +- ?? InfiniCCL allreduce backend ???? +- ?? graph ? custom allreduce buffer ?????? +- ?? prefill ????? custom allreduce ???? +- ???????prefill ? NCCL?decode ????? custom + +### ???? + +```bash +--allreduce-backend=nccl +--allreduce-backend=custom +--allreduce-backend=auto +``` + +### ?? + +- decode ? allreduce ????? NCCL ???? +- prefill ?? NCCL??? custom ??????? +- ? TP=2/4/8 ????????? + +### ???? + +?????? + +- custom allreduce ?? decode ????? +- ?????????custom decode ???? +- prefill ?? custom allreduce ?????? +- ???????? phase-aware?????? custom + +### ???? + +```bash +# NCCL baseline +CUDA_VISIBLE_DEVICES=0,1 python examples/bench.py \ + --device=nvidia \ + --model=/data-aisoft/mechdancer/models/9g_8b_thinking \ + --enable-paged-attn \ + --attn=flash-attn \ + --tp=2 \ + --input-len=16,12800 \ + --output-len=1280 \ + --batch-size=1 \ + --enable-graph \ + --allreduce-backend=nccl + +# Auto strategy +CUDA_VISIBLE_DEVICES=0,1 python examples/bench.py \ + --device=nvidia \ + --model=/data-aisoft/mechdancer/models/9g_8b_thinking \ + --enable-paged-attn \ + --attn=flash-attn \ + --tp=2 \ + --input-len=16,12800 \ + --output-len=1280 \ + --batch-size=1 \ + --enable-graph \ + --allreduce-backend=auto +``` + +## 5. CPU / GPU ??? NUMA ?? + +### ???? + +- ?? CUDA device PCI bus id ???? +- ???? worker thread ??? NUMA CPU ?? +- ?? NVIDIA CUDA runtime ?? +- ????????? + +### ???? + +- ?? device topology ???? +- ? `RankWorker` ???????? worker thread +- ????????? + +????? + +```text +RankInfo: device=NVIDIA:0, tp_size=2, tp_rank=0 bound worker thread to NUMA node 0 CPUs 0-31,64-95 via cuda-runtime pci=0000:16:00.0 +``` + +### ???? + +- ? GPU ??? worker ???????? GPU NUMA ?? +- ??? NUMA CPU ????????? +- ????????????? provider ?? + +### ?? + +- ??????????TP ??????? +- ?? CPU ???? +- ????????????? + +## 6. Chunked Prefill + Decode Priority + +### ???? + +- ?? vLLM chunked prefill ???? +- ? prefill ?? chunk +- decode ?????? +- ?? mixed batch ? phase ?? +- ???????? decode ITL ??? + +### ???? + +- scheduler ?? chunked prefill ?? +- request ?? prefill chunk ?? +- scheduler output ?? prefill/decode flags +- model runner ?? mixed batch ? decode phase / prefill phase ?? +- ????? benchmark ?? + +????? + +```bash +--enable-chunked-prefill +--max-num-batched-tokens +--prefill-chunk-size +--decode-priority +--max-num-partial-prefills +--max-long-partial-prefills +--long-prefill-token-threshold +--min-prefill-chunk-size +``` + +### ???? + +- ? prefill ????????? decode +- ???????? decode P90/P99 ITL +- ???? prefill ??????????? +- ?? cache / scheduler ???? + +### ??????? + +```bash +python examples/chunked_prefill_service_bench.py +``` + +?????? + +```bash +CUDA_VISIBLE_DEVICES=0,1 python examples/chunked_prefill_service_bench.py \ + --device=nvidia \ + --model=/data-aisoft/mechdancer/models/9g_8b_thinking \ + --tp=2 \ + --attn=flash-attn \ + --allreduce-backend=auto \ + --enable-chunked-prefill \ + --max-num-batched-tokens=2048 \ + --prefill-chunk-size=2048 \ + --decode-priority \ + --num-short-requests=4 \ + --short-input-len=128 \ + --short-output-len=256 \ + --long-input-len=12800 \ + --long-output-len=4 \ + --insert-after-tokens=8 \ + --output-json=/data/shared/qinyiqun/chunked_prefill_enabled.json +``` + +Baseline? + +```bash +CUDA_VISIBLE_DEVICES=0,1 python examples/chunked_prefill_service_bench.py \ + --device=nvidia \ + --model=/data-aisoft/mechdancer/models/9g_8b_thinking \ + --tp=2 \ + --attn=flash-attn \ + --allreduce-backend=auto \ + --num-short-requests=4 \ + --short-input-len=128 \ + --short-output-len=256 \ + --long-input-len=12800 \ + --long-output-len=4 \ + --insert-after-tokens=8 \ + --output-json=/data/shared/qinyiqun/chunked_prefill_baseline.json +``` + +??????? + +```text +summary.short_itl_after_long_insert.p50_ms +summary.short_itl_after_long_insert.p90_ms +summary.short_itl_after_long_insert.p99_ms +summary.short_itl_after_long_insert.max_ms +summary.long_ttft_ms +cache_stats.scheduler.mixed_steps +cache_stats.scheduler.execution_split_mixed_steps +``` + +### ?????? + +no-graph ??? smoke ???? + +```text +SERVICE_NOGRAPH_SMOKE_RC=0 +chunked_prefill.enabled: true +decode_priority: true +scheduler.mixed_steps: 1 +scheduler.execution_split_mixed_steps: 1 +``` + +graph ???????????????? + +```text +Check Failed: x_desc->dtype() == dtype +src/infiniop/ops/rearrange/nvidia/rearrange_nvidia.cu:31 +``` + +## 7. ?????????? + +### ???? + +- ? rank ????? +- ????? worker ?? +- ??? join +- ?? TP ? worker ?? close ?? graph / communication / stream sync ?????? + +### ???? + +- `RankWorker::close()` ???? + - `request_close()` + - `join()` +- `InferEngine::close()` ??? + - first pass: all workers `request_close()` + - second pass: all workers `join()` +- InfiniCore `Memory::~Memory()` ? allocator shutdown ??????? + +### ???? + +- TP=2 graph benchmark ????????? +- ?? `Pointer not allocated by this allocator` ???? abort +- ?? `bench.py` ?? `total_time` ?????? + +### ???? + +```bash +CUDA_VISIBLE_DEVICES=1,2 python examples/bench.py \ + --device=nvidia \ + --model=/data-aisoft/mechdancer/models/9g_8b_thinking \ + --enable-paged-attn \ + --attn=flash-attn \ + --tp=2 \ + --input-len=16,1280 \ + --output-len=128 \ + --batch-size=1 \ + --enable-graph +``` + +??? + +```text +BENCH_RC=0 +``` + +## 8. ?? / ?????? + +### bench.py + +```bash +--enable-graph +--allreduce-backend={nccl,custom,auto} +--skip-legacy-moe +--moe-ep-backend={auto,disabled,local_allreduce,allgather_reducescatter} +--ep +--tp +--dp +``` + +### chunked_prefill_service_bench.py + +```bash +--enable-chunked-prefill +--max-num-batched-tokens +--prefill-chunk-size +--decode-priority +--no-decode-priority +--max-num-partial-prefills +--max-long-partial-prefills +--long-prefill-token-threshold +--min-prefill-chunk-size +--skip-legacy-moe +--moe-ep-backend=auto +--ep +--dp +--output-json +``` + +## 9. ?????? + +### ????? / ITL + +| ?? | ?? | ???? | +|---|---|---| +| decode ?? | `examples/bench.py` | `Decode Avg ITL` | +| graph replay | `examples/bench.py --enable-graph` | ITL / ?????? | +| custom allreduce | `--allreduce-backend=nccl/custom/auto` | decode throughput | +| ???? prefill | `input-len=5120/12800` | TTFT / prefill throughput | + +### ????? + +| ?? | ?? | ???? | +|---|---|---| +| chunked prefill | `chunked_prefill_service_bench.py` | short ITL p99 | +| decode priority | `--decode-priority` | short ITL after long insert | +| long prefill ?? | ?? long request | short p90/p99/max | +| scheduler ??? | `cache_stats.scheduler` | mixed steps / split phases | + +## 10. ????? / ???? + +### ?????? + +1. `service + enable_graph` ? `rearrange_nvidia.cu` dtype mismatch +2. chunked prefill ??? A/B ???? +3. ??? benchmark with long-context workloads +4. custom allreduce ? TP=4/8 ?????? +5. allreduce ? GEMM overlap ?????? profiling +6. nsys ????????? decode token ? GPU ?? + +### ????? + +?????? + +1. ?? `chunked_prefill_service_bench.py --enable-graph` ? dtype mismatch +2. ? `9g_8b_thinking` ? chunked prefill A/B +3. ? nsys profile ?? decode gap ? prefill/decode interleave +4. ???????? benchmark ???????? commit ?? + +## 11. ???? + +### InfiniCore + +```bash +cd /root/Infer/InfiniCore + +xmake f \ + --nv-gpu=y \ + --ccl=y \ + --cuda=$CUDA_HOME \ + --aten=y \ + --flash-attn=/root/Infer/InfiniCore/third_party/flash-attention/ \ + --graph=y \ + --cuda_arch=sm_80 \ + -cv + +xmake build +xmake install +xmake build _infinicore +xmake install _infinicore +python -m pip install -e file:///root/Infer/InfiniCore +``` + +### InfiniLM + +```bash +cd /root/Infer/InfiniLM +xmake build _infinilm +python -m pip install -e file:///root/Infer/InfiniLM +``` + +## 12. ?? + +??????? InfiniLM ????? benchmark ?????????????graph?????????????????? + +?????????? + +- decode ???? +- workspace manager +- phase-aware custom allreduce ?? +- NUMA ???? +- chunked prefill + decode priority scheduler +- ??? benchmark ?? +- ? rank ?????? + +?????????????? graph ??? chunked prefill A/B ??????????? 70B ??????????? diff --git a/examples/bench.py b/examples/bench.py index bd424e836..27e8d9b26 100644 --- a/examples/bench.py +++ b/examples/bench.py @@ -6,6 +6,7 @@ import infinicore import numpy as np +from PIL import Image from infinilm.base_config import BaseConfig from infinilm.cache import PagedKVCacheConfig, StaticKVCacheConfig from infinilm.distributed import DistConfig @@ -14,6 +15,7 @@ from infinilm.llm.sampling_params import SamplingParams from infinilm.modeling_utils import load_model_state_dict_by_file from infinilm.moe_config import configure_moe_ep_backend +from infinilm.multimodal.multimodal import load_video from infinilm.processors import AutoInfinilmProcessor from tqdm import tqdm @@ -163,13 +165,66 @@ def get_test_cases( return case_dict -prompt_path = ( - "examples/bench_prompt.md" - if os.path.isfile("examples/bench_prompt.md") - else "InfiniLM/examples/bench_prompt.md" -) -with open(prompt_path, "r") as f: - prompt = f.read() +DEFAULT_PROMPT = "How are you" + + +def load_bench_prompt(prompt_text: str, prompt_file: str | None) -> str: + if prompt_file is not None: + with open(os.path.expanduser(prompt_file), "r") as f: + return f.read() + if prompt_text != DEFAULT_PROMPT: + return prompt_text + prompt_path = ( + "examples/bench_prompt.md" + if os.path.isfile("examples/bench_prompt.md") + else "InfiniLM/examples/bench_prompt.md" + ) + with open(prompt_path, "r") as f: + return f.read() + + +def _build_vl_content(prompt_text: str, image_path: str | None, video_path: str | None): + content = [] + if image_path is not None: + content.append({"type": "image_url", "image_url": {"url": image_path}}) + if video_path is not None: + content.append({"type": "video_url", "video_url": {"url": video_path}}) + content.append({"type": "text", "text": prompt_text}) + return content + + +def _load_vl_media(image_path: str | None, video_path: str | None): + images = None + videos = None + if image_path is not None: + images = [Image.open(image_path).convert("RGB")] + if video_path is not None: + videos = [load_video(video_path)] + return images, videos + + +def get_vl_prompt_length( + model_path: str, image_path: str | None, video_path: str | None, prompt_text: str +) -> int: + processor = AutoInfinilmProcessor.from_pretrained(model_path) + input_content = processor.apply_chat_template( + conversation=[ + { + "role": "user", + "content": _build_vl_content(prompt_text, image_path, video_path), + } + ], + add_generation_prompt=True, + tokenize=False, + ) + images, videos = _load_vl_media(image_path, video_path) + processed_inputs = processor( + input_content, + images=images, + videos=videos, + return_tensors="pt", + ) + return int(processed_inputs["input_ids"].shape[-1]) def repeat_prompt(input_ids: list[int], target_length: int): @@ -192,11 +247,17 @@ def __init__( skip_load=False, cache_config=None, enable_graph=False, + skip_legacy_moe=False, attn_backend="default", use_mla=False, weight_load_mode="async", + allreduce_backend="nccl", moe_ep_backend="disabled", moe_ep_size=1, + prompt_text=DEFAULT_PROMPT, + repeat_prompt_to_length=True, + image_path=None, + video_path=None, ) -> None: model_path = os.path.expanduser(model_path) self.draft_model_path = draft_model_path @@ -231,6 +292,7 @@ def __init__( device=infini_device, distributed_config=DistConfig( tp, + allreduce_backend=allreduce_backend, moe_ep_backend=moe_ep_backend, moe_ep_size=moe_ep_size, ), @@ -240,6 +302,7 @@ def __init__( kv_cache_dtype=cfg.kv_cache_dtype, use_mla=use_mla, weight_load_mode=weight_load_mode, + skip_legacy_moe=skip_legacy_moe, ) # ---------------------------------------------------------------------------- # @@ -257,17 +320,64 @@ def __init__( # ---------------------------------------------------------------------------- # # token编码 # ---------------------------------------------------------------------------- # - input_content = self.processor.apply_chat_template( - conversation=[{"role": "user", "content": prompt}], - add_generation_prompt=True, - tokenize=False, - ) + self.pixel_values = None + self.image_bound = None + self.tgt_sizes = None + self.prefill_position_ids = None - input_ids_list = [ - self.tokenizer.encode( + if image_path is not None or video_path is not None: + input_content = self.processor.apply_chat_template( + conversation=[ + { + "role": "user", + "content": _build_vl_content( + prompt_text, image_path, video_path + ), + } + ], + add_generation_prompt=True, + tokenize=False, + ) + images, videos = _load_vl_media(image_path, video_path) + processed_inputs = self.processor( input_content, + images=images, + videos=videos, + return_tensors="infini", + ) + input_ids_list = [processed_inputs["input_ids"].to_numpy()[0].tolist()] + self.pixel_values = processed_inputs.get("pixel_values") + self.image_bound = processed_inputs.get("image_bound") + self.tgt_sizes = processed_inputs.get("tgt_sizes") + self.prefill_position_ids = processed_inputs.get("position_ids") + print( + "VL input:", + { + "input_ids": processed_inputs["input_ids"].shape, + "position_ids": None + if self.prefill_position_ids is None + else self.prefill_position_ids.shape, + "pixel_values": None + if self.pixel_values is None + else self.pixel_values.shape, + "image_bound": None + if self.image_bound is None + else self.image_bound.shape, + "tgt_sizes": None if self.tgt_sizes is None else self.tgt_sizes.shape, + }, + ) + else: + input_content = self.processor.apply_chat_template( + conversation=[{"role": "user", "content": prompt_text}], + add_generation_prompt=True, + tokenize=False, ) - ] + + input_ids_list = [ + self.tokenizer.encode( + input_content, + ) + ] self.model = model self.input_ids_list = input_ids_list @@ -281,6 +391,7 @@ def __init__( self.use_mla = use_mla self.weight_load_mode = weight_load_mode self.skip_load = skip_load + self.repeat_prompt_to_length = repeat_prompt_to_length def run( self, @@ -291,7 +402,14 @@ def run( top_p=1.0, temperature=1.0, ): - input_ids = repeat_prompt(self.input_ids_list[0], target_length=input_len) + if self.pixel_values is not None: + if batch_size != 1: + raise ValueError("VL benchmark currently supports batch_size=1") + input_ids = self.input_ids_list[0] + elif self.repeat_prompt_to_length: + input_ids = repeat_prompt(self.input_ids_list[0], target_length=input_len) + else: + input_ids = self.input_ids_list[0][:input_len] input_ids_list = [input_ids] * batch_size # ---------------------------------------------------------------------------- # @@ -350,6 +468,10 @@ def run( temperature=temperature, stop_on_eos=False, ), + pixel_values=self.pixel_values, + image_bound=self.image_bound, + tgt_sizes=self.tgt_sizes, + prefill_position_ids=self.prefill_position_ids, _measure_and_log_time=True, ) t2 = time.time() @@ -358,7 +480,15 @@ def run( [output_id.to_numpy()[0] for output_id in output_ids] ) if not skip_load: - print(self.tokenizer.decode(numpy_output_ids, skip_special_tokens=True)) + decoded_ids = numpy_output_ids.reshape(-1).tolist() + decoded_text = self.tokenizer.decode(decoded_ids, skip_special_tokens=True) + print(decoded_text) + if decoded_text == "": + print(f"[decode-empty] output_ids={decoded_ids}") + print( + "[decode-empty] with_special_tokens=" + f"{self.tokenizer.decode(decoded_ids, skip_special_tokens=False)}" + ) print( f"total_time: {round((t2 - t1) * 1000, 2)} ms", @@ -393,12 +523,18 @@ def run( enable_paged_attn = cfg.enable_paged_attn enable_graph = cfg.enable_graph attn_backend = cfg.attn + prompt_text = load_bench_prompt(cfg.prompt, cfg.prompt_file) if isinstance(batch_size, int): batch_size = [batch_size] if isinstance(input_len, int): input_len = [input_len] + if cfg.image is not None or cfg.video is not None: + vl_prompt_len = get_vl_prompt_length( + model_path, cfg.image, cfg.video, prompt_text + ) + input_len = [max(length, vl_prompt_len) for length in input_len] if isinstance(output_len, int): output_len = [output_len] @@ -437,11 +573,17 @@ def run( skip_load=skip_load, cache_config=cache_config, enable_graph=enable_graph, + skip_legacy_moe=cfg.skip_legacy_moe, attn_backend=attn_backend, use_mla=cfg.use_mla, weight_load_mode=cfg.weight_load_mode, + allreduce_backend=cfg.allreduce_backend, moe_ep_backend=moe_ep_backend, moe_ep_size=ep, + prompt_text=prompt_text, + repeat_prompt_to_length=cfg.repeat_prompt, + image_path=cfg.image, + video_path=cfg.video, ) # ---------------------------------------------------------------------------- # @@ -506,28 +648,31 @@ def run( # Warmup done # ---------------------------------------------------------------------------- # - for idx, case in tqdm(cases_dict.items(), desc="Processing cases"): - tqdm.write(f"\033[92mProcessing : {case}\033[0m") + try: + for idx, case in tqdm(cases_dict.items(), desc="Processing cases"): + tqdm.write(f"\033[92mProcessing : {case}\033[0m") - batch_size = case["batch_size"] - input_len = case["input_len"] - output_len = case["output_len"] + batch_size = case["batch_size"] + input_len = case["input_len"] + output_len = case["output_len"] - if not enable_paged_attn: - # reset cache if static kvcache is used - initial_capacity = input_len + output_len - test.model.reset_cache( - StaticKVCacheConfig( - max_batch_size=batch_size, max_cache_len=initial_capacity + if not enable_paged_attn: + # reset cache if static kvcache is used + initial_capacity = input_len + output_len + test.model.reset_cache( + StaticKVCacheConfig( + max_batch_size=batch_size, max_cache_len=initial_capacity + ) ) - ) - # run test one case - test.run( - batch_size=batch_size, - input_len=input_len, - output_len=output_len, - top_k=cfg.top_k, - top_p=cfg.top_p, - temperature=cfg.temperature, - ) + # run test one case + test.run( + batch_size=batch_size, + input_len=input_len, + output_len=output_len, + top_k=cfg.top_k, + top_p=cfg.top_p, + temperature=cfg.temperature, + ) + finally: + test.model.close() diff --git a/examples/chunked_prefill_service_bench.py b/examples/chunked_prefill_service_bench.py new file mode 100644 index 000000000..54d9e2591 --- /dev/null +++ b/examples/chunked_prefill_service_bench.py @@ -0,0 +1,434 @@ +"""Service-style benchmark for chunked prefill scheduling. + +This benchmark exercises the LLMEngine scheduler path, not the lower-level +InferEngine.generate path used by examples/bench.py. It launches several short +decode-heavy requests, inserts one long-prefill request after the short requests +have started decoding, and reports short-request tail ITL plus long-request TTFT. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import time +from dataclasses import dataclass, field +from typing import Any, Optional + +from infinilm.llm import AsyncLLMEngine, SamplingParams +from infinilm.moe_config import configure_moe_ep_backend + + +DEVICE_STR_MAP = { + "cpu": "cpu", + "nvidia": "cuda", + "qy": "cuda", + "cuda": "cuda", + "cambricon": "mlu", + "ascend": "npu", + "metax": "cuda", + "moore": "musa", + "iluvatar": "cuda", + "kunlun": "kunlun", + "hygon": "cuda", + "ali": "cuda", +} + + +@dataclass +class RequestMetrics: + request_id: str + kind: str + prompt_len: int + output_len: int + submit_time: float + first_token_time: Optional[float] = None + finish_time: Optional[float] = None + token_times: list[float] = field(default_factory=list) + token_ids: list[int] = field(default_factory=list) + finish_reason: Optional[str] = None + error: Optional[str] = None + + def to_dict(self, origin: float) -> dict[str, Any]: + ttft_ms = None + if self.first_token_time is not None: + ttft_ms = (self.first_token_time - self.submit_time) * 1000.0 + + total_ms = None + if self.finish_time is not None: + total_ms = (self.finish_time - self.submit_time) * 1000.0 + + itl_ms = [] + for prev, cur in zip(self.token_times, self.token_times[1:]): + itl_ms.append((cur - prev) * 1000.0) + + return { + "request_id": self.request_id, + "kind": self.kind, + "prompt_len": self.prompt_len, + "output_len": self.output_len, + "submit_s": self.submit_time - origin, + "first_token_s": None + if self.first_token_time is None + else self.first_token_time - origin, + "finish_s": None + if self.finish_time is None + else self.finish_time - origin, + "num_tokens": len(self.token_times), + "ttft_ms": ttft_ms, + "total_ms": total_ms, + "itl_ms": itl_ms, + "finish_reason": self.finish_reason, + "error": self.error, + } + + +def percentile(values: list[float], q: float) -> Optional[float]: + if not values: + return None + ordered = sorted(values) + if len(ordered) == 1: + return ordered[0] + pos = (len(ordered) - 1) * q + lo = int(pos) + hi = min(lo + 1, len(ordered) - 1) + frac = pos - lo + return ordered[lo] * (1.0 - frac) + ordered[hi] * frac + + +def summarize(values: list[float]) -> dict[str, Optional[float]]: + return { + "count": len(values), + "avg_ms": sum(values) / len(values) if values else None, + "p50_ms": percentile(values, 0.50), + "p90_ms": percentile(values, 0.90), + "p99_ms": percentile(values, 0.99), + "max_ms": max(values) if values else None, + } + + +def normalize_device(device: str) -> str: + return DEVICE_STR_MAP.get(device.lower(), device) + + +def repeat_tokens(tokens: list[int], target_len: int) -> list[int]: + if target_len <= 0: + return [] + if not tokens: + raise ValueError("Cannot repeat an empty prompt token list") + repeat = (target_len + len(tokens) - 1) // len(tokens) + return (tokens * repeat)[:target_len] + + +def build_prompt_tokens(engine: AsyncLLMEngine, prompt: str, target_len: int) -> list[int]: + tokens = engine.engine.tokenize(prompt) + return repeat_tokens(tokens, target_len) + + +async def collect_request( + engine: AsyncLLMEngine, + request, + metrics: RequestMetrics, + stream_timeout_s: float, + request_timeout_s: float, +) -> RequestMetrics: + try: + async for token_output in engine.stream_request( + request, + timeout=stream_timeout_s, + request_timeout=request_timeout_s, + ): + now = time.perf_counter() + if metrics.first_token_time is None: + metrics.first_token_time = now + if token_output.token_id >= 0: + metrics.token_times.append(now) + metrics.token_ids.append(token_output.token_id) + if token_output.finished: + metrics.finish_reason = ( + token_output.finish_reason.value + if hasattr(token_output.finish_reason, "value") + else str(token_output.finish_reason) + ) + break + metrics.finish_time = time.perf_counter() + except Exception as exc: # noqa: BLE001 - benchmark should report failures. + metrics.error = repr(exc) + metrics.finish_time = time.perf_counter() + return metrics + + +async def wait_for_short_progress( + short_metrics: list[RequestMetrics], + tokens: int, + timeout_s: float, +) -> bool: + deadline = time.perf_counter() + timeout_s + while time.perf_counter() < deadline: + if all(len(metric.token_times) >= tokens for metric in short_metrics): + return True + await asyncio.sleep(0.001) + return False + + +def itls_after_insert( + metrics: list[RequestMetrics], insert_time: float +) -> list[float]: + values = [] + for metric in metrics: + for prev, cur in zip(metric.token_times, metric.token_times[1:]): + if cur >= insert_time: + values.append((cur - prev) * 1000.0) + return values + + +def all_itls(metrics: list[RequestMetrics]) -> list[float]: + values = [] + for metric in metrics: + for prev, cur in zip(metric.token_times, metric.token_times[1:]): + values.append((cur - prev) * 1000.0) + return values + + +def add_engine_args(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--model", required=True) + parser.add_argument("--device", default="nvidia") + parser.add_argument("--dtype", default="bfloat16") + parser.add_argument("--tp", type=int, default=1) + parser.add_argument("--dp", type=int, default=1) + parser.add_argument("--ep", type=int, default=None) + parser.add_argument("--moe-ep-backend", default="auto") + parser.add_argument("--moe-ep-size", type=int, default=None) + parser.add_argument("--skip-legacy-moe", action="store_true") + parser.add_argument("--allreduce-backend", default="nccl", choices=["nccl", "auto", "custom"]) + parser.add_argument("--attn", default="flash-attn", choices=["default", "paged-attn", "flash-attn"]) + parser.add_argument("--enable-graph", action="store_true") + parser.add_argument("--use-mla", action="store_true") + parser.add_argument("--weight-load", dest="weight_load_mode", default="async", choices=["async", "sync"]) + parser.add_argument("--max-batch-size", type=int, default=8) + parser.add_argument("--num-blocks", type=int, default=2048) + parser.add_argument("--block-size", type=int, default=256) + parser.add_argument("--max-cache-len", type=int, default=4096) + + +def add_chunked_args(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--enable-chunked-prefill", action="store_true") + parser.add_argument("--max-num-batched-tokens", type=int, default=None) + parser.add_argument("--prefill-chunk-size", type=int, default=None) + group = parser.add_mutually_exclusive_group() + group.add_argument("--decode-priority", dest="decode_priority", action="store_true", default=None) + group.add_argument("--no-decode-priority", dest="decode_priority", action="store_false") + parser.add_argument("--max-num-partial-prefills", type=int, default=1) + parser.add_argument("--max-long-partial-prefills", type=int, default=1) + parser.add_argument("--long-prefill-token-threshold", type=int, default=None) + parser.add_argument("--min-prefill-chunk-size", type=int, default=None) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + add_engine_args(parser) + add_chunked_args(parser) + parser.add_argument("--num-short-requests", type=int, default=4) + parser.add_argument("--short-input-len", type=int, default=128) + parser.add_argument("--short-output-len", type=int, default=64) + parser.add_argument("--long-input-len", type=int, default=12800) + parser.add_argument("--long-output-len", type=int, default=4) + parser.add_argument("--insert-after-tokens", type=int, default=8) + parser.add_argument("--prompt-file", default="examples/bench_prompt.md") + parser.add_argument("--warmup", action="store_true") + parser.add_argument("--stream-timeout-s", type=float, default=120.0) + parser.add_argument("--request-timeout-s", type=float, default=1200.0) + parser.add_argument("--progress-timeout-s", type=float, default=120.0) + parser.add_argument("--output-json", default=None) + return parser.parse_args() + + +def load_prompt(path: str) -> str: + if path and os.path.exists(path): + with open(path, "r", encoding="utf-8") as f: + return f.read() + return "Explain the history and geography of Mount Tai in detail." + + +async def run_once(args: argparse.Namespace) -> dict[str, Any]: + required_cache_len = max( + args.short_input_len + args.short_output_len, + args.long_input_len + args.long_output_len, + ) + if args.max_cache_len < required_cache_len: + args.max_cache_len = required_cache_len + + requested_ep = args.ep if args.ep is not None else args.moe_ep_size + moe_ep_backend, moe_ep_size = configure_moe_ep_backend( + args.tp, args.dp, requested_ep, args.moe_ep_backend, args.model + ) + args.moe_ep_backend = moe_ep_backend + args.moe_ep_size = moe_ep_size + print(f"MoE EP backend: {moe_ep_backend} TP={args.tp} DP={args.dp} EP={moe_ep_size}") + + engine = AsyncLLMEngine( + model_path=args.model, + device=normalize_device(args.device), + dtype=args.dtype, + tensor_parallel_size=args.tp, + moe_ep_backend=moe_ep_backend, + moe_ep_size=moe_ep_size, + allreduce_backend=args.allreduce_backend, + cache_type="paged", + max_batch_size=args.max_batch_size, + max_tokens=max(args.short_output_len, args.long_output_len), + num_blocks=args.num_blocks, + block_size=args.block_size, + max_cache_len=args.max_cache_len, + temperature=1.0, + top_p=1.0, + top_k=1, + enable_graph=args.enable_graph, + attn_backend=args.attn, + use_mla=args.use_mla, + weight_load_mode=args.weight_load_mode, + skip_legacy_moe=args.skip_legacy_moe, + max_num_batched_tokens=args.max_num_batched_tokens, + enable_chunked_prefill=args.enable_chunked_prefill, + prefill_chunk_size=args.prefill_chunk_size, + decode_priority=args.decode_priority, + max_num_partial_prefills=args.max_num_partial_prefills, + max_long_partial_prefills=args.max_long_partial_prefills, + long_prefill_token_threshold=args.long_prefill_token_threshold, + min_prefill_chunk_size=args.min_prefill_chunk_size, + ) + + origin = time.perf_counter() + tasks: list[asyncio.Task] = [] + short_metrics: list[RequestMetrics] = [] + long_metric: Optional[RequestMetrics] = None + insert_time: Optional[float] = None + progress_ready = False + + try: + engine.start() + prompt = load_prompt(args.prompt_file) + short_tokens = build_prompt_tokens(engine, prompt, args.short_input_len) + long_tokens = build_prompt_tokens(engine, prompt, args.long_input_len) + + if args.warmup: + warmup_req = engine.add_request( + messages=None, + prompt_token_ids=short_tokens, + sampling_params=SamplingParams(max_tokens=2, top_k=1, top_p=1.0, ignore_eos=True), + request_id="warmup", + ) + warmup_metrics = RequestMetrics("warmup", "warmup", len(short_tokens), 2, time.perf_counter()) + await collect_request(engine, warmup_req, warmup_metrics, args.stream_timeout_s, args.request_timeout_s) + + for idx in range(args.num_short_requests): + req_id = f"short-{idx}" + submit = time.perf_counter() + req = engine.add_request( + messages=None, + prompt_token_ids=short_tokens, + sampling_params=SamplingParams( + max_tokens=args.short_output_len, + top_k=1, + top_p=1.0, + temperature=1.0, + ignore_eos=True, + ), + request_id=req_id, + ) + metric = RequestMetrics(req_id, "short", len(short_tokens), args.short_output_len, submit) + short_metrics.append(metric) + tasks.append( + asyncio.create_task( + collect_request( + engine, + req, + metric, + args.stream_timeout_s, + args.request_timeout_s, + ) + ) + ) + + progress_ready = await wait_for_short_progress( + short_metrics, + args.insert_after_tokens, + args.progress_timeout_s, + ) + + insert_time = time.perf_counter() + long_req = engine.add_request( + messages=None, + prompt_token_ids=long_tokens, + sampling_params=SamplingParams( + max_tokens=args.long_output_len, + top_k=1, + top_p=1.0, + temperature=1.0, + ignore_eos=True, + ), + request_id="long", + ) + long_metric = RequestMetrics( + "long", + "long", + len(long_tokens), + args.long_output_len, + insert_time, + ) + tasks.append( + asyncio.create_task( + collect_request( + engine, + long_req, + long_metric, + args.stream_timeout_s, + args.request_timeout_s, + ) + ) + ) + + await asyncio.wait_for(asyncio.gather(*tasks), timeout=args.request_timeout_s) + finally: + engine.stop() + + short_after = itls_after_insert(short_metrics, insert_time or origin) + short_all = all_itls(short_metrics) + long_ttft = None + long_total = None + if long_metric and long_metric.first_token_time: + long_ttft = (long_metric.first_token_time - long_metric.submit_time) * 1000.0 + if long_metric and long_metric.finish_time: + long_total = (long_metric.finish_time - long_metric.submit_time) * 1000.0 + + return { + "config": vars(args), + "progress_ready_before_insert": progress_ready, + "insert_time_s": None if insert_time is None else insert_time - origin, + "summary": { + "short_itl_all": summarize(short_all), + "short_itl_after_long_insert": summarize(short_after), + "long_ttft_ms": long_ttft, + "long_total_ms": long_total, + }, + "cache_stats": engine.get_cache_stats(), + "requests": [m.to_dict(origin) for m in short_metrics] + + ([] if long_metric is None else [long_metric.to_dict(origin)]), + } + + +def main() -> None: + args = parse_args() + result = asyncio.run(run_once(args)) + text = json.dumps(result, indent=2, ensure_ascii=False) + print(text) + if args.output_json: + os.makedirs(os.path.dirname(os.path.abspath(args.output_json)), exist_ok=True) + with open(args.output_json, "w", encoding="utf-8") as f: + f.write(text) + f.write("\n") + + +if __name__ == "__main__": + main() diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index 3b8f6533d..f64d5aaf2 100644 --- a/python/infinilm/base_config.py +++ b/python/infinilm/base_config.py @@ -67,6 +67,7 @@ def __init__(self): self.ep = self.args.ep self.moe_ep_backend = self.args.moe_ep_backend self.skip_legacy_moe = self.args.skip_legacy_moe + self.allreduce_backend = self.args.allreduce_backend self.attn = self.args.attn self.enable_graph = self.args.enable_graph @@ -81,10 +82,20 @@ def __init__(self): self.batch_size = self.args.batch_size self.max_batch_size = self.args.max_batch_size + self.max_num_batched_tokens = self.args.max_num_batched_tokens + self.enable_chunked_prefill = self.args.enable_chunked_prefill + self.prefill_chunk_size = self.args.prefill_chunk_size + self.decode_priority = self.args.decode_priority + self.max_num_partial_prefills = self.args.max_num_partial_prefills + self.max_long_partial_prefills = self.args.max_long_partial_prefills + self.long_prefill_token_threshold = self.args.long_prefill_token_threshold + self.min_prefill_chunk_size = self.args.min_prefill_chunk_size self.input_len = self.args.input_len self.output_len = self.args.output_len self.max_new_tokens = self.args.max_new_tokens self.prompt = self.args.prompt + self.prompt_file = self.args.prompt_file + self.repeat_prompt = self.args.repeat_prompt self.top_k = self.args.top_k self.top_p = self.args.top_p self.temperature = self.args.temperature @@ -214,6 +225,13 @@ def _add_common_args(self): action="store_true", help="use the new fused MoE implementation instead of the legacy Qwen3 MoE MLP", ) + self.parser.add_argument( + "--allreduce-backend", + type=str, + default="nccl", + choices=["nccl", "auto", "custom"], + help="Tensor-parallel AllReduce backend selector.", + ) # --- Infer backend optimization --- self.parser.add_argument( @@ -269,6 +287,59 @@ def _add_common_args(self): default=8, help="maximum batch size for server", ) + self.parser.add_argument( + "--max-num-batched-tokens", + type=int, + default=None, + help="maximum scheduled tokens per step for paged-cache scheduler", + ) + self.parser.add_argument( + "--enable-chunked-prefill", + action="store_true", + help="split long prefills across scheduler steps", + ) + self.parser.add_argument( + "--prefill-chunk-size", + type=int, + default=None, + help="maximum tokens per prefill chunk when chunked prefill is enabled", + ) + self.parser.add_argument( + "--decode-priority", + action="store_true", + default=None, + help=( + "schedule ready decode batches before waiting prefill chunks; " + "defaults to enabled when chunked prefill is enabled" + ), + ) + self.parser.add_argument( + "--max-num-partial-prefills", + type=int, + default=1, + help="maximum partial prefill chunks scheduled in one step", + ) + self.parser.add_argument( + "--max-long-partial-prefills", + type=int, + default=1, + help="maximum long partial prefill chunks scheduled in one step", + ) + self.parser.add_argument( + "--long-prefill-token-threshold", + type=int, + default=None, + help=( + "prompt length threshold for long partial prefill limiting; " + "defaults to the resolved prefill chunk size" + ), + ) + self.parser.add_argument( + "--min-prefill-chunk-size", + type=int, + default=None, + help="minimum non-final prefill chunk size when mixed with decode", + ) self.parser.add_argument( "--input-len", type=parse_list, default=10, help="input sequence length" ) @@ -284,6 +355,18 @@ def _add_common_args(self): self.parser.add_argument( "--prompt", type=str, default="How are you", help="default prompt text" ) + self.parser.add_argument( + "--prompt-file", + type=str, + default=None, + help="read prompt text from a file instead of examples/bench_prompt.md", + ) + self.parser.add_argument( + "--no-repeat-prompt", + action="store_false", + dest="repeat_prompt", + help="truncate the prompt to input-len instead of repeating it", + ) self.parser.add_argument("--top-k", type=int, default=1) self.parser.add_argument("--top-p", type=float, default=1.0) self.parser.add_argument("--temperature", type=float, default=1.0) diff --git a/python/infinilm/config/engine_config.py b/python/infinilm/config/engine_config.py index 1bb1f733e..81a98706f 100644 --- a/python/infinilm/config/engine_config.py +++ b/python/infinilm/config/engine_config.py @@ -17,11 +17,20 @@ class EngineConfig: tensor_parallel_size: Number of devices for tensor parallelism. moe_ep_backend: MoE expert-parallel backend. moe_ep_size: MoE expert-parallel size. + allreduce_backend: Tensor-parallel AllReduce backend selector. cache_type: Cache type ('paged' or 'static'). max_batch_size: Maximum batch size for inference (only for paged cache). max_tokens: Default maximum tokens to generate. num_blocks: Number of KV cache blocks (only for paged cache). block_size: Size of each KV cache block (only for paged cache). + max_num_batched_tokens: Token budget per paged-cache scheduler step. + enable_chunked_prefill: Whether to split long prefills across scheduler steps. + prefill_chunk_size: Maximum tokens per chunked prefill step. + decode_priority: Whether decode requests are scheduled before prefill chunks. + max_num_partial_prefills: Maximum partial prefill chunks in one scheduler step. + max_long_partial_prefills: Maximum long partial prefill chunks in one scheduler step. + long_prefill_token_threshold: Prompt length threshold for long partial prefills. + min_prefill_chunk_size: Minimum non-final prefill chunk when decode is batched. max_cache_len: Maximum sequence length (only for static cache). temperature: Default sampling temperature. top_p: Default top-p sampling parameter. @@ -42,11 +51,20 @@ class EngineConfig: tensor_parallel_size: int = 1 moe_ep_backend: str = "disabled" moe_ep_size: int = 1 + allreduce_backend: str = "nccl" cache_type: str = "paged" # "paged" or "static" max_batch_size: int = 16 max_tokens: int = 4096 num_blocks: int = 512 block_size: int = 256 + max_num_batched_tokens: Optional[int] = None + enable_chunked_prefill: bool = False + prefill_chunk_size: Optional[int] = None + decode_priority: Optional[bool] = None + max_num_partial_prefills: int = 1 + max_long_partial_prefills: int = 1 + long_prefill_token_threshold: Optional[int] = None + min_prefill_chunk_size: Optional[int] = None max_cache_len: int = 4096 temperature: float = 1.0 top_p: float = 0.8 @@ -65,6 +83,26 @@ def __post_init__(self) -> None: if self.weight_load_mode not in {"async", "sync"}: raise ValueError("weight_load_mode must be either 'async' or 'sync'") + self.allreduce_backend = self.allreduce_backend.lower().replace("-", "_") + if self.allreduce_backend not in {"nccl", "auto", "custom"}: + raise ValueError( + "allreduce_backend must be one of 'nccl', 'auto', or 'custom'" + ) + if self.max_num_partial_prefills < 1: + raise ValueError("max_num_partial_prefills must be >= 1") + if self.max_long_partial_prefills < 1: + raise ValueError("max_long_partial_prefills must be >= 1") + if self.max_long_partial_prefills > self.max_num_partial_prefills: + raise ValueError( + "max_long_partial_prefills must be <= max_num_partial_prefills" + ) + if ( + self.long_prefill_token_threshold is not None + and self.long_prefill_token_threshold < 0 + ): + raise ValueError("long_prefill_token_threshold must be >= 0") + if self.min_prefill_chunk_size is not None and self.min_prefill_chunk_size < 1: + raise ValueError("min_prefill_chunk_size must be >= 1") if ( self.kv_transfer_config is not None diff --git a/python/infinilm/distributed/dist_config.py b/python/infinilm/distributed/dist_config.py index 1d245d364..e0b2035c6 100644 --- a/python/infinilm/distributed/dist_config.py +++ b/python/infinilm/distributed/dist_config.py @@ -7,6 +7,7 @@ def __init__( self, tp_size=None, tp_device_ids=None, + allreduce_backend="nccl", moe_ep_backend="disabled", moe_ep_size=1, ): @@ -15,15 +16,32 @@ def __init__( if tp_size is not None and tp_device_ids is not None: raise ValueError("Provide either tp_size OR tp_device_ids, not both") + backend = self._parse_allreduce_backend(_infinilm, allreduce_backend) if tp_size is not None: - self._underlying = _infinilm.DistConfig(tp_size) + self._underlying = _infinilm.DistConfig(tp_size, backend) elif tp_device_ids is not None: - self._underlying = _infinilm.DistConfig(tp_device_ids) + self._underlying = _infinilm.DistConfig(tp_device_ids, backend) else: self._underlying = _infinilm.DistConfig() + self._underlying.allreduce_backend = backend self.moe_ep_backend = moe_ep_backend self.moe_ep_size = moe_ep_size + @staticmethod + def _parse_allreduce_backend(_infinilm, value): + if hasattr(_infinilm, "AllReduceBackend") and isinstance(value, _infinilm.AllReduceBackend): + return value + normalized = str(value).lower().replace("-", "_") + if normalized == "auto": + return _infinilm.AllReduceBackend.AUTO + if normalized == "nccl": + return _infinilm.AllReduceBackend.NCCL + if normalized == "custom": + return _infinilm.AllReduceBackend.CUSTOM + raise ValueError( + f"Unsupported allreduce_backend={value!r}; expected one of auto/nccl/custom" + ) + @property def tp_device_ids(self): return self._underlying.tp_device_ids @@ -48,6 +66,16 @@ def moe_ep_size(self): def moe_ep_size(self, value): self._underlying.moe_ep_size = int(value) + @property + def allreduce_backend(self): + return self._underlying.allreduce_backend + + @allreduce_backend.setter + def allreduce_backend(self, value): + from infinilm.lib import _infinilm + + self._underlying.allreduce_backend = self._parse_allreduce_backend(_infinilm, value) + def __repr__(self): return repr(self._underlying) diff --git a/python/infinilm/generation/utils.py b/python/infinilm/generation/utils.py index bad9c2613..ed74cc184 100644 --- a/python/infinilm/generation/utils.py +++ b/python/infinilm/generation/utils.py @@ -22,6 +22,7 @@ def infini_to_ctype_dtype(infini_dtype): def infini_to_numpy(infini_tensor: infinicore.Tensor): if infini_tensor.device.type != "cpu": infini_tensor_cpu = infini_tensor.to(infinicore.device("cpu", 0)) + infinicore.sync_stream() else: infini_tensor_cpu = infini_tensor diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 3e874046e..b89cde60b 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -455,15 +455,18 @@ def generate( input_ids, generation_config, *, + prefill_position_ids=None, pixel_values=None, image_bound=None, tgt_sizes=None, _measure_and_log_time=False, ): eos_token_id = self.eos_token_id + self.reset_request_state() past_seq_len = 0 output_ids = [] + output_history = None initial_batch_size, initial_seqlen = input_ids.shape[:2] seq_len = initial_seqlen batch_size = initial_batch_size @@ -474,7 +477,15 @@ def generate( ) if _measure_and_log_time: - time_measurements = [] + generation_start_time = time.perf_counter() + prefill_latency = None + decode_start_time = None + + mrope_position_delta = None + if prefill_position_ids is not None: + mrope_position_delta = ( + int(prefill_position_ids.to_numpy().max()) + 1 - initial_seqlen + ) block_tables = None max_blocks_per_batch = 0 @@ -509,17 +520,22 @@ def generate( ) for iter in range(0, generation_config.max_new_tokens): - if _measure_and_log_time: - start_time = time.perf_counter() - batch_size, seq_len = input_ids.shape[:2] if self.enable_paged_attn: input_ids = input_ids.view([1, batch_size * seq_len]) - position_ids = infinicore.from_list( - list(range(past_seq_len, past_seq_len + seq_len)) * batch_size, - dtype=infinicore.int64, - ) + if iter > 0 and mrope_position_delta is not None: + mrope_position = past_seq_len + mrope_position_delta + position_ids = infinicore.from_list( + [[mrope_position] * 3 for _ in range(batch_size)], + dtype=infinicore.int64, + ) + else: + position_ids = infinicore.from_list( + list(range(past_seq_len, past_seq_len + seq_len)) + * batch_size, + dtype=infinicore.int64, + ) if iter == 0: slot_mapping_list = [] @@ -547,16 +563,26 @@ def generate( dtype=infinicore.int64, ) else: - position_ids = infinicore.from_list( - [ - list(range(past_seq_len, past_seq_len + seq_len)) - for _ in range(batch_size) - ], - dtype=infinicore.int64, - ) + if iter > 0 and mrope_position_delta is not None: + mrope_position = past_seq_len + mrope_position_delta + position_ids = infinicore.from_list( + [[[mrope_position] * 3] for _ in range(batch_size)], + dtype=infinicore.int64, + ) + else: + position_ids = infinicore.from_list( + [ + list(range(past_seq_len, past_seq_len + seq_len)) + for _ in range(batch_size) + ], + dtype=infinicore.int64, + ) slot_mapping = None + if iter == 0 and prefill_position_ids is not None: + position_ids = prefill_position_ids + past_kv_lengths = infinicore.from_list( [past_seq_len] * batch_size, dtype=infinicore.int32 ) @@ -602,7 +628,27 @@ def generate( top_p=generation_config.top_p, ) - output_ids.append(output_id) + if output_history is None and generation_config.max_new_tokens is not None: + output_history = infinicore.empty( + [generation_config.max_new_tokens, batch_size], + dtype=output_id.dtype, + device=output_id.device, + ) + + if output_history is not None: + output_slot = output_history.narrow(0, iter, 1).view([batch_size]) + else: + output_slot = infinicore.empty( + [batch_size], dtype=output_id.dtype, device=output_id.device + ) + self.copy_last_output_to(output_slot) + output_ids.append(output_slot) + + if _measure_and_log_time and iter == 0: + self.sync_last_output() + prefill_end_time = time.perf_counter() + prefill_latency = prefill_end_time - generation_start_time + decode_start_time = prefill_end_time if ( initial_batch_size == 1 @@ -617,24 +663,31 @@ def generate( past_seq_len = past_seq_len + seq_len - if _measure_and_log_time: - end_time = time.perf_counter() - - time_measurements.append((end_time - start_time)) + if output_ids: + self.sync_last_output() if _measure_and_log_time: + generation_end_time = time.perf_counter() + total_latency = generation_end_time - generation_start_time + if prefill_latency is None: + prefill_latency = total_latency + decode_latency = 0.0 + else: + decode_latency = generation_end_time - decode_start_time + decode_tokens = max(0, len(output_ids) - 1) + print( - f"\n\n\n Generation completed in {round(sum(time_measurements) * 1000, 2)} ms" + f"\n\n\n Generation completed in {round(total_latency * 1000, 2)} ms" ) print( - f" Batchsize={initial_batch_size} Per_Batch_Input_Len={initial_seqlen} Per_Batch_New_Tokens={len(time_measurements)}\n" + f" Batchsize={initial_batch_size} Per_Batch_Input_Len={initial_seqlen} Per_Batch_New_Tokens={len(output_ids)}\n" ) print( - f" Prefill TTFT: {round(time_measurements[0] * 1000, 2)} ms Throughput: {round((initial_batch_size * initial_seqlen) / time_measurements[0], 2)} tok/s\n", + f" Prefill TTFT: {round(prefill_latency * 1000, 2)} ms Throughput: {round((initial_batch_size * initial_seqlen) / prefill_latency, 2)} tok/s\n", ) - if len(time_measurements) > 1: + if decode_tokens > 0: print( - f" Decode Avg ITL: {round(sum(time_measurements[1:]) * 1000 / (len(time_measurements) - 1), 2)} ms Throughput: {round((initial_batch_size * (len(time_measurements) - 1)) / sum(time_measurements[1:]), 2)} tok/s\n", + f" Decode Avg ITL: {round(decode_latency * 1000 / decode_tokens, 2)} ms Throughput: {round((initial_batch_size * decode_tokens) / decode_latency, 2)} tok/s\n", ) return output_ids @@ -644,6 +697,9 @@ def reset_cache(self, cache_config): self.enable_paged_attn = isinstance(cache_config, PagedKVCacheConfig) super().reset_cache(cache_config) + def copy_last_output_to(self, dst): + super().copy_last_output_to(dst._underlying) + def state_dict_keyname(self): return list(super().state_dict_keyname()) diff --git a/python/infinilm/llm/cache_manager.py b/python/infinilm/llm/cache_manager.py index 00abd541d..27ff7455b 100644 --- a/python/infinilm/llm/cache_manager.py +++ b/python/infinilm/llm/cache_manager.py @@ -112,6 +112,14 @@ def __init__(self, num_blocks: int, block_size: int): self.free_block_ids: deque = deque(range(num_blocks)) self.used_block_ids: Set[int] = set() self.pending_block_ids: Set[int] = set() + self._stats = { + "prefix_cache_queries": 0, + "prefix_cache_hit_requests": 0, + "prefix_cache_hit_blocks": 0, + "prefix_cache_query_blocks": 0, + "prefix_cache_hit_tokens": 0, + "prefix_cache_evicted_blocks": 0, + } def __repr__(self): return ( @@ -150,6 +158,7 @@ def _deallocate_block(self, block_id: int): if block.hash != -1 and self.hash_to_block_id.get(block.hash) == block_id: del self.hash_to_block_id[block.hash] + self._stats["prefix_cache_evicted_blocks"] += 1 block.free() self.used_block_ids.remove(block_id) @@ -283,6 +292,13 @@ def get_computed_blocks( num_local_cached_tokens = max_blocks_to_reuse * self.block_size + self._stats["prefix_cache_queries"] += 1 + self._stats["prefix_cache_query_blocks"] += num_full_blocks + self._stats["prefix_cache_hit_blocks"] += max_blocks_to_reuse + self._stats["prefix_cache_hit_tokens"] += num_local_cached_tokens + if max_blocks_to_reuse > 0: + self._stats["prefix_cache_hit_requests"] += 1 + for block_id in range(max_blocks_to_reuse): block = self.blocks[blocks_blueprint[block_id]["block_id"]] block.ref_count += 1 @@ -325,9 +341,11 @@ def allocate_slots( total_tokens = num_computed_tokens + num_new_tokens - num_blocks_needed = ( - total_tokens + self.block_size - 1 - ) // self.block_size - len(cached_block_table) + num_blocks_needed = max( + (total_tokens + self.block_size - 1) // self.block_size + - len(cached_block_table), + 0, + ) if not self.can_allocate(num_blocks_needed): if not self.try_free_blocks(num_blocks_needed): @@ -341,7 +359,7 @@ def allocate_slots( for block_idx in range(start_block_idx, total_blocks): start_tok = block_idx * self.block_size - end_tok = min(start_tok + self.block_size, len(token_ids)) + end_tok = min(start_tok + self.block_size, total_tokens) block_tokens = token_ids[start_tok:end_tok] is_full_block = len(block_tokens) == self.block_size @@ -533,6 +551,36 @@ def try_free_blocks(self, num_required: int) -> bool: # PD-disaggregation specific + def commit_computed_blocks( + self, + block_table: List[int], + token_ids: List[int], + start_token: int, + end_token: int, + ) -> None: + """Register hashes for full blocks whose KV has just been computed.""" + if end_token <= start_token: + return + + first_block = start_token // self.block_size + last_full_block = end_token // self.block_size + + for block_idx in range(first_block, last_full_block): + block_start = block_idx * self.block_size + block_end = block_start + self.block_size + if block_end > end_token: + break + + block_id = block_table[block_idx] + block = self.blocks[block_id] + if block.hash == -1: + prefix_hash = -1 + if block_idx > 0: + prefix_hash = self.blocks[block_table[block_idx - 1]].hash + block_tokens = token_ids[block_start:block_end] + block.update(self.compute_hash(block_tokens, prefix_hash), block_tokens) + self.hash_to_block_id[block.hash] = block_id + def update_blocks_hash(self, block_table: List[int], num_local_cached_tokens: int): """Register hashes for blocks beyond the locally cached prefix into the lookup table. @@ -601,3 +649,16 @@ def update_blocks_slot( new_slot_mapping.extend(range(base, base + end_offset)) return new_slot_mapping + + def get_stats(self) -> dict: + """Return prefix-cache statistics without mutating allocator state.""" + stats = dict(self._stats) + query_blocks = stats["prefix_cache_query_blocks"] + stats["prefix_cache_block_hit_rate"] = ( + stats["prefix_cache_hit_blocks"] / query_blocks if query_blocks else 0.0 + ) + queries = stats["prefix_cache_queries"] + stats["prefix_cache_request_hit_rate"] = ( + stats["prefix_cache_hit_requests"] / queries if queries else 0.0 + ) + return stats diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py index 59e9e94ba..d9c868b03 100644 --- a/python/infinilm/llm/llm.py +++ b/python/infinilm/llm/llm.py @@ -12,6 +12,7 @@ import threading import time import uuid +from dataclasses import dataclass from typing import AsyncIterator, List, Optional, Union import janus @@ -34,6 +35,138 @@ logger = logging.getLogger(__name__) +_TRUE_ENV_VALUES = {"1", "true", "yes", "on"} + + +def _parse_env_bool(value: str) -> bool: + return value.strip().lower() in _TRUE_ENV_VALUES + + +def _env_flag(name: str, default: bool = False) -> bool: + value = os.getenv(name) + if value is None: + return default + return _parse_env_bool(value) + + +def _env_optional_flag(name: str) -> Optional[bool]: + value = os.getenv(name) + if value is None or value == "": + return None + return _parse_env_bool(value) + + +def _env_optional_int(name: str, default: Optional[int] = None) -> Optional[int]: + value = os.getenv(name) + if value is None or value == "": + return default + return int(value) + + +@dataclass(frozen=True) +class PagedSchedulerOptions: + max_num_batched_tokens: int + enable_chunked_prefill: bool + prefill_chunk_size: Optional[int] + decode_priority: bool + max_num_partial_prefills: int + max_long_partial_prefills: int + long_prefill_token_threshold: Optional[int] + min_prefill_chunk_size: Optional[int] + + +def _resolve_paged_scheduler_options( + config: EngineConfig, + max_position_embeddings: int, +) -> PagedSchedulerOptions: + max_num_batched_tokens = _env_optional_int( + "INFINILM_MAX_NUM_BATCHED_TOKENS", + config.max_num_batched_tokens or max_position_embeddings, + ) + if not 1024 <= max_num_batched_tokens <= max_position_embeddings: + raise ValueError( + "max_num_batched_tokens must be between 1024 and " + f"max_position_embeddings ({max_position_embeddings}), got " + f"{max_num_batched_tokens}" + ) + + enable_chunked_prefill = _env_flag( + "INFINILM_ENABLE_CHUNKED_PREFILL", + config.enable_chunked_prefill, + ) + + prefill_chunk_size = _env_optional_int( + "INFINILM_PREFILL_CHUNK_SIZE", + config.prefill_chunk_size, + ) + if prefill_chunk_size is not None and not ( + 0 < prefill_chunk_size <= max_position_embeddings + ): + raise ValueError( + "prefill_chunk_size must be between 1 and " + f"max_position_embeddings ({max_position_embeddings}), got " + f"{prefill_chunk_size}" + ) + + decode_priority = _env_optional_flag("INFINILM_DECODE_PRIORITY") + if decode_priority is None: + decode_priority = config.decode_priority + if decode_priority is None: + decode_priority = enable_chunked_prefill + + max_num_partial_prefills = _env_optional_int( + "INFINILM_MAX_NUM_PARTIAL_PREFILLS", + config.max_num_partial_prefills, + ) + max_long_partial_prefills = _env_optional_int( + "INFINILM_MAX_LONG_PARTIAL_PREFILLS", + config.max_long_partial_prefills, + ) + if max_num_partial_prefills < 1: + raise ValueError("max_num_partial_prefills must be >= 1") + if not 1 <= max_long_partial_prefills <= max_num_partial_prefills: + raise ValueError( + "max_long_partial_prefills must be between 1 and max_num_partial_prefills" + ) + + long_prefill_token_threshold = _env_optional_int( + "INFINILM_LONG_PREFILL_TOKEN_THRESHOLD", + config.long_prefill_token_threshold, + ) + if long_prefill_token_threshold is None: + long_prefill_token_threshold = prefill_chunk_size or max_num_batched_tokens + if not 0 <= long_prefill_token_threshold <= max_position_embeddings: + raise ValueError( + "long_prefill_token_threshold must be between 0 and " + f"max_position_embeddings ({max_position_embeddings}), got " + f"{long_prefill_token_threshold}" + ) + + min_prefill_chunk_size = _env_optional_int( + "INFINILM_MIN_PREFILL_CHUNK_SIZE", + config.min_prefill_chunk_size, + ) + if min_prefill_chunk_size is not None and not ( + 0 < min_prefill_chunk_size <= max_position_embeddings + ): + raise ValueError( + "min_prefill_chunk_size must be between 1 and " + f"max_position_embeddings ({max_position_embeddings}), got " + f"{min_prefill_chunk_size}" + ) + + return PagedSchedulerOptions( + max_num_batched_tokens=max_num_batched_tokens, + enable_chunked_prefill=enable_chunked_prefill, + prefill_chunk_size=prefill_chunk_size, + decode_priority=decode_priority, + max_num_partial_prefills=max_num_partial_prefills, + max_long_partial_prefills=max_long_partial_prefills, + long_prefill_token_threshold=long_prefill_token_threshold, + min_prefill_chunk_size=min_prefill_chunk_size, + ) + + class LLMEngine: """Low-level LLM engine that handles inference execution.""" @@ -82,19 +215,26 @@ def __init__(self, config: EngineConfig): ) num_mamba_cache_blocks = max(2, config.num_blocks // 4) - max_num_batched_tokens = int( - os.getenv("INFINILM_MAX_NUM_BATCHED_TOKENS", max_position_embeddings) + scheduler_options = _resolve_paged_scheduler_options( + config, + max_position_embeddings, ) - assert 1024 <= max_num_batched_tokens <= max_position_embeddings self.scheduler = Scheduler( max_batch_size=config.max_batch_size, num_blocks=config.num_blocks, block_size=config.block_size, - max_num_batched_tokens=max_num_batched_tokens, + max_num_batched_tokens=scheduler_options.max_num_batched_tokens, connector=connector, has_mamba_cache=has_mamba_cache, num_mamba_cache_blocks=num_mamba_cache_blocks, + enable_chunked_prefill=scheduler_options.enable_chunked_prefill, + prefill_chunk_size=scheduler_options.prefill_chunk_size, + max_num_partial_prefills=scheduler_options.max_num_partial_prefills, + max_long_partial_prefills=scheduler_options.max_long_partial_prefills, + long_prefill_token_threshold=scheduler_options.long_prefill_token_threshold, + min_prefill_chunk_size=scheduler_options.min_prefill_chunk_size, + decode_priority=scheduler_options.decode_priority, ) logger.info(f"Using Paged KV Cache with num_blocks={config.num_blocks}") if has_mamba_cache: @@ -122,6 +262,10 @@ def add_request(self, request: InferenceRequest): """Add a request to the scheduler.""" self.scheduler.add_request(request) + def get_cache_stats(self) -> dict: + """Return scheduler and KV cache statistics.""" + return self.scheduler.get_cache_stats() + def step(self) -> tuple[bool, list[tuple]]: """Run one inference step. @@ -146,6 +290,7 @@ def step(self) -> tuple[bool, list[tuple]]: scheduler_output.is_prefill, scheduler_output.scheduled_requests, sampled_tokens_list, + scheduler_output.scheduled_prefill_flags, ) # Return False (no immediate work) only when no requests were scheduled @@ -166,6 +311,7 @@ def _update_requests( is_prefill: bool, requests: List[InferenceRequest], sampled_tokens: List[int], + request_prefill_flags: Optional[List[bool]] = None, ) -> List[tuple]: """Update request status after inference step.""" if is_prefill: @@ -177,7 +323,19 @@ def _update_requests( case _: raise ValueError(f"Unsupported cache_type: {self.cache_type}") pending = [] - for req, token_ids in zip(requests, sampled_tokens): + requests_to_complete = [] + if request_prefill_flags is None: + request_prefill_flags = [is_prefill] * len(requests) + if len(request_prefill_flags) != len(requests): + raise RuntimeError( + "request_prefill_flags must match scheduled request count" + ) + if len(sampled_tokens) != len(requests): + raise RuntimeError("sampled token count must match scheduled requests") + + for req, token_ids, request_is_prefill in zip( + requests, sampled_tokens, request_prefill_flags + ): if req.is_aborted(): logger.info( f"Request {req.request_id} aborted by client, skipping update" @@ -186,10 +344,17 @@ def _update_requests( # (status still RUNNING). if not req.is_finished(): req.mark_canceled() + requests_to_complete.append(req) continue + if request_is_prefill and self.cache_type == "paged": + can_sample = self.scheduler.complete_prefill_chunk(req) + if not can_sample: + continue + if not isinstance(token_ids, list): token_ids = [token_ids] + requests_to_complete.append(req) for token_id in token_ids: if req.is_finished(): @@ -246,7 +411,7 @@ def _update_requests( continue pending.append((req.output_queue.async_q, output)) - self.scheduler.complete_requests(requests) + self.scheduler.complete_requests(requests_to_complete) return pending def _check_request_finished(self, req: InferenceRequest, token_id: int) -> bool: @@ -318,6 +483,7 @@ def __init__( tensor_parallel_size: int = 1, moe_ep_backend: str = "disabled", moe_ep_size: int = 1, + allreduce_backend: str = "nccl", cache_type: str = "paged", max_batch_size: int = 16, max_tokens: int = 4096, @@ -333,6 +499,14 @@ def __init__( weight_load_mode: str = "async", skip_load: bool = False, skip_legacy_moe: bool = False, + max_num_batched_tokens: Optional[int] = None, + enable_chunked_prefill: bool = False, + prefill_chunk_size: Optional[int] = None, + decode_priority: Optional[bool] = None, + max_num_partial_prefills: int = 1, + max_long_partial_prefills: int = 1, + long_prefill_token_threshold: Optional[int] = None, + min_prefill_chunk_size: Optional[int] = None, ): """Initialize LLM. @@ -364,11 +538,20 @@ def __init__( tensor_parallel_size=tensor_parallel_size, moe_ep_backend=moe_ep_backend, moe_ep_size=moe_ep_size, + allreduce_backend=allreduce_backend, cache_type=cache_type, max_batch_size=max_batch_size, max_tokens=max_tokens, num_blocks=num_blocks, block_size=block_size, + max_num_batched_tokens=max_num_batched_tokens, + enable_chunked_prefill=enable_chunked_prefill, + prefill_chunk_size=prefill_chunk_size, + decode_priority=decode_priority, + max_num_partial_prefills=max_num_partial_prefills, + max_long_partial_prefills=max_long_partial_prefills, + long_prefill_token_threshold=long_prefill_token_threshold, + min_prefill_chunk_size=min_prefill_chunk_size, max_cache_len=max_cache_len, temperature=temperature, top_p=top_p, @@ -487,6 +670,10 @@ def generate( outputs = [req.to_request_output() for req in requests] return outputs + def get_cache_stats(self) -> dict: + """Return scheduler and KV cache statistics.""" + return self.engine.get_cache_stats() + def chat( self, messages: Union[List[dict], List[List[dict]]], @@ -525,6 +712,7 @@ def __init__( tensor_parallel_size: int = 1, moe_ep_backend: str = "disabled", moe_ep_size: int = 1, + allreduce_backend: str = "nccl", cache_type: str = "paged", max_batch_size: int = 16, max_tokens: int = 512, @@ -540,6 +728,14 @@ def __init__( use_mla: bool = False, weight_load_mode: str = "async", skip_legacy_moe: bool = False, + max_num_batched_tokens: Optional[int] = None, + enable_chunked_prefill: bool = False, + prefill_chunk_size: Optional[int] = None, + decode_priority: Optional[bool] = None, + max_num_partial_prefills: int = 1, + max_long_partial_prefills: int = 1, + long_prefill_token_threshold: Optional[int] = None, + min_prefill_chunk_size: Optional[int] = None, ): """Initialize AsyncLLMEngine. @@ -574,11 +770,20 @@ def __init__( tensor_parallel_size=tensor_parallel_size, moe_ep_backend=moe_ep_backend, moe_ep_size=moe_ep_size, + allreduce_backend=allreduce_backend, cache_type=cache_type, max_batch_size=max_batch_size, max_tokens=max_tokens, num_blocks=num_blocks, block_size=block_size, + max_num_batched_tokens=max_num_batched_tokens, + enable_chunked_prefill=enable_chunked_prefill, + prefill_chunk_size=prefill_chunk_size, + decode_priority=decode_priority, + max_num_partial_prefills=max_num_partial_prefills, + max_long_partial_prefills=max_long_partial_prefills, + long_prefill_token_threshold=long_prefill_token_threshold, + min_prefill_chunk_size=min_prefill_chunk_size, max_cache_len=max_cache_len, temperature=temperature, top_p=top_p, @@ -602,6 +807,10 @@ def __init__( def is_healthy(self) -> bool: return bool(self._healthy) + def get_cache_stats(self) -> dict: + """Return scheduler and KV cache statistics.""" + return self.engine.get_cache_stats() + def start(self): """Start the background inference loop.""" if self._running: diff --git a/python/infinilm/llm/model_runner/model_runner.py b/python/infinilm/llm/model_runner/model_runner.py index 45f1eaee1..1971e058d 100644 --- a/python/infinilm/llm/model_runner/model_runner.py +++ b/python/infinilm/llm/model_runner/model_runner.py @@ -39,6 +39,7 @@ class ModelRunnerOutput: req_ids: list[str] = field(default_factory=list) sampled_token_ids: list[int | list[int]] = field(default_factory=list) kv_connector_output: KVConnectorOutput | None = None + execution_stats: dict[str, Any] = field(default_factory=dict) class ModelRunner: @@ -71,6 +72,7 @@ def __init__(self, config: EngineConfig): device=self.device, distributed_config=DistConfig( config.tensor_parallel_size, + allreduce_backend=config.allreduce_backend, moe_ep_backend=config.moe_ep_backend, moe_ep_size=config.moe_ep_size, ), @@ -163,16 +165,19 @@ def _init_device(self): def execute_model(self, scheduler_output) -> ModelRunnerOutput: sampled_tokens_list = [] + execution_stats = {} kv_connector_output = None if self.kv_connector is None: - sampled_tokens_list = self._model_forward(scheduler_output) + sampled_tokens_list, execution_stats = self._model_forward(scheduler_output) else: with self.maybe_get_kv_connector_output( scheduler_output, ) as kv_connector_output: if scheduler_output.num_requests > 0: - sampled_tokens_list = self._model_forward(scheduler_output) + sampled_tokens_list, execution_stats = self._model_forward( + scheduler_output + ) # model_runner_output req_ids = [] @@ -183,9 +188,67 @@ def execute_model(self, scheduler_output) -> ModelRunnerOutput: req_ids=req_ids, sampled_token_ids=sampled_tokens_list, kv_connector_output=kv_connector_output, + execution_stats=execution_stats, ) def _model_forward(self, scheduler_output): + if self._should_split_mixed_batch(scheduler_output): + return self._model_forward_by_phase(scheduler_output) + + sampled_tokens_list = self._model_forward_single(scheduler_output) + return sampled_tokens_list, self._make_execution_stats( + scheduler_output, split_mixed_batch=False + ) + + def _should_split_mixed_batch(self, scheduler_output) -> bool: + return ( + self.config.cache_type == "paged" + and getattr(scheduler_output, "is_mixed", False) + and bool(getattr(scheduler_output, "execution_phases", None)) + ) + + def _model_forward_by_phase(self, scheduler_output): + sampled_tokens_by_index = [None] * scheduler_output.num_requests + phase_records = [] + + for phase in scheduler_output.execution_phases: + sub_output = scheduler_output.make_subset(phase.request_indices) + phase_tokens = self._model_forward_single(sub_output) + if len(phase_tokens) != len(phase.request_indices): + raise RuntimeError( + f"{phase.kind} phase sampled token count does not match request count" + ) + for output_index, token_id in zip(phase.request_indices, phase_tokens): + sampled_tokens_by_index[output_index] = token_id + phase_records.append( + { + "kind": phase.kind, + "num_requests": len(phase.request_indices), + "num_tokens": phase.num_tokens, + } + ) + + if any(token_id is None for token_id in sampled_tokens_by_index): + raise RuntimeError("phase execution did not cover every scheduled request") + + return sampled_tokens_by_index, { + "split_mixed_batch": True, + "phases": phase_records, + } + + def _make_execution_stats(self, scheduler_output, split_mixed_batch: bool) -> dict: + phases = [] + for phase in getattr(scheduler_output, "execution_phases", []): + phases.append( + { + "kind": phase.kind, + "num_requests": len(phase.request_indices), + "num_tokens": phase.num_tokens, + } + ) + return {"split_mixed_batch": split_mixed_batch, "phases": phases} + + def _model_forward_single(self, scheduler_output): # Build model inputs model_input = self.processor.build_model_inputs( scheduler_output, diff --git a/python/infinilm/llm/request.py b/python/infinilm/llm/request.py index abccfe539..f842f5b66 100644 --- a/python/infinilm/llm/request.py +++ b/python/infinilm/llm/request.py @@ -161,6 +161,9 @@ def __init__( ) self.num_computed_tokens: int = 0 # Total tokens computed (local + remote) self.num_blocks: int = 0 + self.prefill_chunk_start: int = 0 + self.prefill_chunk_end: int = 0 + self.prefill_chunk_is_final: bool = True # Mamba cache management. None means no mamba cache row is currently owned. self.mamba_cache_index: Optional[int] = None diff --git a/python/infinilm/llm/scheduler.py b/python/infinilm/llm/scheduler.py index a1153e916..1ccf87b11 100644 --- a/python/infinilm/llm/scheduler.py +++ b/python/infinilm/llm/scheduler.py @@ -4,6 +4,7 @@ import logging import queue +from dataclasses import dataclass from typing import List, Optional import janus @@ -44,6 +45,13 @@ def commit_accepted_tokens( self._cache_manager.commit_blocks_hash(block_table, token_ids, num_tokens) +@dataclass(frozen=True) +class ExecutionPhase: + kind: str + request_indices: tuple[int, ...] + num_tokens: int + + class SchedulerOutput: """Scheduler output containing scheduled requests and execution phase info.""" @@ -51,14 +59,89 @@ def __init__( self, scheduled_requests: List[InferenceRequest], is_prefill: bool = False, + scheduled_prefill_flags: Optional[List[bool]] = None, speculative_cache_ops: Optional[SpeculativeCacheOps] = None, ): self.scheduled_requests = scheduled_requests self.num_requests = len(scheduled_requests) - self.is_prefill = is_prefill + if scheduled_prefill_flags is None: + scheduled_prefill_flags = [is_prefill] * self.num_requests + if len(scheduled_prefill_flags) != self.num_requests: + raise ValueError( + "scheduled_prefill_flags must match scheduled_requests length" + ) + self.scheduled_prefill_flags = scheduled_prefill_flags self.speculative_cache_ops = speculative_cache_ops + # C++ attention currently chooses the prefill/extend path at batch level. + # A mixed decode+prefill batch must therefore advertise prefill=True. + self.is_prefill = ( + any(scheduled_prefill_flags) if scheduled_requests else is_prefill + ) + self.num_prefill_tokens = 0 + self.num_decode_tokens = 0 + for req, request_is_prefill in zip( + scheduled_requests, scheduled_prefill_flags + ): + if request_is_prefill: + self.num_prefill_tokens += max( + 0, req.prefill_chunk_end - req.prefill_chunk_start + ) + else: + self.num_decode_tokens += 1 + self.num_batched_tokens = self.num_prefill_tokens + self.num_decode_tokens + self.execution_phases = self._build_execution_phases() self.kv_connector_metadata = None + def request_is_prefill(self, index: int) -> bool: + return self.scheduled_prefill_flags[index] + + @property + def is_mixed(self) -> bool: + return self.num_decode_tokens > 0 and self.num_prefill_tokens > 0 + + def _build_execution_phases(self) -> List[ExecutionPhase]: + decode_indices = tuple( + i + for i, is_prefill in enumerate(self.scheduled_prefill_flags) + if not is_prefill + ) + prefill_indices = tuple( + i + for i, is_prefill in enumerate(self.scheduled_prefill_flags) + if is_prefill + ) + phases = [] + if decode_indices: + phases.append( + ExecutionPhase( + kind="decode", + request_indices=decode_indices, + num_tokens=len(decode_indices), + ) + ) + if prefill_indices: + phases.append( + ExecutionPhase( + kind="prefill", + request_indices=prefill_indices, + num_tokens=self.num_prefill_tokens, + ) + ) + return phases + + def make_subset(self, request_indices: List[int]) -> "SchedulerOutput": + scheduled_requests = [self.scheduled_requests[i] for i in request_indices] + scheduled_prefill_flags = [ + self.scheduled_prefill_flags[i] for i in request_indices + ] + subset = SchedulerOutput( + scheduled_requests=scheduled_requests, + scheduled_prefill_flags=scheduled_prefill_flags, + speculative_cache_ops=self.speculative_cache_ops, + ) + subset.kv_connector_metadata = self.kv_connector_metadata + return subset + class Scheduler: """Request scheduler with integrated BlockManager for KV cache management. @@ -78,6 +161,13 @@ def __init__( connector=None, has_mamba_cache: bool = False, num_mamba_cache_blocks: int | None = None, + enable_chunked_prefill: bool = False, + prefill_chunk_size: Optional[int] = None, + max_num_partial_prefills: int = 1, + max_long_partial_prefills: int = 1, + long_prefill_token_threshold: Optional[int] = None, + min_prefill_chunk_size: Optional[int] = None, + decode_priority: bool = False, ): self.waiting_queue = janus.Queue() self.running_queue = janus.Queue() @@ -100,6 +190,77 @@ def __init__( self.block_size = block_size self.max_num_batched_tokens = max_num_batched_tokens self.connector = connector + self.enable_chunked_prefill = enable_chunked_prefill + self.prefill_chunk_size = prefill_chunk_size or max_num_batched_tokens + self.max_num_partial_prefills = max(1, max_num_partial_prefills) + self.max_long_partial_prefills = min( + max(1, max_long_partial_prefills), self.max_num_partial_prefills + ) + self.long_prefill_token_threshold = long_prefill_token_threshold + self.min_prefill_chunk_size = min_prefill_chunk_size or min( + block_size, self.prefill_chunk_size + ) + self.decode_priority = decode_priority + self._stats = { + "schedule_steps": 0, + "decode_steps": 0, + "prefill_steps": 0, + "mixed_steps": 0, + "decode_tokens": 0, + "prefill_tokens": 0, + "partial_prefill_chunks": 0, + "full_prefill_chunks": 0, + "prefill_deferred_by_token_budget": 0, + "prefill_deferred_by_partial_limit": 0, + "prefill_deferred_by_long_partial_limit": 0, + "execution_split_mixed_steps": 0, + "execution_decode_phases": 0, + "execution_prefill_phases": 0, + "execution_decode_phase_tokens": 0, + "execution_prefill_phase_tokens": 0, + } + + def _is_partial_prefill( + self, + request: InferenceRequest, + num_computed_tokens: int, + num_new_tokens: int, + ) -> bool: + return ( + self.enable_chunked_prefill + and num_new_tokens > 0 + and num_computed_tokens + num_new_tokens < request.get_prompt_length() + ) + + def _is_long_prefill( + self, request: InferenceRequest, num_computed_tokens: int + ) -> bool: + threshold = self.long_prefill_token_threshold + if threshold is None or threshold <= 0: + return False + return ( + request.get_prompt_length() > threshold + and request.get_prompt_length() - num_computed_tokens > threshold + ) + + def _partial_prefill_limit_reason( + self, + request: InferenceRequest, + num_computed_tokens: int, + num_new_tokens: int, + partial_prefill_count: int, + long_partial_prefill_count: int, + ) -> Optional[str]: + if not self._is_partial_prefill(request, num_computed_tokens, num_new_tokens): + return None + if partial_prefill_count >= self.max_num_partial_prefills: + return "partial_limit" + if ( + self._is_long_prefill(request, num_computed_tokens) + and long_partial_prefill_count >= self.max_long_partial_prefills + ): + return "long_partial_limit" + return None def add_request(self, request: InferenceRequest): if request is not None: @@ -125,17 +286,234 @@ def _exceeds_token_budget( ) def schedule(self) -> Optional[SchedulerOutput]: - """Schedule and return batch of requests to execute.""" - deferred_requests = [] - scheduled_requests = [] - is_prefill = False + """Schedule and return the next batch of requests to execute. + + The default policy preserves the original prefill-first behavior. When + chunked prefill and decode priority are enabled, it follows vLLM's shape: + schedule ready decodes first, then fill the remaining token budget with + prefill chunks in the same batch. + """ + scheduled_requests: List[InferenceRequest] = [] + scheduled_prefill_flags: List[bool] = [] current_num_batched_tokens = 0 - current_prefill_extra_blocks = 0 - # Process Waiting queue (prefill phase) - while ( - len(scheduled_requests) < self.max_batch_size - and current_num_batched_tokens < self.max_num_batched_tokens + if self.decode_priority and self.enable_chunked_prefill: + current_num_batched_tokens = self._schedule_decode_requests( + scheduled_requests, + scheduled_prefill_flags, + current_num_batched_tokens, + ) + current_prefill_extra_blocks = sum( + self._get_prefill_extra_blocks(req) for req in scheduled_requests + ) + self._schedule_prefill_requests( + scheduled_requests, + scheduled_prefill_flags, + current_num_batched_tokens, + current_prefill_extra_blocks, + ) + else: + self._schedule_prefill_requests( + scheduled_requests, + scheduled_prefill_flags, + current_num_batched_tokens, + ) + if not scheduled_requests: + scheduler_output = self._schedule_decode_batch() + if scheduler_output is not None: + return scheduler_output + + if scheduled_requests: + return self._make_scheduler_output( + scheduled_requests, + scheduled_prefill_flags, + ) + + if self.connector is not None: + return self._make_connector_only_output() + + return None + + def update_waiting_for_remote_kv(self, request: InferenceRequest): + self.remote_kv_requests.pop(request.request_id, None) + self.pending_kv_decode_blocks -= ( + request.sampling_params.max_tokens + self.block_size - 1 + ) // self.block_size + if request.request_id in self.failed_receiving_kv_req_ids: + if request.num_computed_tokens: + valid_block_count = request.num_computed_tokens // self.block_size + self.cache_manager.update_blocks_hash( + request.block_table[:valid_block_count], + request.num_local_cached_tokens, + ) + request.slot_mapping = self.cache_manager.update_blocks_slot( + request.block_table, + request.num_computed_tokens, + request.get_prompt_length(), + ) + request.num_local_cached_tokens = request.num_computed_tokens + else: + self.cache_manager.free_blocks(request.block_table) + request.block_table = [] + request.slot_mapping = [] + request.num_local_cached_tokens = 0 + self.failed_receiving_kv_req_ids.discard(request.request_id) + else: + self.cache_manager.update_blocks_hash( + request.block_table, request.num_local_cached_tokens + ) + request.num_local_cached_tokens = request.num_computed_tokens + self.finished_receiving_kv_req_ids.discard(request.request_id) + + def complete_prefill_chunk(self, request: InferenceRequest) -> bool: + """Finalize a prefill chunk and return True when generation may sample.""" + chunk_start = request.prefill_chunk_start + chunk_end = request.prefill_chunk_end + self.cache_manager.commit_computed_blocks( + request.block_table, + request.get_input_tokens(), + chunk_start, + chunk_end, + ) + request.num_computed_tokens = chunk_end + request.num_local_cached_tokens = chunk_end + request.slot_mapping = [] + + if chunk_end < request.get_prompt_length(): + request.status = RequestStatus.WAITING + self.waiting_queue.sync_q.put(request) + return False + + request.status = RequestStatus.RUNNING + return True + + def _get_prefill_chunk_num_tokens( + self, + request: InferenceRequest, + num_computed_tokens: int, + current_num_batched_tokens: int, + num_scheduled_requests: int, + ) -> int: + prompt_len = request.get_prompt_length() + remaining_tokens = prompt_len - num_computed_tokens + if remaining_tokens <= 0: + return 0 + + if not self.enable_chunked_prefill: + return remaining_tokens + + token_budget_left = self.max_num_batched_tokens - current_num_batched_tokens + if num_scheduled_requests == 0: + token_budget_left = max(token_budget_left, 1) + if token_budget_left <= 0: + return 0 + + num_tokens = min( + remaining_tokens, + self.prefill_chunk_size, + token_budget_left, + ) + + if num_computed_tokens % self.block_size != 0: + next_block_boundary = ( + (num_computed_tokens // self.block_size) + 1 + ) * self.block_size + tokens_to_block_boundary = next_block_boundary - num_computed_tokens + num_tokens = min(num_tokens, tokens_to_block_boundary) + else: + tokens_to_block_boundary = None + + if 0 < num_tokens < self.min_prefill_chunk_size and num_scheduled_requests > 0: + final_chunk = num_tokens >= remaining_tokens + completes_partial_block = ( + tokens_to_block_boundary is not None + and num_tokens == tokens_to_block_boundary + ) + if not final_chunk and not completes_partial_block: + return 0 + + return max(num_tokens, 0) + + def _set_prefill_chunk( + self, request: InferenceRequest, chunk_start: int, chunk_end: int + ) -> None: + request.prefill_chunk_start = chunk_start + request.prefill_chunk_end = chunk_end + request.prefill_chunk_is_final = chunk_end >= request.get_prompt_length() + + def _make_scheduler_output( + self, + scheduled_requests: List[InferenceRequest], + scheduled_prefill_flags: List[bool], + ) -> SchedulerOutput: + scheduler_output = SchedulerOutput( + scheduled_requests=scheduled_requests, + scheduled_prefill_flags=scheduled_prefill_flags, + speculative_cache_ops=self.speculative_cache_ops, + ) + self._record_scheduler_output(scheduler_output) + if self.connector is not None: + scheduler_output.kv_connector_metadata = ( + self.connector.build_connector_meta() + ) + return scheduler_output + + def _make_connector_only_output(self) -> SchedulerOutput: + self._stats["schedule_steps"] += 1 + scheduler_output = SchedulerOutput( + scheduled_requests=[], + speculative_cache_ops=self.speculative_cache_ops, + ) + scheduler_output.kv_connector_metadata = self.connector.build_connector_meta() + return scheduler_output + + def _record_scheduler_output(self, scheduler_output: SchedulerOutput) -> None: + self._stats["schedule_steps"] += 1 + has_decode = scheduler_output.num_decode_tokens > 0 + has_prefill = scheduler_output.num_prefill_tokens > 0 + if has_decode: + self._stats["decode_steps"] += 1 + self._stats["decode_tokens"] += scheduler_output.num_decode_tokens + if has_prefill: + self._stats["prefill_steps"] += 1 + self._stats["prefill_tokens"] += scheduler_output.num_prefill_tokens + if has_decode and has_prefill: + self._stats["mixed_steps"] += 1 + + for req, request_is_prefill in zip( + scheduler_output.scheduled_requests, + scheduler_output.scheduled_prefill_flags, + ): + if not request_is_prefill: + continue + if req.prefill_chunk_end < req.get_prompt_length(): + self._stats["partial_prefill_chunks"] += 1 + else: + self._stats["full_prefill_chunks"] += 1 + + def _schedule_prefill_requests( + self, + scheduled_requests: List[InferenceRequest], + scheduled_prefill_flags: List[bool], + current_num_batched_tokens: int, + current_prefill_extra_blocks: int = 0, + ) -> int: + deferred_requests = [] + partial_prefill_count = sum( + 1 + for req, is_prefill in zip(scheduled_requests, scheduled_prefill_flags) + if is_prefill and req.prefill_chunk_end < req.get_prompt_length() + ) + long_partial_prefill_count = sum( + 1 + for req, is_prefill in zip(scheduled_requests, scheduled_prefill_flags) + if is_prefill + and req.prefill_chunk_end < req.get_prompt_length() + and self._is_long_prefill(req, req.prefill_chunk_start) + ) + + while self._has_schedule_capacity( + len(scheduled_requests), current_num_batched_tokens ): try: req = self.waiting_queue.sync_q.get_nowait() @@ -181,17 +559,43 @@ def schedule(self) -> Optional[SchedulerOutput]: num_computed_tokens -= 1 num_new_tokens = req.get_prompt_length() - num_computed_tokens + if self.enable_chunked_prefill and not load_kv_async: + num_new_tokens = self._get_prefill_chunk_num_tokens( + req, + num_computed_tokens, + current_num_batched_tokens, + len(scheduled_requests), + ) + + limit_reason = self._partial_prefill_limit_reason( + req, + num_computed_tokens, + num_new_tokens, + partial_prefill_count, + long_partial_prefill_count, + ) + if limit_reason is not None: + if num_local_computed_tokens > 0: + self.cache_manager.free_blocks(cached_block_table) + if limit_reason == "partial_limit": + self._stats["prefill_deferred_by_partial_limit"] += 1 + else: + self._stats[ + "prefill_deferred_by_long_partial_limit" + ] += 1 + deferred_requests.append(req) + break + # Early token budget check: skip can_accept_request and allocate_slots # for requests that would exceed the per-schedule token budget. if not load_kv_async: - num_tokens_this_step = ( - req.get_prompt_length() - num_local_computed_tokens - ) - if self._exceeds_token_budget( + num_tokens_this_step = num_new_tokens + if num_tokens_this_step <= 0 or self._exceeds_token_budget( current_num_batched_tokens, num_tokens_this_step, len(scheduled_requests), ): + self._stats["prefill_deferred_by_token_budget"] += 1 if num_local_computed_tokens > 0: self.cache_manager.free_blocks(cached_block_table) deferred_requests.append(req) @@ -218,7 +622,7 @@ def schedule(self) -> Optional[SchedulerOutput]: num_computed_tokens=num_computed_tokens, cached_block_table=cached_block_table, blocks_blueprint=blocks_blueprint, - delay_cache_blocks=load_kv_async, + delay_cache_blocks=True, ) if req_blocks is None: @@ -257,20 +661,61 @@ def schedule(self) -> Optional[SchedulerOutput]: ) else: load_kv_async = False - num_tokens_this_step = ( - req.get_prompt_length() - req.num_local_cached_tokens + num_computed_tokens = req.num_computed_tokens + num_new_tokens = self._get_prefill_chunk_num_tokens( + req, + num_computed_tokens, + current_num_batched_tokens, + len(scheduled_requests), + ) + num_tokens_this_step = num_new_tokens + + limit_reason = self._partial_prefill_limit_reason( + req, + num_computed_tokens, + num_new_tokens, + partial_prefill_count, + long_partial_prefill_count, ) - if self._exceeds_token_budget( + if limit_reason is not None: + if limit_reason == "partial_limit": + self._stats["prefill_deferred_by_partial_limit"] += 1 + else: + self._stats[ + "prefill_deferred_by_long_partial_limit" + ] += 1 + deferred_requests.append(req) + break + + if num_tokens_this_step <= 0 or self._exceeds_token_budget( current_num_batched_tokens, num_tokens_this_step, len(scheduled_requests), ): + self._stats["prefill_deferred_by_token_budget"] += 1 deferred_requests.append(req) break - self.cache_manager.update_blocks_hash( - req.block_table, req.num_local_cached_tokens + + req_blocks, slot_mapping = self.cache_manager.allocate_slots( + req_tokens, + num_new_tokens, + num_computed_tokens=num_computed_tokens, + cached_block_table=req.block_table, + delay_cache_blocks=True, ) + if req_blocks is None: + logger.warning( + "Failed to allocate KV cache blocks for request: %s", + req.request_id, + ) + deferred_requests.append(req) + break + + req.block_table = req_blocks + req.slot_mapping = slot_mapping + req.num_blocks = len(req_blocks) + if load_kv_async: req.status = RequestStatus.WAITING_FOR_REMOTE_KVS self.remote_kv_requests[req.request_id] = req @@ -280,42 +725,57 @@ def schedule(self) -> Optional[SchedulerOutput]: continue current_prefill_extra_blocks += self._get_prefill_extra_blocks(req) + self._set_prefill_chunk( + req, + req.num_computed_tokens, + req.num_computed_tokens + num_new_tokens, + ) scheduled_requests.append(req) + scheduled_prefill_flags.append(True) - num_tokens_this_step = req.get_prompt_length() - req.num_local_cached_tokens + num_tokens_this_step = num_new_tokens current_num_batched_tokens += num_tokens_this_step + if self._is_partial_prefill(req, req.num_computed_tokens, num_new_tokens): + partial_prefill_count += 1 + if self._is_long_prefill(req, req.num_computed_tokens): + long_partial_prefill_count += 1 + req.status = RequestStatus.RUNNING if deferred_requests: for req in deferred_requests: self.waiting_queue.sync_q.put(req) - # Return prefill batch if any waiting requests were scheduled - if scheduled_requests: - is_prefill = True - scheduler_output = SchedulerOutput( - scheduled_requests=scheduled_requests, - is_prefill=is_prefill, - speculative_cache_ops=self.speculative_cache_ops, - ) - if self.connector is not None: - meta = self.connector.build_connector_meta() - scheduler_output.kv_connector_metadata = meta - return scheduler_output + return current_num_batched_tokens - # Process Running queue (decode phase) - while len(scheduled_requests) < self.max_batch_size: + def _has_schedule_capacity( + self, + num_scheduled_requests: int, + current_num_batched_tokens: int, + ) -> bool: + return ( + num_scheduled_requests < self.max_batch_size + and current_num_batched_tokens < self.max_num_batched_tokens + ) + + def _schedule_decode_requests( + self, + scheduled_requests: List[InferenceRequest], + scheduled_prefill_flags: List[bool], + current_num_batched_tokens: int = 0, + ) -> int: + while self._has_schedule_capacity( + len(scheduled_requests), current_num_batched_tokens + ): try: req = self.running_queue.sync_q.get_nowait() except queue.Empty: break - # Skip requests that were already finished (e.g., timed out/canceled while running) if req.is_finished(): self.complete_requests([req]) continue - # Decode phase: allocate slot for newly generated token try: req.block_table, new_slot = self.cache_manager.append_slot( req.block_table, req.get_total_length(), req.get_all_token_ids() @@ -324,91 +784,62 @@ def schedule(self) -> Optional[SchedulerOutput]: req.num_blocks = len(req.block_table) req.num_local_cached_tokens = req.get_total_length() - 1 scheduled_requests.append(req) - + scheduled_prefill_flags.append(False) + current_num_batched_tokens += 1 except RuntimeError as e: raise RuntimeError("No available cache blocks for new token") from e - # Promote completed remote KV transfers (lower priority than running queue). - # Cleanup (is_finished, failed re-queue) runs unconditionally; batch append only if slots remain. - if self.connector is not None and self.remote_kv_requests: - for req_id in list(self.remote_kv_requests.keys()): - req = self.remote_kv_requests[req_id] - if req.is_finished(): - self.complete_requests([req]) - continue - if req_id in self.failed_receiving_kv_req_ids: - logger.warning( - f"Request {req_id[:8]}... failed receiving KV, re-queuing for prefill." - ) - self.update_waiting_for_remote_kv(req) - req.status = RequestStatus.WAITING - self.waiting_queue.sync_q.put(req) - elif req_id in self.finished_receiving_kv_req_ids: - if len(scheduled_requests) < self.max_batch_size: - logger.info( - f"Request {req_id[:8]}... finished receiving KV, scheduling for decode." - ) - self.update_waiting_for_remote_kv(req) - req.status = RequestStatus.RUNNING - scheduled_requests.append(req) - else: - break # Defer promotion to next schedule() if batch is full - - # Return decode batch if any running requests were scheduled - if scheduled_requests: - is_prefill = False - scheduler_output = SchedulerOutput( - scheduled_requests=scheduled_requests, - is_prefill=is_prefill, - speculative_cache_ops=self.speculative_cache_ops, - ) + return self._promote_remote_kv_requests( + scheduled_requests, + scheduled_prefill_flags, + current_num_batched_tokens, + ) - if self.connector is not None: - meta = self.connector.build_connector_meta() - scheduler_output.kv_connector_metadata = meta - return scheduler_output + def _schedule_decode_batch(self) -> Optional[SchedulerOutput]: + scheduled_requests: List[InferenceRequest] = [] + scheduled_prefill_flags: List[bool] = [] + self._schedule_decode_requests(scheduled_requests, scheduled_prefill_flags) - if self.connector is not None: - scheduler_output = SchedulerOutput( - scheduled_requests=[], - speculative_cache_ops=self.speculative_cache_ops, - ) - meta = self.connector.build_connector_meta() - scheduler_output.kv_connector_metadata = meta - return scheduler_output + if not scheduled_requests: + return None - return None + return self._make_scheduler_output(scheduled_requests, scheduled_prefill_flags) - def update_waiting_for_remote_kv(self, request: InferenceRequest): - self.remote_kv_requests.pop(request.request_id, None) - self.pending_kv_decode_blocks -= ( - request.sampling_params.max_tokens + self.block_size - 1 - ) // self.block_size - if request.request_id in self.failed_receiving_kv_req_ids: - if request.num_computed_tokens: - valid_block_count = request.num_computed_tokens // self.block_size - self.cache_manager.update_blocks_hash( - request.block_table[:valid_block_count], - request.num_local_cached_tokens, + def _promote_remote_kv_requests( + self, + scheduled_requests: List[InferenceRequest], + scheduled_prefill_flags: Optional[List[bool]] = None, + current_num_batched_tokens: int = 0, + ) -> int: + if self.connector is None or not self.remote_kv_requests: + return current_num_batched_tokens + for req_id in list(self.remote_kv_requests.keys()): + req = self.remote_kv_requests[req_id] + if req.is_finished(): + self.complete_requests([req]) + continue + if req_id in self.failed_receiving_kv_req_ids: + logger.warning( + f"Request {req_id[:8]}... failed receiving KV, re-queuing for prefill." ) - request.slot_mapping = self.cache_manager.update_blocks_slot( - request.block_table, - request.num_computed_tokens, - request.get_prompt_length(), + self.update_waiting_for_remote_kv(req) + req.status = RequestStatus.WAITING + self.waiting_queue.sync_q.put(req) + elif req_id in self.finished_receiving_kv_req_ids: + if not self._has_schedule_capacity( + len(scheduled_requests), current_num_batched_tokens + ): + break + logger.info( + f"Request {req_id[:8]}... finished receiving KV, scheduling for decode." ) - request.num_local_cached_tokens = request.num_computed_tokens - else: - self.cache_manager.free_blocks(request.block_table) - request.block_table = [] - request.slot_mapping = [] - request.num_local_cached_tokens = 0 - self.failed_receiving_kv_req_ids.discard(request.request_id) - else: - self.cache_manager.update_blocks_hash( - request.block_table, request.num_local_cached_tokens - ) - request.num_local_cached_tokens = request.num_computed_tokens - self.finished_receiving_kv_req_ids.discard(request.request_id) + self.update_waiting_for_remote_kv(req) + req.status = RequestStatus.RUNNING + scheduled_requests.append(req) + current_num_batched_tokens += 1 + if scheduled_prefill_flags is not None: + scheduled_prefill_flags.append(False) + return current_num_batched_tokens def complete_requests(self, requests: List[InferenceRequest]): """Handle completed requests and free their blocks.""" @@ -510,6 +941,7 @@ def _get_prefill_extra_blocks(self, request: InferenceRequest) -> int: return max(total_required_blocks - len(request.block_table), 0) def update_from_output(self, model_output): + self._record_execution_stats(getattr(model_output, "execution_stats", None)) if self.connector is None or model_output.kv_connector_output is None: return @@ -558,8 +990,23 @@ def update_from_output(self, model_output): req = self.remote_kv_requests[req_id] req.num_computed_tokens = req.num_local_cached_tokens + def _record_execution_stats(self, execution_stats: Optional[dict]) -> None: + if not execution_stats: + return + if execution_stats.get("split_mixed_batch"): + self._stats["execution_split_mixed_steps"] += 1 + for phase in execution_stats.get("phases", []): + kind = phase.get("kind") + num_tokens = int(phase.get("num_tokens") or 0) + if kind == "decode": + self._stats["execution_decode_phases"] += 1 + self._stats["execution_decode_phase_tokens"] += num_tokens + elif kind == "prefill": + self._stats["execution_prefill_phases"] += 1 + self._stats["execution_prefill_phase_tokens"] += num_tokens + def get_cache_stats(self) -> dict: - """Get cache statistics.""" + """Get cache and scheduler statistics.""" stats = { "num_blocks": self.cache_manager.num_blocks, "block_size": self.cache_manager.block_size, @@ -567,7 +1014,19 @@ def get_cache_stats(self) -> dict: "usable_blocks": self.cache_manager.get_total_usable_blocks(), "num_pending_blocks": len(self.cache_manager.pending_block_ids), "num_used_blocks": len(self.cache_manager.used_block_ids), + "chunked_prefill": { + "enabled": self.enable_chunked_prefill, + "decode_priority": self.decode_priority, + "max_num_batched_tokens": self.max_num_batched_tokens, + "prefill_chunk_size": self.prefill_chunk_size, + "min_prefill_chunk_size": self.min_prefill_chunk_size, + "max_num_partial_prefills": self.max_num_partial_prefills, + "max_long_partial_prefills": self.max_long_partial_prefills, + "long_prefill_token_threshold": self.long_prefill_token_threshold, + }, + "scheduler": dict(self._stats), } + stats.update(self.cache_manager.get_stats()) if self.mamba_cache_manager is not None: stats.update( { diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index 8f84c9b97..80426e5d9 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -59,8 +59,11 @@ def _is_internal_moe_packed_weight(key: str) -> bool: # InfiniLM registers packed MoE parameters internally. HF checkpoints # provide per-expert gate/up/down weights instead, so these packed tensors # are expected missing keys during non-strict checkpoint loading. - return key.endswith(".mlp.experts.w13_weight") or key.endswith( - ".mlp.experts.w2_weight" + return ( + key.endswith(".mlp.experts.w13_weight") + or key.endswith(".mlp.experts.w2_weight") + or key.endswith(".mlp.vision_experts.w13_weight") + or key.endswith(".mlp.vision_experts.w2_weight") ) @@ -705,6 +708,69 @@ def _remap_videonsa(state_dict, config=None): return state_dict +def _remap_ernie4_5_moe_vl(state_dict, config=None): + """Map ERNIE-4.5-VL-MoE weights to InfiniLM's module layout.""" + text_num_experts = 0 + if config is not None: + moe_num_experts = config.get("moe_num_experts", 0) + if isinstance(moe_num_experts, (list, tuple)): + text_num_experts = int(moe_num_experts[0]) + else: + text_num_experts = int(moe_num_experts) + + remapped = {} + for key, tensor in state_dict.items(): + new_key = key + if new_key.startswith("language_model."): + new_key = "model." + new_key[len("language_model.") :] + elif new_key.startswith("model.language_model."): + new_key = "model." + new_key[len("model.language_model.") :] + elif new_key.startswith("model.resampler_model."): + new_key = "resampler_model." + new_key[len("model.resampler_model.") :] + elif new_key.startswith("model.vision_model."): + new_key = "vision_model." + new_key[len("model.vision_model.") :] + + if new_key.endswith(".mlp.moe_statics.e_score_correction_bias"): + prefix = new_key.removesuffix(".mlp.moe_statics.e_score_correction_bias") + remapped[f"{prefix}.mlp.gate.e_score_correction_bias"] = tensor[ + 0 + ].contiguous() + remapped[f"{prefix}.mlp.vision_gate.e_score_correction_bias"] = tensor[ + 1 + ].contiguous() + continue + + if ".mlp.moe_statics." in new_key: + continue + + if ".mlp.experts." in new_key and text_num_experts > 0: + parts = new_key.split(".") + try: + expert_pos = parts.index("experts") + 1 + expert_idx = int(parts[expert_pos]) + except (ValueError, IndexError): + expert_idx = -1 + if expert_idx >= text_num_experts: + parts[expert_pos - 1] = "vision_experts" + parts[expert_pos] = str(expert_idx - text_num_experts) + new_key = ".".join(parts) + + if new_key.endswith(".mlp.gate.weight"): + remapped[new_key] = tensor.transpose(0, 1).contiguous() + continue + + if new_key.endswith(".mlp.gate.weight_1"): + remapped[ + new_key.removesuffix(".mlp.gate.weight_1") + + ".mlp.vision_gate.weight" + ] = tensor.transpose(0, 1).contiguous() + continue + + remapped[new_key] = tensor + + return remapped + + # Model type → remap function mapping def _remap_qwen3_5(state_dict, config): """Apply Qwen3.5-specific load-time weight fixes.""" @@ -751,4 +817,6 @@ def _remap_qwen3_5(state_dict, config): "mamba": _remap_mamba, "videonsa": _remap_videonsa, "qwen3_5": _remap_qwen3_5, + "ernie4_5_moe_vl": _remap_ernie4_5_moe_vl, + "ernie4_5_vl_moe": _remap_ernie4_5_moe_vl, } diff --git a/python/infinilm/multimodal/multimodal.py b/python/infinilm/multimodal/multimodal.py index 72fea63c6..97fcff149 100644 --- a/python/infinilm/multimodal/multimodal.py +++ b/python/infinilm/multimodal/multimodal.py @@ -1,4 +1,8 @@ +import os from typing import List, Union +from urllib.parse import urlparse + +import numpy as np from PIL import Image @@ -19,6 +23,108 @@ def has_multimodal_inputs(messages: Union[List[dict], dict]) -> bool: return False +def _resolve_local_path(url: str) -> str: + parsed = urlparse(url) + if parsed.scheme == "file": + return os.path.expanduser(parsed.path) + if parsed.scheme in ("", None): + return os.path.expanduser(url) + raise NotImplementedError(f"Unsupported multimodal URL scheme: {parsed.scheme}") + + +def _get_media_url(item: dict, key: str) -> str: + payload = item.get(key) + if isinstance(payload, dict): + payload = payload.get("url") + if not isinstance(payload, str) or payload == "": + raise ValueError(f"Missing {key}.url in multimodal input") + return payload + + +def _sample_frame_indices( + frame_count: int, + fps: float, + *, + target_frames: int = -1, + target_fps: float = 2.0, + min_frames: int = 16, + max_frames: int = 180, + frames_sample: str = "leading", +) -> list[int]: + if frame_count <= 0: + raise ValueError("Video contains no frames") + if fps <= 0: + fps = 1.0 + + if target_frames <= 0: + if target_fps <= 0: + raise ValueError("Either target_frames or target_fps must be positive") + duration = frame_count / fps + frame_target = int(duration * target_fps) + if min_frames > 0 and frame_target < min_frames: + target_frames = min_frames + elif max_frames > 0 and frame_target > max_frames: + target_frames = max_frames + else: + seconds = np.arange(0, duration, 1.0 / target_fps) + indices = np.around(seconds * fps).astype(np.int64).tolist() + return [idx for idx in indices if 0 <= idx < frame_count] or [0] + + samples = min(max(target_frames, 1), frame_count) + intervals = np.linspace(0, frame_count, samples + 1).astype(np.int64) + indices = [] + for start, end in zip(intervals[:-1], intervals[1:]): + end = max(start, end - 1) + if frames_sample == "middle": + indices.append(int((start + end) // 2)) + else: + indices.append(int(start)) + return indices + + +def load_video( + url: str, + *, + target_frames: int = -1, + target_fps: float = 2.0, + min_frames: int = 16, + max_frames: int = 180, + frames_sample: str = "leading", +): + try: + import decord + except ImportError as exc: + raise ImportError("Video input requires decord to be installed") from exc + + path = _resolve_local_path(url) + reader = decord.VideoReader(path, num_threads=1) + frame_count = len(reader) + fps = float(reader.get_avg_fps() or 1.0) + frame_indices = _sample_frame_indices( + frame_count, + fps, + target_frames=target_frames, + target_fps=target_fps, + min_frames=min_frames, + max_frames=max_frames, + frames_sample=frames_sample, + ) + frames = reader.get_batch(frame_indices).asnumpy() + if frames.shape[0] % 2 != 0: + frames = np.concatenate([frames, frames[-1:]], axis=0) + frame_indices.append(frame_indices[-1]) + metadata = { + "fps": fps, + "duration": frame_count / fps, + "num_of_frame": frame_count, + "total_num_frames": frame_count, + "frames_indices": frame_indices, + "video_backend": "decord", + "do_sample_frames": False, + } + return frames, metadata + + def resolve_multimodal_inputs(messages: Union[List[dict], dict]): """Get images, videos, audios from the messages.""" if isinstance(messages, dict): @@ -40,20 +146,38 @@ def resolve_multimodal_inputs(messages: Union[List[dict], dict]): if item.get("type") == "text": pass elif item.get("type") == "image_url": - # TODO support other image url formats - images.append(Image.open(item["image_url"]["url"])) - image_urls.append(item["image_url"]["url"]) + image_url = _get_media_url(item, "image_url") + images.append(Image.open(_resolve_local_path(image_url))) + image_urls.append(image_url) elif item.get("type") == "video_url": - video = item["video_url"]["url"] - videos.append(video) + payload = item.get("video_url") + video = payload.get("url") if isinstance(payload, dict) else payload if isinstance(video, str): + video_args = payload if isinstance(payload, dict) else {} + + def pick_arg(key, default): + return item.get(key, video_args.get(key, default)) + + videos.append( + load_video( + video, + target_frames=int(pick_arg("target_frames", -1)), + target_fps=float(pick_arg("fps", 2.0)), + min_frames=int(pick_arg("min_frames", 16)), + max_frames=int(pick_arg("max_frames", 180)), + frames_sample=str(pick_arg("frames_sample", "leading")), + ) + ) video_urls.append(video) else: + videos.append(video) video_urls.append( f"predecoded_video:{len(video_urls)}:{len(video)}" ) else: # TODO support audio - raise NotImplementedError("Only image/video input is supported for now") + raise NotImplementedError( + "Only image and video inputs are supported for now" + ) return { "images": images, diff --git a/python/infinilm/processors/basic_llm_processor.py b/python/infinilm/processors/basic_llm_processor.py index a6fbc33ac..3f202ae50 100644 --- a/python/infinilm/processors/basic_llm_processor.py +++ b/python/infinilm/processors/basic_llm_processor.py @@ -137,7 +137,7 @@ def _build_model_input_from_static_scheduler_output( input_offsets = [0, 1] return { - "input_ids": infinicore.from_list(input_ids, dtype=infinicore.int64), + "input_ids": infinicore.from_list(input_ids, dtype=infinicore.int32), "position_ids": infinicore.from_list(position_ids, dtype=infinicore.int64), "past_kv_lengths": infinicore.from_list( [past_kv_len], dtype=infinicore.int32 @@ -202,24 +202,26 @@ def _build_model_input_from_batch_scheduler_output( ) current_offset = 0 - for req in scheduler_output.scheduled_requests: + for req_id, req in enumerate(scheduler_output.scheduled_requests): num_cached = req.num_local_cached_tokens - if scheduler_output.is_prefill: + if scheduler_output.request_is_prefill(req_id): # Prefill phase req_tokens = req.get_input_tokens() - tokens_to_compute = req_tokens[num_cached:] + chunk_start = req.prefill_chunk_start + chunk_end = req.prefill_chunk_end or len(req_tokens) + tokens_to_compute = req_tokens[chunk_start:chunk_end] tokens.extend(tokens_to_compute) compute_len = len(tokens_to_compute) - seq_len = len(req_tokens) + seq_len = chunk_end seq_lens.append(seq_len) current_offset += compute_len seq_offsets.append(current_offset) slot_mapping.extend(req.slot_mapping) - cached_lens.append(num_cached) - position_ids.extend(range(num_cached, num_cached + compute_len)) + cached_lens.append(chunk_start) + position_ids.extend(range(chunk_start, chunk_end)) else: # Decode phase @@ -247,7 +249,7 @@ def _build_model_input_from_batch_scheduler_output( cu_seqlens.append(cu_seqlens[-1] + seq_len) return { - "input_ids": infinicore.from_list([tokens], dtype=infinicore.int64), + "input_ids": infinicore.from_list([tokens], dtype=infinicore.int32), "position_ids": infinicore.from_list(position_ids, dtype=infinicore.int64), "past_kv_lengths": infinicore.from_list( cached_lens, dtype=infinicore.int32 diff --git a/python/infinilm/processors/ernie4_5_vl_processor.py b/python/infinilm/processors/ernie4_5_vl_processor.py new file mode 100644 index 000000000..5365e047d --- /dev/null +++ b/python/infinilm/processors/ernie4_5_vl_processor.py @@ -0,0 +1,317 @@ +import json +from pathlib import Path + +import torch +from transformers import AutoProcessor + +from .basic_llm_processor import BasicLLMProcessor +from .processor import InfinilmProcessor, register_processor + + +_IMAGE_PLACEHOLDER = "<|IMAGE_START|><|image@placeholder|><|IMAGE_END|>" +_VIDEO_PLACEHOLDER = "<|VIDEO_START|><|video@placeholder|><|VIDEO_END|>" + + +@register_processor("ernie4_5_moe_vl") +@register_processor("ernie4_5_vl_moe") +class Ernie4_5VLProcessor(InfinilmProcessor): + def __init__(self, model_dir_path: str): + self.model_dir_path = model_dir_path + self.processor = AutoProcessor.from_pretrained( + model_dir_path, trust_remote_code=True + ) + if hasattr(self.processor, "eval"): + self.processor.eval() + self.tokenizer = self.processor.tokenizer + self.image_processor = self.processor.image_processor + + # Keep the initial VL path bounded while the vision tower runs in eager kernels. + max_patches = 1024 + if hasattr(self.processor, "image_max_pixels"): + self.processor.image_max_pixels = max_patches * 28 * 28 + + with open(Path(model_dir_path) / "config.json", "r") as f: + self.config = json.load(f) + self.image_patch_id = int(self.config["im_patch_id"]) + + def _ensure_media_placeholders( + self, prompt: str, has_images: bool, has_videos: bool + ) -> str: + prefix = "" + if has_images and "<|image@placeholder|>" not in prompt: + prefix += _IMAGE_PLACEHOLDER + if has_videos and "<|video@placeholder|>" not in prompt: + prefix += _VIDEO_PLACEHOLDER + if prefix == "": + return prompt + return f"User: {prefix}{prompt}\nAssistant:" + + def __call__( + self, + prompt, + images=None, + videos=None, + audios=None, + return_tensors: str = None, + **kwargs, + ) -> dict: + if audios: + raise NotImplementedError("ERNIE-4.5-VL audio input is not wired yet") + + if images or videos: + prompt = self._ensure_media_placeholders( + prompt, bool(images), bool(videos) + ) + result = self.processor( + text=[prompt], + images=images, + videos=videos, + padding=True, + return_tensors="pt", + **kwargs, + ) + result = dict(result) + result["image_bound"] = self._build_image_bound(result["input_ids"]) + self._prepare_visual_outputs(result) + return self._convert_return_tensors(result, return_tensors) + + if return_tensors is None: + return self.tokenizer(prompt, add_special_tokens=False) + if return_tensors == "pt": + return self.tokenizer( + prompt, return_tensors="pt", add_special_tokens=False + ) + if return_tensors == "infini": + return self._convert_return_tensors( + self.tokenizer(prompt, return_tensors="pt", add_special_tokens=False), + return_tensors, + ) + return self.tokenizer( + prompt, return_tensors=return_tensors, add_special_tokens=False + ) + + def _normalize_pixel_values(self, pixel_values: torch.Tensor) -> torch.Tensor: + if not torch.is_tensor(pixel_values): + pixel_values = torch.tensor(pixel_values) + image_mean = torch.tensor( + self.image_processor.image_mean, dtype=torch.float32 + ).reshape(1, 3) + image_std = torch.tensor( + self.image_processor.image_std, dtype=torch.float32 + ).reshape(1, 3) + patch_size = int(self.config["vision_config"]["patch_size"]) + patch_size_squared = patch_size * patch_size + image_mean = image_mean.repeat_interleave(patch_size_squared, dim=-1) + image_std = image_std.repeat_interleave(patch_size_squared, dim=-1) + rescale_factor = float(self.image_processor.rescale_factor) + pixel_values = ( + rescale_factor * pixel_values.to(torch.float32) - image_mean + ) / image_std + return pixel_values.to(torch.bfloat16).contiguous() + + def _prepare_visual_outputs(self, result: dict) -> None: + pixel_values = self._pop_visual_tensor( + result, ["images", "pixel_values", "pixel_values_videos"] + ) + if pixel_values is not None: + result["pixel_values"] = self._normalize_pixel_values(pixel_values) + + grid_thw = self._pop_visual_tensor( + result, ["grid_thw", "image_grid_thw", "video_grid_thw"] + ) + if grid_thw is not None: + result["tgt_sizes"] = grid_thw.to(torch.int64).contiguous() + + @staticmethod + def _pop_visual_tensor(result: dict, keys: list[str]): + tensors = [] + for key in keys: + value = result.pop(key, None) + if value is None: + continue + if not torch.is_tensor(value): + value = torch.tensor(value) + tensors.append(value) + if not tensors: + return None + if len(tensors) == 1: + return tensors[0] + return torch.cat(tensors, dim=0) + + def _build_image_bound(self, input_ids: torch.Tensor) -> torch.Tensor: + ids = input_ids[0].tolist() + ranges = [] + start = None + for idx, token_id in enumerate(ids): + if token_id == self.image_patch_id: + if start is None: + start = idx + elif start is not None: + ranges.append([start, idx]) + start = None + if start is not None: + ranges.append([start, len(ids)]) + if not ranges: + raise ValueError("ERNIE VL input did not contain visual patch tokens") + return torch.tensor([ranges], dtype=torch.int64) + + def _convert_return_tensors(self, result: dict, return_tensors: str): + if return_tensors != "infini": + return result + import infinicore + + converted = {} + for key, value in result.items(): + if isinstance(value, torch.Tensor): + converted[key] = infinicore.from_torch(value.contiguous()) + else: + converted[key] = value + return converted + + def apply_chat_template( + self, + conversation, + add_generation_prompt: bool = False, + tokenize: bool = True, + **kwargs, + ): + return self.tokenizer.apply_chat_template( + conversation=conversation, + add_generation_prompt=add_generation_prompt, + tokenize=tokenize, + **kwargs, + ) + + def build_model_inputs( + self, + scheduler_output, + temperature: float = 1.0, + top_p: float = 0.8, + top_k: int = 1, + **kwargs, + ) -> dict: + helper = BasicLLMProcessor.__new__(BasicLLMProcessor) + helper.tokenizer = self.tokenizer + model_inputs = BasicLLMProcessor.build_model_inputs( + helper, scheduler_output, temperature, top_p, top_k, **kwargs + ) + if not scheduler_output.is_prefill: + import infinicore + import torch + + mrope_position_ids = [] + has_mrope_position_ids = False + for req in scheduler_output.scheduled_requests: + processed_inputs = req.processed_inputs + if processed_inputs is not None and "position_ids" in processed_inputs: + pos = processed_inputs["position_ids"] + if pos.ndim == 3: + pos = pos[0] + decode_pos = int(pos.max().item()) + len(req.generated_token_ids) + mrope_position_ids.append( + torch.tensor([decode_pos, decode_pos, decode_pos], dtype=torch.int64) + ) + has_mrope_position_ids = True + else: + decode_pos = req.get_total_length() - 1 + mrope_position_ids.append( + torch.tensor([decode_pos, decode_pos, decode_pos], dtype=torch.int64) + ) + if has_mrope_position_ids: + model_inputs["position_ids"] = infinicore.from_torch( + torch.vstack(mrope_position_ids).contiguous() + ) + return model_inputs + + import infinicore + import torch + + pixel_values = [] + image_bound = [] + tgt_sizes = [] + image_req_ids = [] + mrope_position_ids = [] + has_mrope_position_ids = False + for req_id, req in enumerate(scheduler_output.scheduled_requests): + processed_inputs = req.processed_inputs + num_cached = req.num_local_cached_tokens + tokens_to_compute = req.get_input_tokens()[num_cached:] + if processed_inputs is not None and "position_ids" in processed_inputs: + pos = processed_inputs["position_ids"] + if pos.ndim == 3: + pos = pos[0] + mrope_position_ids.append(pos[num_cached : num_cached + len(tokens_to_compute)].to(torch.int64)) + has_mrope_position_ids = True + else: + mrope_position_ids.append( + torch.tensor( + [[pos, pos, pos] for pos in range(num_cached, num_cached + len(tokens_to_compute))], + dtype=torch.int64, + ) + ) + + if processed_inputs is None or "pixel_values" not in processed_inputs: + continue + if num_cached != 0: + raise NotImplementedError( + "ERNIE-4.5-VL prefix-cache reuse with image inputs is not wired yet" + ) + pixel_values.append( + infinicore.from_torch(processed_inputs["pixel_values"].contiguous()) + ) + image_bound.append( + infinicore.from_torch(processed_inputs["image_bound"].contiguous()) + ) + tgt_sizes.append( + infinicore.from_torch(processed_inputs["tgt_sizes"].contiguous()) + ) + image_req_ids.append(req_id) + + if pixel_values: + model_inputs["pixel_values"] = pixel_values + model_inputs["image_bound"] = image_bound + model_inputs["tgt_sizes"] = tgt_sizes + model_inputs["image_req_ids"] = image_req_ids + if has_mrope_position_ids: + model_inputs["position_ids"] = infinicore.from_torch( + torch.vstack(mrope_position_ids).contiguous() + ) + return model_inputs + + def get_tokenizer(self): + return self.tokenizer + + def get_mm_token_index_list( + self, prompt_token_ids, image_ids=None, video_ids=None, audio_ids=None, **kwargs + ): + if audio_ids: + raise NotImplementedError( + "ERNIE-4.5-VL audio service input is not wired yet" + ) + image_ids = image_ids or [] + video_ids = video_ids or [] + media_ids = list(image_ids) + list(video_ids) + ranges = [] + start = None + for idx, token_id in enumerate(prompt_token_ids): + if token_id == self.image_patch_id: + if start is None: + start = idx + elif start is not None: + ranges.append((start, idx)) + start = None + if start is not None: + ranges.append((start, len(prompt_token_ids))) + + if len(ranges) != len(media_ids): + raise ValueError( + "The number of ERNIE media patch ranges does not match the number of media inputs" + ) + return [ + { + "start_index": start, + "end_index": end, + "identifier": media_ids[idx], + } + for idx, (start, end) in enumerate(ranges) + ] diff --git a/python/infinilm/processors/minicpmv_processor.py b/python/infinilm/processors/minicpmv_processor.py index e4f879c9c..27b853fc2 100644 --- a/python/infinilm/processors/minicpmv_processor.py +++ b/python/infinilm/processors/minicpmv_processor.py @@ -130,22 +130,24 @@ def build_model_inputs( for req_id, req in enumerate(scheduler_output.scheduled_requests): num_cached = req.num_local_cached_tokens - if scheduler_output.is_prefill: + if scheduler_output.request_is_prefill(req_id): # Prefill phase req_tokens = req.get_input_tokens() - tokens_to_compute = req_tokens[num_cached:] + chunk_start = req.prefill_chunk_start + chunk_end = req.prefill_chunk_end or len(req_tokens) + tokens_to_compute = req_tokens[chunk_start:chunk_end] tokens.extend(tokens_to_compute) compute_len = len(tokens_to_compute) - seq_len = len(req_tokens) + seq_len = chunk_end seq_lens.append(seq_len) current_offset += compute_len seq_offsets.append(current_offset) slot_mapping.extend(req.slot_mapping) - cached_lens.append(num_cached) - position_ids.extend(range(num_cached, num_cached + compute_len)) + cached_lens.append(chunk_start) + position_ids.extend(range(chunk_start, chunk_end)) if ( req.processed_inputs is not None @@ -154,7 +156,7 @@ def build_model_inputs( import torch num_cached_patch = ( - (req.processed_inputs["image_bound"][0][:, 1] <= num_cached) + (req.processed_inputs["image_bound"][0][:, 1] <= chunk_start) .sum() .item() ) @@ -215,11 +217,8 @@ def build_model_inputs( image_bound_infini = infinicore.from_torch(bound) - def append_mm_data(mm_data__: dict, key__: str, value__): - if mm_data__.get(key__) is None: - mm_data[key__] = [value__] - else: - mm_data[key__].append(value__) + def append_mm_data(target: dict, key: str, value): + target.setdefault(key, []).append(value) append_mm_data(mm_data, "pixel_values", pixel_values_infini) append_mm_data(mm_data, "tgt_sizes", tgt_sizes_infini) diff --git a/python/infinilm/server/inference_server.py b/python/infinilm/server/inference_server.py index 242593e25..4c62f159c 100644 --- a/python/infinilm/server/inference_server.py +++ b/python/infinilm/server/inference_server.py @@ -100,9 +100,18 @@ def __init__( tensor_parallel_size: int = 1, moe_ep_backend: str = "disabled", moe_ep_size: int = 1, + allreduce_backend: str = "nccl", cache_type: str = "paged", max_tokens: int = 4096, max_batch_size: int = 16, + max_num_batched_tokens: Optional[int] = None, + enable_chunked_prefill: bool = False, + prefill_chunk_size: Optional[int] = None, + decode_priority: Optional[bool] = None, + max_num_partial_prefills: int = 1, + max_long_partial_prefills: int = 1, + long_prefill_token_threshold: Optional[int] = None, + min_prefill_chunk_size: Optional[int] = None, num_blocks: int = 512, block_size: int = 256, max_cache_len: int = 4096, @@ -130,6 +139,14 @@ def __init__( cache_type: Cache type ('paged' or 'static'). max_tokens: Default maximum tokens to generate. max_batch_size: Maximum batch size for inference (only for paged cache). + max_num_batched_tokens: Token budget per scheduler step. + enable_chunked_prefill: Whether to split long prefills across steps. + prefill_chunk_size: Maximum tokens per prefill chunk. + decode_priority: Whether to schedule decode before waiting prefill chunks. + max_num_partial_prefills: Maximum partial prefill chunks per scheduler step. + max_long_partial_prefills: Maximum long partial prefill chunks per step. + long_prefill_token_threshold: Prompt length threshold for long partial prefills. + min_prefill_chunk_size: Minimum non-final mixed prefill chunk size. num_blocks: Number of KV cache blocks (only for paged cache). block_size: Size of each KV cache block (only for paged cache). max_cache_len: Maximum sequence length (only for static cache). @@ -153,9 +170,18 @@ def __init__( self.tensor_parallel_size = tensor_parallel_size self.moe_ep_backend = moe_ep_backend self.moe_ep_size = moe_ep_size + self.allreduce_backend = allreduce_backend self.cache_type = cache_type self.max_tokens = max_tokens self.max_batch_size = max_batch_size + self.max_num_batched_tokens = max_num_batched_tokens + self.enable_chunked_prefill = enable_chunked_prefill + self.prefill_chunk_size = prefill_chunk_size + self.decode_priority = decode_priority + self.max_num_partial_prefills = max_num_partial_prefills + self.max_long_partial_prefills = max_long_partial_prefills + self.long_prefill_token_threshold = long_prefill_token_threshold + self.min_prefill_chunk_size = min_prefill_chunk_size self.num_blocks = num_blocks self.block_size = block_size self.max_cache_len = max_cache_len @@ -192,9 +218,18 @@ async def lifespan(app: FastAPI): tensor_parallel_size=self.tensor_parallel_size, moe_ep_backend=self.moe_ep_backend, moe_ep_size=self.moe_ep_size, + allreduce_backend=self.allreduce_backend, cache_type=self.cache_type, max_batch_size=self.max_batch_size, max_tokens=self.max_tokens, + max_num_batched_tokens=self.max_num_batched_tokens, + enable_chunked_prefill=self.enable_chunked_prefill, + prefill_chunk_size=self.prefill_chunk_size, + decode_priority=self.decode_priority, + max_num_partial_prefills=self.max_num_partial_prefills, + max_long_partial_prefills=self.max_long_partial_prefills, + long_prefill_token_threshold=self.long_prefill_token_threshold, + min_prefill_chunk_size=self.min_prefill_chunk_size, num_blocks=self.num_blocks, block_size=self.block_size, max_cache_len=self.max_cache_len, @@ -268,6 +303,13 @@ async def health(): return JSONResponse(content={"status": "unhealthy"}, status_code=503) return {"status": "healthy"} + @app.get("/cache_stats") + @app.get("/v1/cache_stats") + async def cache_stats(): + if self.engine is None: + return JSONResponse(content={"error": "engine not initialized"}, status_code=503) + return self.engine.get_cache_stats() + def _models_payload(): return { "object": "list", @@ -607,9 +649,18 @@ def main(): tensor_parallel_size=cfg.tp, moe_ep_backend=moe_ep_backend, moe_ep_size=ep, + allreduce_backend=cfg.allreduce_backend, cache_type="paged" if cfg.enable_paged_attn else "static", max_tokens=cfg.max_new_tokens, max_batch_size=cfg.max_batch_size, + max_num_batched_tokens=cfg.max_num_batched_tokens, + enable_chunked_prefill=cfg.enable_chunked_prefill, + prefill_chunk_size=cfg.prefill_chunk_size, + decode_priority=cfg.decode_priority, + max_num_partial_prefills=cfg.max_num_partial_prefills, + max_long_partial_prefills=cfg.max_long_partial_prefills, + long_prefill_token_threshold=cfg.long_prefill_token_threshold, + min_prefill_chunk_size=cfg.min_prefill_chunk_size, num_blocks=cfg.num_blocks, block_size=cfg.block_size, max_cache_len=cfg.max_cache_len, diff --git a/test/llm/test_phase_aware_execution.py b/test/llm/test_phase_aware_execution.py new file mode 100644 index 000000000..2ea0f3e5d --- /dev/null +++ b/test/llm/test_phase_aware_execution.py @@ -0,0 +1,291 @@ +from __future__ import annotations + +import importlib.util +import sys +import types +import unittest +from dataclasses import dataclass +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +@dataclass +class FakeRequest: + request_id: str + prefill_chunk_start: int = 0 + prefill_chunk_end: int = 0 + + +def _scheduler_stubs(): + infinilm_pkg = types.ModuleType("infinilm") + llm_pkg = types.ModuleType("infinilm.llm") + request_mod = types.ModuleType("infinilm.llm.request") + cache_manager_mod = types.ModuleType("infinilm.llm.cache_manager") + + class RequestStatus: + WAITING = "waiting" + RUNNING = "running" + WAITING_FOR_REMOTE_KVS = "waiting_for_remote_kvs" + FINISHED = "finished" + CANCELED = "canceled" + FAILED = "failed" + TIMEOUT = "timeout" + + class BlockManager: + def __init__(self, *args, **kwargs): + raise AssertionError("BlockManager is not used by these tests") + + class MambaCacheManager: + def __init__(self, *args, **kwargs): + raise AssertionError("MambaCacheManager is not used by these tests") + + request_mod.RequestStatus = RequestStatus + request_mod.InferenceRequest = object + cache_manager_mod.BlockManager = BlockManager + cache_manager_mod.MambaCacheManager = MambaCacheManager + + return { + "infinilm": infinilm_pkg, + "infinilm.llm": llm_pkg, + "infinilm.llm.request": request_mod, + "infinilm.llm.cache_manager": cache_manager_mod, + } + + +def _model_runner_stubs(): + stubs = { + "infinicore": types.ModuleType("infinicore"), + "infinilm.distributed": types.ModuleType("infinilm.distributed"), + "infinilm.infer_engine": types.ModuleType("infinilm.infer_engine"), + "infinilm.cache": types.ModuleType("infinilm.cache"), + "infinilm.cache.cache": types.ModuleType("infinilm.cache.cache"), + "infinilm.modeling_utils": types.ModuleType("infinilm.modeling_utils"), + "infinilm.config": types.ModuleType("infinilm.config"), + "infinilm.config.engine_config": types.ModuleType( + "infinilm.config.engine_config" + ), + "infinilm.kv_connector": types.ModuleType("infinilm.kv_connector"), + "infinilm.processors": types.ModuleType("infinilm.processors"), + } + stubs["infinilm.distributed"].DistConfig = object + stubs["infinilm.infer_engine"].InferEngine = object + stubs["infinilm.cache.cache"].PagedKVCacheConfig = object + stubs["infinilm.cache.cache"].StaticKVCacheConfig = object + stubs["infinilm.modeling_utils"].load_model_state_dict_by_file = ( + lambda *args, **kwargs: None + ) + stubs["infinilm.config.engine_config"].EngineConfig = object + stubs["infinilm.kv_connector"].KVConnectorRole = SimpleNamespace(WORKER=1) + stubs["infinilm.kv_connector"].KVConnectorFactory = object + stubs["infinilm.processors"].AutoInfinilmProcessor = object + return stubs + + +def _load_module(module_name: str, relative_path: str, stubs: dict[str, types.ModuleType]): + sys.modules.pop(module_name, None) + spec = importlib.util.spec_from_file_location( + module_name, REPO_ROOT / relative_path + ) + module = importlib.util.module_from_spec(spec) + with patch.dict(sys.modules, stubs): + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +def load_scheduler_module(): + return _load_module( + "_phase_test_scheduler", + "python/infinilm/llm/scheduler.py", + _scheduler_stubs(), + ) + + +def load_model_runner_module(): + stubs = _scheduler_stubs() + stubs.update(_model_runner_stubs()) + return _load_module( + "_phase_test_model_runner", + "python/infinilm/llm/model_runner/model_runner.py", + stubs, + ) + + +class PhaseAwareExecutionTest(unittest.TestCase): + def test_scheduler_output_builds_decode_first_phase_plan(self): + scheduler_module = load_scheduler_module() + output = scheduler_module.SchedulerOutput( + scheduled_requests=[ + FakeRequest("decode-0"), + FakeRequest( + "prefill-0", prefill_chunk_start=32, prefill_chunk_end=96 + ), + FakeRequest("decode-1"), + ], + scheduled_prefill_flags=[False, True, False], + ) + + self.assertTrue(output.is_mixed) + self.assertTrue(output.is_prefill) + self.assertEqual(output.num_decode_tokens, 2) + self.assertEqual(output.num_prefill_tokens, 64) + self.assertEqual(output.num_batched_tokens, 66) + self.assertEqual( + [ + (phase.kind, phase.request_indices, phase.num_tokens) + for phase in output.execution_phases + ], + [ + ("decode", (0, 2), 2), + ("prefill", (1,), 64), + ], + ) + + decode_subset = output.make_subset((0, 2)) + self.assertFalse(decode_subset.is_prefill) + self.assertEqual( + [req.request_id for req in decode_subset.scheduled_requests], + ["decode-0", "decode-1"], + ) + self.assertEqual(decode_subset.num_decode_tokens, 2) + self.assertEqual(decode_subset.num_prefill_tokens, 0) + + prefill_subset = output.make_subset((1,)) + self.assertTrue(prefill_subset.is_prefill) + self.assertEqual( + [req.request_id for req in prefill_subset.scheduled_requests], + ["prefill-0"], + ) + self.assertEqual(prefill_subset.num_decode_tokens, 0) + self.assertEqual(prefill_subset.num_prefill_tokens, 64) + + def test_model_runner_splits_mixed_batch_and_restores_output_order(self): + scheduler_module = load_scheduler_module() + model_runner_module = load_model_runner_module() + output = scheduler_module.SchedulerOutput( + scheduled_requests=[ + FakeRequest("decode-0"), + FakeRequest("prefill-0", prefill_chunk_start=0, prefill_chunk_end=128), + FakeRequest("decode-1"), + ], + scheduled_prefill_flags=[False, True, False], + ) + + runner = object.__new__(model_runner_module.ModelRunner) + runner.config = SimpleNamespace(cache_type="paged") + calls = [] + + def fake_forward_single(sub_output): + calls.append( + { + "is_prefill": sub_output.is_prefill, + "request_ids": [ + req.request_id for req in sub_output.scheduled_requests + ], + "decode_tokens": sub_output.num_decode_tokens, + "prefill_tokens": sub_output.num_prefill_tokens, + } + ) + return [100 + len(calls)] * sub_output.num_requests + + runner._model_forward_single = fake_forward_single + + sampled_tokens, stats = runner._model_forward(output) + + self.assertEqual( + calls, + [ + { + "is_prefill": False, + "request_ids": ["decode-0", "decode-1"], + "decode_tokens": 2, + "prefill_tokens": 0, + }, + { + "is_prefill": True, + "request_ids": ["prefill-0"], + "decode_tokens": 0, + "prefill_tokens": 128, + }, + ], + ) + self.assertEqual(sampled_tokens, [101, 102, 101]) + self.assertEqual( + stats, + { + "split_mixed_batch": True, + "phases": [ + {"kind": "decode", "num_requests": 2, "num_tokens": 2}, + {"kind": "prefill", "num_requests": 1, "num_tokens": 128}, + ], + }, + ) + + def test_model_runner_keeps_static_cache_mixed_batch_unsplit(self): + scheduler_module = load_scheduler_module() + model_runner_module = load_model_runner_module() + output = scheduler_module.SchedulerOutput( + scheduled_requests=[ + FakeRequest("decode-0"), + FakeRequest("prefill-0", prefill_chunk_start=0, prefill_chunk_end=16), + ], + scheduled_prefill_flags=[False, True], + ) + + runner = object.__new__(model_runner_module.ModelRunner) + runner.config = SimpleNamespace(cache_type="static") + calls = [] + + def fake_forward_single(sub_output): + calls.append([req.request_id for req in sub_output.scheduled_requests]) + return [7] * sub_output.num_requests + + runner._model_forward_single = fake_forward_single + + sampled_tokens, stats = runner._model_forward(output) + + self.assertEqual(calls, [["decode-0", "prefill-0"]]) + self.assertEqual(sampled_tokens, [7, 7]) + self.assertFalse(stats["split_mixed_batch"]) + self.assertEqual( + stats["phases"], + [ + {"kind": "decode", "num_requests": 1, "num_tokens": 1}, + {"kind": "prefill", "num_requests": 1, "num_tokens": 16}, + ], + ) + + def test_scheduler_records_phase_execution_stats(self): + scheduler_module = load_scheduler_module() + scheduler = object.__new__(scheduler_module.Scheduler) + scheduler._stats = { + "execution_split_mixed_steps": 0, + "execution_decode_phases": 0, + "execution_prefill_phases": 0, + "execution_decode_phase_tokens": 0, + "execution_prefill_phase_tokens": 0, + } + + scheduler._record_execution_stats( + { + "split_mixed_batch": True, + "phases": [ + {"kind": "decode", "num_tokens": 2}, + {"kind": "prefill", "num_tokens": 128}, + ], + } + ) + + self.assertEqual(scheduler._stats["execution_split_mixed_steps"], 1) + self.assertEqual(scheduler._stats["execution_decode_phases"], 1) + self.assertEqual(scheduler._stats["execution_prefill_phases"], 1) + self.assertEqual(scheduler._stats["execution_decode_phase_tokens"], 2) + self.assertEqual(scheduler._stats["execution_prefill_phase_tokens"], 128) + + +if __name__ == "__main__": + unittest.main() diff --git a/xmake.lua b/xmake.lua index aab1a0c70..88effb5bc 100644 --- a/xmake.lua +++ b/xmake.lua @@ -50,6 +50,9 @@ target("_infinilm") add_linkdirs(INFINI_ROOT.."/lib") add_links("infinicore_cpp_api", "infiniop", "infinirt", "infiniccl") + if is_plat("linux") then + add_links("dl") + end -- Add src files add_files("csrc/**.cpp")