diff --git a/README.md b/README.md index 4c1eaf1e1..bac239816 100644 --- a/README.md +++ b/README.md @@ -76,6 +76,8 @@ ```bash CUDA_VISIBLE_DEVICES=0,1,2,3 python python/infinilm/server/inference_server.py --device nvidia --model=/models/9G7B_MHA/ --enable-paged-attn --attn=flash-attn --enable-graph ``` + + - Qwen3-MoE 执行后端可通过 `--moe-backend=auto|legacy|fused` 选择；默认 `auto` 仅对支持的 CUDA FP16/BF16 非量化模型启用 fused 路径。 - 测试推理服务性能： ```bash diff --git a/csrc/cache/kv_cache.cpp b/csrc/cache/kv_cache.cpp index 88d7367e9..1d30b72de 100644 --- a/csrc/cache/kv_cache.cpp +++ b/csrc/cache/kv_cache.cpp @@ -77,11 +77,19 @@ infinicore::Tensor create_layer_kv_cache( // ========================== // PagedKVCacheConfig // ========================== +// Preserve the legacy Graph capture ceiling for two-argument C++ callers. PagedKVCacheConfig::PagedKVCacheConfig( size_t num_blocks, size_t block_size) + : PagedKVCacheConfig(num_blocks, block_size, 512) {} + +PagedKVCacheConfig::PagedKVCacheConfig( + size_t num_blocks, + size_t block_size, + size_t max_batch_size) : num_blocks_(num_blocks), - block_size_(block_size) { + block_size_(block_size), + max_batch_size_(max_batch_size) { } std::unique_ptr @@ -99,6 +107,11 @@ PagedKVCacheConfig::block_size() const { return block_size_; } +size_t +PagedKVCacheConfig::max_batch_size() const { + return max_batch_size_; +} + namespace PagedKVCache { // ========================== // PagedKVCache diff --git a/csrc/cache/kv_cache.hpp b/csrc/cache/kv_cache.hpp index 4d0a9a704..7ca94a586 100644 --- a/csrc/cache/kv_cache.hpp +++ b/csrc/cache/kv_cache.hpp @@ -40,14 +40,20 @@ class PagedKVCacheConfig final : public CacheConfig { PagedKVCacheConfig( size_t num_blocks, size_t block_size = 256); + PagedKVCacheConfig( + size_t num_blocks, + size_t block_size, + size_t max_batch_size); std::unique_ptr unique_copy() const override; size_t num_blocks() const; size_t block_size() const; + size_t max_batch_size() const; private: size_t num_blocks_; size_t block_size_; + size_t max_batch_size_; }; namespace PagedKVCache { diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp index f267794e2..c170a96b9 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -5,7 +5,20 @@ namespace infinilm::engine { PagedCompiler::PagedCompiler(const std::shared_ptr &model, RankBarrier *barrier) - : GraphCompiler(model, barrier) { + : GraphCompiler(model, barrier) {} + +void PagedCompiler::compile() { + compiled_map_decode_.clear(); + const auto *config = dynamic_cast(model_->get_cache_config()); + if (config == nullptr) { + return; + } + const size_t max_batch_size = config->max_batch_size(); + if (max_batch_size == 0) { + throw std::invalid_argument("Paged Graph max_batch_size must be greater than zero"); + } + + decode_batch_sizes_.clear(); for (size_t b = 1; b < 64; ++b) { decode_batch_sizes_.push_back(b); } @@ -18,86 +31,98 @@ PagedCompiler::PagedCompiler(const std::shared_ptr &model, RankBa for (size_t b = 256; b <= 512; b += 64) { decode_batch_sizes_.push_back(b); } -} - -void PagedCompiler::compile() { - if (model_->get_cache_config() != nullptr && dynamic_cast(model_->get_cache_config())) { - 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(); - 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) { - 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.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()); - set_zeros(input.total_sequence_lengths.value()); - std::vector total_sequence_lengths_vec(b, 1); - infinicore::context::memcpyH2D(input.total_sequence_lengths.value()->data(), total_sequence_lengths_vec.data(), b * sizeof(int32_t), false); - input.input_offsets = infinicore::Tensor::empty({b + 1}, infinicore::DataType::I32, infinicore::context::getDevice()); - std::vector input_offsets_vec(b + 1, 0); - for (size_t i = 0; i <= b; i++) { - input_offsets_vec[i] = i; - } - infinicore::context::memcpyH2D(input.input_offsets.value()->data(), input_offsets_vec.data(), (b + 1) * sizeof(int32_t), false); - input.cu_seqlens = infinicore::Tensor::empty({b + 1}, infinicore::DataType::I32, infinicore::context::getDevice()); - infinicore::context::memcpyH2D(input.cu_seqlens.value()->data(), input_offsets_vec.data(), (b + 1) * sizeof(int32_t), false); - const size_t block_per_req = nblocks; - input.block_tables = block_tables_holder_->as_strided({b, block_per_req}, {(ptrdiff_t)block_per_req, 1}); - input.slot_mapping = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice()); - set_zeros(input.slot_mapping.value()); - - // Attention reads attn_metadata from thread-local forward context. - infinilm::global_state::get_forward_context().attn_metadata = { - input.past_sequence_lengths, - input.total_sequence_lengths, - input.input_offsets, - input.cu_seqlens, - input.block_tables, - input.slot_mapping, - }; - return input; - }; + decode_batch_sizes_.erase( + std::remove_if(decode_batch_sizes_.begin(), decode_batch_sizes_.end(), + [max_batch_size](size_t b) { return b > max_batch_size; }), + decode_batch_sizes_.end()); + if (decode_batch_sizes_.empty() || decode_batch_sizes_.back() != max_batch_size) { + decode_batch_sizes_.push_back(max_batch_size); + } - { - 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); - infinicore::context::syncStream(); - // Warmup runs the eager Marlin path and may leave per-layer lock - // workspaces dirty. Reset before CUDA graph capture so capture - // starts from the same all-zero lock state as normal execution. - model_->reset_runtime_state(); - infinicore::context::syncStream(); + // Keep 1-16 dense, retain multiples of eight from the legacy bucket list, + // and always capture the configured maximum. + if (max_batch_size > 16) { + decode_batch_sizes_.erase( + std::remove_if(decode_batch_sizes_.begin(), decode_batch_sizes_.end(), + [max_batch_size](size_t b) { + return b > 16 && b != max_batch_size && b % 8 != 0; + }), + decode_batch_sizes_.end()); + } + size_t nblocks = config->num_blocks(); + 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) { + InfinilmModel::Input input; + input.use_local_vocab_logits = true; + 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.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()); + set_zeros(input.total_sequence_lengths.value()); + std::vector total_sequence_lengths_vec(b, 1); + infinicore::context::memcpyH2D(input.total_sequence_lengths.value()->data(), total_sequence_lengths_vec.data(), b * sizeof(int32_t), false); + input.input_offsets = infinicore::Tensor::empty({b + 1}, infinicore::DataType::I32, infinicore::context::getDevice()); + std::vector input_offsets_vec(b + 1, 0); + for (size_t i = 0; i <= b; i++) { + input_offsets_vec[i] = i; } + infinicore::context::memcpyH2D(input.input_offsets.value()->data(), input_offsets_vec.data(), (b + 1) * sizeof(int32_t), false); + input.cu_seqlens = infinicore::Tensor::empty({b + 1}, infinicore::DataType::I32, infinicore::context::getDevice()); + infinicore::context::memcpyH2D(input.cu_seqlens.value()->data(), input_offsets_vec.data(), (b + 1) * sizeof(int32_t), false); + const size_t block_per_req = nblocks; + input.block_tables = block_tables_holder_->as_strided({b, block_per_req}, {(ptrdiff_t)block_per_req, 1}); + input.slot_mapping = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice()); + set_zeros(input.slot_mapping.value()); + + // Attention reads attn_metadata from thread-local forward context. + infinilm::global_state::get_forward_context().attn_metadata = { + input.past_sequence_lengths, + input.total_sequence_lengths, + input.input_offsets, + input.cu_seqlens, + input.block_tables, + input.slot_mapping, + }; + return input; + }; + + { + 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); + infinicore::context::syncStream(); + // Warmup runs the eager Marlin path and may leave per-layer lock + // workspaces dirty. Reset before CUDA graph capture so capture + // starts from the same all-zero lock state as normal execution. + model_->reset_runtime_state(); + infinicore::context::syncStream(); + } - for (size_t b : decode_batch_sizes_) { - auto input = make_decode_input(b); - - barrier_->wait(); - (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 - // graph capture; the current implementation still pays a memset - // before every graph replay in get_compiled(). - model_->reset_runtime_state(); - infinicore::context::syncStream(); - infinicore::context::startGraphRecording(); - auto output = model_->forward(input); - auto graph = infinicore::context::stopGraphRecording(); - barrier_->wait(); - - 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)}; - } + for (size_t b : decode_batch_sizes_) { + auto input = make_decode_input(b); + + barrier_->wait(); + (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 + // graph capture; the current implementation still pays a memset + // before every graph replay in get_compiled(). + model_->reset_runtime_state(); + infinicore::context::syncStream(); + infinicore::context::startGraphRecording(); + auto output = model_->forward(input); + auto graph = infinicore::context::stopGraphRecording(); + barrier_->wait(); + + 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)}; } } @@ -106,8 +131,10 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & size_t batch_size = input.block_tables.value()->size(0); size_t block_per_req = input.block_tables.value()->size(1); - // only support decode only batch - if (batch_size != input.input_ids.value()->size(1)) { + // Vocab-parallel graphs are valid only for local-logit sampling. + if ((model_->uses_vocab_parallel_logits() + && !input.use_local_vocab_logits) + || batch_size != input.input_ids.value()->size(1)) { return {nullptr, nullptr}; } else { auto result = compiled_map_decode_.find(batch_size); diff --git a/csrc/engine/compiler/static_batching_compiler.cpp b/csrc/engine/compiler/static_batching_compiler.cpp index af2f4799f..76f35dcf3 100644 --- a/csrc/engine/compiler/static_batching_compiler.cpp +++ b/csrc/engine/compiler/static_batching_compiler.cpp @@ -11,6 +11,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.use_local_vocab_logits = true; input.input_ids = infinicore::Tensor::empty({b, 1}, infinicore::DataType::I64, 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()); @@ -46,6 +47,10 @@ void StaticBatchingCompiler::compile() { StaticBatchingCompiler::Compiled StaticBatchingCompiler::get_compiled( const InfinilmModel::Input &input) { if (model_->get_cache_config() != nullptr && dynamic_cast(model_->get_cache_config())) { + if (model_->uses_vocab_parallel_logits() + && !input.use_local_vocab_logits) { + return std::make_tuple(nullptr, nullptr); + } size_t batch_size = input.input_ids.value()->size(0); size_t seqlen = input.input_ids.value()->size(1); auto result = compiled_map_.find(std::make_tuple(batch_size, seqlen)); diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp index a5221eb37..a7eded92d 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -147,6 +147,11 @@ std::vector InferEngine::state_dict_keys() { infinilm::InfinilmModel::Input InferEngine::Input::to_model_input(infinicore::Device device) const { + if (input_offsets.has_value() + && input_offsets.value()->device().getType() != infinicore::Device::Type::CPU) { + throw std::invalid_argument("`input_offsets` must reside on CPU."); + } + auto to_device = [&](const std::optional &t) -> std::optional { return t.has_value() ? t.value()->to(device) : t; @@ -163,6 +168,34 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { } return result; }; + std::optional> last_token_indices; + if (!sample_all_positions + && input_offsets.has_value() + && input_ids.has_value() + && input_ids.value()->shape().size() == 2 + && input_ids.value()->size(0) == 1 + && input_offsets.value()->size(0) > 1 + && input_ids.value()->size(1) > input_offsets.value()->size(0) - 1) { + const auto &offsets_tensor = input_offsets.value(); + const size_t num_requests = offsets_tensor->size(0) - 1; + const auto *offsets = reinterpret_cast(offsets_tensor->data()); + if (offsets[0] != 0 + || static_cast(offsets[num_requests]) != input_ids.value()->size(1)) { + throw std::invalid_argument("input_offsets must cover the packed input"); + } + std::vector indices; + indices.reserve(num_requests); + for (size_t i = 0; i < num_requests; ++i) { + if (offsets[i + 1] <= offsets[i]) { + throw std::invalid_argument("input_offsets must describe non-empty requests"); + } + indices.push_back(static_cast(offsets[i + 1] - 1)); + } + last_token_indices = std::move(indices); + } + const bool use_local_vocab_logits = allow_local_vocab_logits + && !sample_all_positions + && is_greedy_sampling(); infinilm::InfinilmModel::Input input = { to_device(input_ids), // @todo: on device in the future @@ -181,7 +214,10 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { to_device_vec(image_grid_thw), image_req_ids, visual_token_ranges, - to_device(target_hidden_states)}; + to_device(target_hidden_states), + std::move(last_token_indices), + use_local_vocab_logits, + sample_all_positions}; infinilm::global_state::get_forward_context().attn_metadata = { input.past_sequence_lengths, diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp index 11696e1eb..4d2c9a474 100644 --- a/csrc/engine/rank_worker.cpp +++ b/csrc/engine/rank_worker.cpp @@ -1,6 +1,8 @@ #include "rank_worker.hpp" #include "../models/model_factory.hpp" #include "infinicore/ops.hpp" +#include "infinicore/ops/distributed/allgather.hpp" +#include "infinicore/ops/take.hpp" #include #include @@ -419,7 +421,9 @@ void RankWorker::thread_loop() { 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) { + if (!local_args.sample_all_positions + && local_args.allow_local_vocab_logits + && compiler_ != nullptr) { auto [graph, output] = compiler_->get_compiled(local_args.to_model_input(infinicore::Device::cpu())); if (graph != nullptr && output != nullptr) { graph->run(); @@ -434,43 +438,135 @@ void RankWorker::thread_loop() { hidden_states = model_output.hidden_states; } - // Random sampling (rank 0 only) - if (rank_info_.tp_rank == 0) { - auto temperature{local_args.temperature}; - auto top_p{local_args.top_p}; - auto top_k{local_args.top_k}; - - const auto &logits_shape{logits->shape()}; - const auto &vocab_size{logits_shape[2]}; - const auto &total_len{logits_shape[1]}; - const auto &batch_size{logits_shape[0]}; - - auto n_req = local_args.input_offsets.value()->size(0) - 1; - int32_t *input_offsets = (int32_t *)local_args.input_offsets.value()->data(); + auto temperature{local_args.temperature}; + auto top_p{local_args.top_p}; + auto top_k{local_args.top_k}; + const auto &logits_shape{logits->shape()}; + const size_t vocab_size{logits_shape[2]}; + const size_t total_len{logits_shape[1]}; + const size_t batch_size{logits_shape[0]}; + const size_t logits_positions = batch_size * total_len; + const size_t n_req = local_args.input_offsets.value()->size(0) - 1; + int32_t *input_offsets = reinterpret_cast( + local_args.input_offsets.value()->data()); + const bool sample_all_positions = local_args.sample_all_positions; + const bool logits_are_last_token_only = !sample_all_positions && logits_positions == n_req; + const size_t n_out = sample_all_positions + ? static_cast(input_offsets[n_req]) + : n_req; + + const size_t global_vocab_size = model_config_->get("vocab_size"); + const bool distributed_greedy = !sample_all_positions + && local_args.is_greedy_sampling() + && rank_info_.tp_size > 1 + && vocab_size * static_cast(rank_info_.tp_size) + == global_vocab_size; + + if (distributed_greedy) { + // Exchange one local argmax candidate per TP rank. + auto local_ids = infinicore::Tensor::empty( + {n_out}, infinicore::DataType::I64, rank_info_.device); + auto local_values = infinicore::Tensor::empty( + {n_out}, logits->dtype(), rank_info_.device); + infinicore::Tensor output_ids; + infinicore::Tensor winning_ranks; + if (rank_info_.tp_rank == 0) { + output_ids = infinicore::Tensor::empty( + {n_out}, infinicore::DataType::I64, rank_info_.device); + winning_ranks = infinicore::Tensor::empty( + {n_out}, infinicore::DataType::I64, rank_info_.device); + } + auto flat_logits = logits->view({logits_positions, vocab_size}); + for (size_t i = 0; i < n_out; ++i) { + size_t score_idx = i; + if (!logits_are_last_token_only) { + score_idx = static_cast( + input_offsets[i + 1] - 1); + } + auto score = flat_logits->narrow( + {{0, score_idx, 1}}) + ->view({vocab_size}); + auto local_id = local_ids->narrow({{0, i, 1}}); + infinicore::op::random_sample_( + local_id->view({}), score, 0.0f, + top_p, top_k, temperature); + infinicore::op::take_( + local_values->narrow({{0, i, 1}}), score, local_id); + } - 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 tp_size = static_cast(rank_info_.tp_size); + auto candidate_values = infinicore::op::distributed::allgather( + local_values, tp_size, rank_info_.comm); + auto candidate_ids = infinicore::op::distributed::allgather( + local_ids, tp_size, rank_info_.comm); + if (n_out > 1) { + candidate_values = candidate_values + ->view({tp_size, n_out}) + ->permute({1, 0}) + ->contiguous(); + candidate_ids = candidate_ids + ->view({tp_size, n_out}) + ->permute({1, 0}) + ->contiguous(); + } else { + candidate_values = candidate_values->view({1, tp_size}); + candidate_ids = candidate_ids->view({1, tp_size}); + } + if (rank_info_.tp_rank == 0) { + for (size_t i = 0; i < n_out; ++i) { + auto winning_rank = winning_ranks->narrow( + {{0, i, 1}}); + auto values = candidate_values->narrow( + {{0, i, 1}}) + ->view({tp_size}); + infinicore::op::random_sample_( + winning_rank->view({}), values, + 0.0f, 0.0f, 1, 0.0f); + auto ids = candidate_ids->narrow( + {{0, i, 1}}) + ->view({tp_size}); + infinicore::op::take_( + output_ids->narrow({{0, i, 1}}), ids, + winning_rank); + } + auto output_ids_cpu = output_ids->to( + infinicore::Device::cpu()); + auto winning_ranks_cpu = winning_ranks->to( + infinicore::Device::cpu()); + infinicore::context::syncStream(); + auto *output_ptr = reinterpret_cast( + output_ids_cpu->data()); + const auto *rank_ptr = reinterpret_cast( + winning_ranks_cpu->data()); + for (size_t i = 0; i < n_out; ++i) { + output_ptr[i] += rank_ptr[i] + * static_cast(vocab_size); + } + output_ = Output{ + output_ids_cpu, logits, hidden_states}; + } + } else if (rank_info_.tp_rank == 0) { + auto output_ids = infinicore::Tensor::empty( + {n_out}, infinicore::DataType::I64, rank_info_.device); + auto flat_logits = logits->view({logits_positions, vocab_size}); for (size_t i{0}; i < n_out; ++i) { size_t score_idx = i; - if (!sample_all_positions) { + if (!sample_all_positions && !logits_are_last_token_only) { score_idx = static_cast(input_offsets[i + 1] - 1); } - auto score{logits->view({batch_size * total_len, vocab_size})->narrow({{0, score_idx, 1}})->view({vocab_size})}; - auto out{output_ids->narrow({{0, i, 1}})->view({})}; + auto score = flat_logits->narrow( + {{0, score_idx, 1}}) + ->view({vocab_size}); + auto out = output_ids->narrow({{0, i, 1}})->view({}); float random_val = std::uniform_real_distribution(0, 1)(rng_); infinicore::op::random_sample_( 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}}; - - output_ = std::move(out); + output_ = Output{output_ids, logits, hidden_states}; } job_done_ = true; diff --git a/csrc/engine/rank_worker.hpp b/csrc/engine/rank_worker.hpp index d396ef6f1..085f6aca8 100644 --- a/csrc/engine/rank_worker.hpp +++ b/csrc/engine/rank_worker.hpp @@ -73,12 +73,19 @@ class RankWorker { /// Sample logits at every packed input position instead of one token per request. bool sample_all_positions{false}; + /// Allow callers to consume vocabulary-parallel logits. + bool allow_local_vocab_logits{false}; + float temperature{1}; int top_k{50}; float top_p{1}; + bool is_greedy_sampling() const { + return top_k == 1 || top_p == 0.0f || temperature == 0.0f; + } + infinilm::InfinilmModel::Input to_model_input(infinicore::Device device) const; }; diff --git a/csrc/layers/causal_lm_templates/text_causal_lm.hpp b/csrc/layers/causal_lm_templates/text_causal_lm.hpp index eb4f2b47f..e1469a858 100644 --- a/csrc/layers/causal_lm_templates/text_causal_lm.hpp +++ b/csrc/layers/causal_lm_templates/text_causal_lm.hpp @@ -1,8 +1,13 @@ #pragma once +#include "../../global_state/global_state.hpp" #include "../../models/infinilm_model.hpp" #include "../linear/linear.hpp" #include "infinicore/device.hpp" +#include "infinicore/ops/cat.hpp" +#include "infinicore/ops/distributed/allgather.hpp" +#include +#include namespace infinilm::layers::causal_lm_templates { @@ -36,7 +41,21 @@ class TextCausalLM : public InfinilmModel { const auto &dtype{model_config->get_dtype()}; model_ = this->register_module("model", model_config, device); - lm_head_ = this->register_module("lm_head", hidden_size, vocab_size, false, dtype, device); + const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); + tp_size_ = static_cast(rank_info.tp_size); + tp_communicator_ = rank_info.comm; + vocab_parallel_ = tp_size_ > 1 + && model_config->get("model_type") == "qwen3_moe" + && vocab_size % tp_size_ == 0 + && model_config->get_quant_scheme() == infinilm::quantization::QuantScheme::NONE; + if (vocab_parallel_) { + lm_head_ = this->register_module( + "lm_head", hidden_size, vocab_size, false, dtype, device, + rank_info.tp_rank, rank_info.tp_size); + } else { + lm_head_ = this->register_module( + "lm_head", hidden_size, vocab_size, false, dtype, device); + } } /** @@ -44,19 +63,76 @@ class TextCausalLM : public InfinilmModel { */ Output forward(const Input &input) const override { auto hidden_states = model_->forward(input); - auto logits = lm_head_->forward(hidden_states); + auto lm_head_input = hidden_states; + // Ordinary prefill needs only the final-position logits. + if (!input.sample_all_positions + && input.last_token_indices.has_value() + && !input.last_token_indices->empty() + && hidden_states->size(0) == 1 + && hidden_states->size(1) > input.last_token_indices->size()) { + const auto &indices = input.last_token_indices.value(); + for (const size_t index : indices) { + if (index >= hidden_states->size(1)) { + throw std::out_of_range("last-token index exceeds packed hidden states"); + } + } + if (indices.size() == 1) { + lm_head_input = hidden_states->narrow({{1, indices.front(), 1}}); + } else { + std::vector final_hidden_states; + final_hidden_states.reserve(indices.size()); + for (const size_t index : indices) { + final_hidden_states.push_back( + hidden_states->narrow({{1, index, 1}})); + } + lm_head_input = infinicore::op::cat(std::move(final_hidden_states), 1); + } + } + auto logits = lm_head_->forward(lm_head_input); + if (!input.use_local_vocab_logits) { + logits = gather_vocab_logits_(std::move(logits)); + } return {logits}; } + /// Return full-vocabulary logits, gathering rank-local shards when needed. infinicore::Tensor logits_from_hidden(const infinicore::Tensor &hidden_states) const { - return lm_head_->forward(const_cast(hidden_states)); + return gather_vocab_logits_( + lm_head_->forward(const_cast(hidden_states))); } Model &model() { return *model_; } + bool uses_vocab_parallel_logits() const override { + return vocab_parallel_; + } + protected: INFINICORE_NN_MODULE(Model, model); - INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, lm_head); + infinicore::Tensor gather_vocab_logits_(infinicore::Tensor logits) const { + if (!vocab_parallel_) { + return logits; + } + auto output_shape = logits->shape(); + const size_t local_vocab_size = output_shape.back(); + const size_t positions = logits->numel() / local_vocab_size; + auto gathered = infinicore::op::distributed::allgather( + logits, tp_size_, tp_communicator_); + output_shape.back() = local_vocab_size * tp_size_; + if (positions == 1) { + return gathered->view(output_shape); + } + return gathered->view({tp_size_, positions, local_vocab_size}) + ->permute({1, 0, 2}) + ->contiguous() + ->view(output_shape); + } + + // The base type holds either a replicated or column-parallel registered head. + INFINICORE_NN_MODULE(infinilm::layers::linear::BaseLinear, lm_head); + bool vocab_parallel_{false}; + size_t tp_size_{1}; + infinicclComm_t tp_communicator_{nullptr}; }; } // namespace infinilm::layers::causal_lm_templates diff --git a/csrc/models/infinilm_model.cpp b/csrc/models/infinilm_model.cpp index c2cc68762..2a77be12a 100644 --- a/csrc/models/infinilm_model.cpp +++ b/csrc/models/infinilm_model.cpp @@ -94,6 +94,12 @@ void InfinilmModel::process_weights_after_loading() { } void InfinilmModel::reset_runtime_state() const { + // Non-quantized Qwen3-MoE currently has no resettable module state. + if (model_config_ + && model_config_->get_or("model_type", "") == "qwen3_moe" + && model_config_->get_quant_scheme() == quantization::QuantScheme::NONE) { + return; + } for (const auto &[_, sub] : children()) { reset_runtime_state_recursive_(sub.get()); } diff --git a/csrc/models/infinilm_model.hpp b/csrc/models/infinilm_model.hpp index ac994fd6d..609f83f04 100644 --- a/csrc/models/infinilm_model.hpp +++ b/csrc/models/infinilm_model.hpp @@ -55,6 +55,12 @@ class InfinilmModel : public infinicore::nn::Module { std::optional> visual_token_ranges; /// Target model hidden states consumed by draft/MTP models. std::optional target_hidden_states; + /// Packed hidden-state positions sampled by ordinary prefill. + std::optional> last_token_indices; + /// Keep vocabulary-parallel logits sharded for distributed greedy sampling. + bool use_local_vocab_logits{false}; + /// Preserve logits for every packed position for speculative/MTP callers. + bool sample_all_positions{false}; }; struct Output { @@ -66,6 +72,9 @@ class InfinilmModel : public infinicore::nn::Module { virtual ~InfinilmModel() = default; virtual Output forward(const Input &input) const = 0; + virtual bool uses_vocab_parallel_logits() const { + return false; + } virtual void reset_cache(const cache::CacheConfig *cache_config); virtual const cache::CacheConfig *get_cache_config() const { return cache_config_.get(); diff --git a/csrc/pybind11/cache/cache.hpp b/csrc/pybind11/cache/cache.hpp index 492f6c302..439c9da2c 100644 --- a/csrc/pybind11/cache/cache.hpp +++ b/csrc/pybind11/cache/cache.hpp @@ -34,15 +34,19 @@ inline void bind_cache(py::module &m) { infinilm::cache::CacheConfig, std::shared_ptr>(m, "PagedKVCacheConfig") .def( - py::init(), + py::init(), py::arg("num_blocks"), - py::arg("block_size") = 256) + py::arg("block_size") = 256, + py::arg("max_batch_size") = 512) .def( "num_blocks", &infinilm::cache::PagedKVCacheConfig::num_blocks) .def( "block_size", &infinilm::cache::PagedKVCacheConfig::block_size) + .def( + "max_batch_size", + &infinilm::cache::PagedKVCacheConfig::max_batch_size) .def("__repr__", [](const infinilm::cache::PagedKVCacheConfig &) { return ""; }); diff --git a/csrc/pybind11/engine/engine.hpp b/csrc/pybind11/engine/engine.hpp index b73ce06ae..b3f5dac57 100644 --- a/csrc/pybind11/engine/engine.hpp +++ b/csrc/pybind11/engine/engine.hpp @@ -239,6 +239,7 @@ inline void bind_infer_engine(py::module &m) { .def_readwrite("visual_token_ranges", &InferEngine::Input::visual_token_ranges) .def_readwrite("target_hidden_states", &InferEngine::Input::target_hidden_states) .def_readwrite("sample_all_positions", &InferEngine::Input::sample_all_positions) + .def_readwrite("allow_local_vocab_logits", &InferEngine::Input::allow_local_vocab_logits) .def_readwrite("temperature", &InferEngine::Input::temperature) .def_readwrite("top_k", &InferEngine::Input::top_k) .def_readwrite("top_p", &InferEngine::Input::top_p); diff --git a/examples/test_infer.py b/examples/test_infer.py index b02500422..334c1251b 100644 --- a/examples/test_infer.py +++ b/examples/test_infer.py @@ -31,7 +31,8 @@ def test( video_num_frames=None, skip_load=False, weight_load_mode="async", - skip_legacy_moe=False, + skip_legacy_moe=None, + moe_backend="auto", ): model_path = os.path.expanduser(model_path) # ---------------------------------------------------------------------------- # @@ -62,6 +63,7 @@ def test( skip_load=skip_load, weight_load_mode=weight_load_mode, skip_legacy_moe=skip_legacy_moe, + moe_backend=moe_backend, ) conversations = [ @@ -150,4 +152,5 @@ def test( skip_load=cfg.skip_load, weight_load_mode=cfg.weight_load_mode, skip_legacy_moe=cfg.skip_legacy_moe, + moe_backend=cfg.moe_backend, ) diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index 3b8f6533d..05ac361ee 100644 --- a/python/infinilm/base_config.py +++ b/python/infinilm/base_config.py @@ -66,6 +66,7 @@ def __init__(self): self.dp = self.args.dp self.ep = self.args.ep self.moe_ep_backend = self.args.moe_ep_backend + self.moe_backend = self.args.moe_backend self.skip_legacy_moe = self.args.skip_legacy_moe self.attn = self.args.attn @@ -209,9 +210,16 @@ def _add_common_args(self): default="auto", help=MOE_EP_BACKEND_HELP, ) + self.parser.add_argument( + "--moe-backend", + choices=["auto", "legacy", "fused"], + default="auto", + help="MoE execution backend", + ) self.parser.add_argument( "--skip-legacy-moe", action="store_true", + default=None, help="use the new fused MoE implementation instead of the legacy Qwen3 MoE MLP", ) diff --git a/python/infinilm/cache/cache.py b/python/infinilm/cache/cache.py index 4a9bcc446..b8b8e41c4 100644 --- a/python/infinilm/cache/cache.py +++ b/python/infinilm/cache/cache.py @@ -18,5 +18,9 @@ def __init__( class PagedKVCacheConfig(CacheConfig, _infinilm.PagedKVCacheConfig): - def __init__(self, num_blocks: int, block_size: int = 256): - _infinilm.PagedKVCacheConfig.__init__(self, num_blocks, block_size) + def __init__( + self, num_blocks: int, block_size: int = 256, max_batch_size: int = 512 + ): + _infinilm.PagedKVCacheConfig.__init__( + self, num_blocks, block_size, max_batch_size + ) diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 3e874046e..4d1937c2a 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -230,6 +230,7 @@ def _build_input( visual_token_ranges=None, target_hidden_states=None, sample_all_positions=False, + allow_local_vocab_logits=False, temperature=None, top_k=None, top_p=None, @@ -269,7 +270,7 @@ def convert_tensor_list(tensor_list_): top_k = 1 if top_k is None else top_k top_p = 1.0 if top_p is None else top_p - return super().Input( + model_input = super().Input( input_ids, position_ids=position_ids, past_sequence_lengths=past_kv_lengths, @@ -292,6 +293,8 @@ def convert_tensor_list(tensor_list_): top_k=top_k, top_p=top_p, ) + model_input.allow_local_vocab_logits = allow_local_vocab_logits + return model_input def forward( self, @@ -385,6 +388,7 @@ def convert_tensor_list(tensor_list_): image_req_ids=image_req_ids, visual_token_ranges=visual_token_ranges, target_hidden_states=target_hidden_states, + allow_local_vocab_logits=True, temperature=temperature, top_k=top_k, top_p=top_p, @@ -419,28 +423,27 @@ def forward_raw( top_p=None, ): try: - output = super().forward( - self._build_input( - input_ids, - position_ids=position_ids, - past_kv_lengths=past_kv_lengths, - total_kv_lengths=total_kv_lengths, - input_offsets=input_offsets, - cu_seqlens=cu_seqlens, - block_tables=block_tables, - slot_mapping=slot_mapping, - pixel_values=pixel_values, - image_bound=image_bound, - tgt_sizes=tgt_sizes, - image_req_ids=image_req_ids, - visual_token_ranges=visual_token_ranges, - target_hidden_states=target_hidden_states, - sample_all_positions=sample_all_positions, - temperature=temperature, - top_k=top_k, - top_p=top_p, - ) + model_input = self._build_input( + input_ids, + position_ids=position_ids, + past_kv_lengths=past_kv_lengths, + total_kv_lengths=total_kv_lengths, + input_offsets=input_offsets, + cu_seqlens=cu_seqlens, + block_tables=block_tables, + slot_mapping=slot_mapping, + pixel_values=pixel_values, + image_bound=image_bound, + tgt_sizes=tgt_sizes, + image_req_ids=image_req_ids, + visual_token_ranges=visual_token_ranges, + target_hidden_states=target_hidden_states, + sample_all_positions=sample_all_positions, + temperature=temperature, + top_k=top_k, + top_p=top_p, ) + output = super().forward(model_input) return { "output_ids": infinicore.Tensor(output.output_ids), "logits": infinicore.Tensor(output.logits), diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py index 59e9e94ba..ddb19a985 100644 --- a/python/infinilm/llm/llm.py +++ b/python/infinilm/llm/llm.py @@ -7,6 +7,7 @@ """ import asyncio +import json import logging import os import threading @@ -33,6 +34,84 @@ logger = logging.getLogger(__name__) +_MOE_BACKENDS = {"auto", "legacy", "fused"} +# Apply the conservative prefill budget from 32 concurrent requests onward. +# It still admits 64 x 32-token prompts in one wave while bounding longer packed +# prefills on 24 GiB GPUs. `INFINILM_MAX_NUM_BATCHED_TOKENS` overrides it. +_QWEN3_MOE_HIGH_CONCURRENCY_BATCH_SIZE = 32 +_QWEN3_MOE_MAX_BATCHED_TOKENS = 4096 + + +def _resolve_moe_backend( + model_path: str, + device: str, + dtype: str, + moe_backend: str, + skip_legacy_moe: Optional[bool], +): + """Resolve the MoE policy while preserving the legacy boolean override.""" + if moe_backend not in _MOE_BACKENDS: + raise ValueError( + f"Unsupported moe_backend {moe_backend!r}; expected one of " + f"{sorted(_MOE_BACKENDS)}" + ) + + if skip_legacy_moe is not None: + resolved_backend = "fused" if skip_legacy_moe else "legacy" + if moe_backend != "auto" and moe_backend != resolved_backend: + raise ValueError( + "Conflicting MoE options: " + f"moe_backend={moe_backend!r}, skip_legacy_moe={skip_legacy_moe!r}" + ) + logger.info( + "Resolved MoE backend to %s from explicit skip_legacy_moe=%s", + resolved_backend, + skip_legacy_moe, + ) + return skip_legacy_moe, resolved_backend + + if moe_backend != "auto": + return moe_backend == "fused", moe_backend + + config_path = os.path.join(model_path, "config.json") + with open(config_path, encoding="utf-8") as config_file: + hf_config = json.load(config_file) + + text_config = hf_config.get("text_config") or hf_config + quantization_config = text_config.get("quantization_config") or {} + quant_method = quantization_config.get("quant_method") or text_config.get( + "quant_method" + ) + model_type = text_config.get("model_type", "") + fused_supported = ( + model_type == "qwen3_moe" + and device == "cuda" + and dtype in {"float16", "bfloat16"} + and not quant_method + ) + resolved_backend = "fused" if fused_supported else "legacy" + logger.info( + "Resolved moe_backend=auto to %s (model_type=%s, device=%s, " + "dtype=%s, quant_method=%s)", + resolved_backend, + model_type or "unknown", + device, + dtype, + quant_method or "none", + ) + return fused_supported, resolved_backend + + +def _default_max_num_batched_tokens( + llm_config: dict, max_batch_size: int, max_position_embeddings: int +) -> int: + if ( + llm_config.get("model_type") == "qwen3_moe" + and max_batch_size >= _QWEN3_MOE_HIGH_CONCURRENCY_BATCH_SIZE + ): + return min(max_position_embeddings, _QWEN3_MOE_MAX_BATCHED_TOKENS) + return max_position_embeddings + class LLMEngine: """Low-level LLM engine that handles inference execution.""" @@ -82,10 +161,18 @@ def __init__(self, config: EngineConfig): ) num_mamba_cache_blocks = max(2, config.num_blocks // 4) + # Bound packed prefill activations for high-concurrency Qwen3-MoE. + default_max_num_batched_tokens = _default_max_num_batched_tokens( + llm_config, config.max_batch_size, max_position_embeddings + ) max_num_batched_tokens = int( - os.getenv("INFINILM_MAX_NUM_BATCHED_TOKENS", max_position_embeddings) + os.getenv( + "INFINILM_MAX_NUM_BATCHED_TOKENS", + default_max_num_batched_tokens, + ) ) assert 1024 <= max_num_batched_tokens <= max_position_embeddings + logger.info("Scheduler max_num_batched_tokens=%d", max_num_batched_tokens) self.scheduler = Scheduler( max_batch_size=config.max_batch_size, @@ -332,7 +419,8 @@ def __init__( use_mla: bool = False, weight_load_mode: str = "async", skip_load: bool = False, - skip_legacy_moe: bool = False, + skip_legacy_moe: Optional[bool] = None, + moe_backend: str = "auto", ): """Initialize LLM. @@ -354,7 +442,16 @@ def __init__( attn_backend: Attention backend to use ('default', 'flash-attn'). use_mla: Whether to use DeepSeek V2 MLA attention when supported. weight_load_mode: Weight loading mode across tensor-parallel workers. + skip_legacy_moe: Compatibility override for the legacy boolean API. + moe_backend: MoE execution policy ('auto', 'legacy', or 'fused'). """ + skip_legacy_moe, resolved_moe_backend = _resolve_moe_backend( + model_path=model_path, + device=device, + dtype=dtype, + moe_backend=moe_backend, + skip_legacy_moe=skip_legacy_moe, + ) config = EngineConfig( model_path=model_path, draft_model_path=draft_model_path, @@ -382,6 +479,7 @@ def __init__( ) self.engine = LLMEngine(config) self.config = config + self.moe_backend = resolved_moe_backend def generate( self, @@ -539,7 +637,8 @@ def __init__( kv_transfer_config: Optional[KVTransferConfig] = None, use_mla: bool = False, weight_load_mode: str = "async", - skip_legacy_moe: bool = False, + skip_legacy_moe: Optional[bool] = None, + moe_backend: str = "auto", ): """Initialize AsyncLLMEngine. @@ -564,7 +663,16 @@ def __init__( kv_connector_extra_config: Extra config dict for KV connector. use_mla: Whether to use DeepSeek V2 MLA attention when supported. weight_load_mode: Weight loading mode across tensor-parallel workers. + skip_legacy_moe: Compatibility override for the legacy boolean API. + moe_backend: MoE execution policy ('auto', 'legacy', or 'fused'). """ + skip_legacy_moe, resolved_moe_backend = _resolve_moe_backend( + model_path=model_path, + device=device, + dtype=dtype, + moe_backend=moe_backend, + skip_legacy_moe=skip_legacy_moe, + ) config = EngineConfig( model_path=model_path, draft_model_path=draft_model_path, @@ -592,6 +700,7 @@ def __init__( ) self.engine = LLMEngine(config) self.config = config + self.moe_backend = resolved_moe_backend self._running = False self._step_thread: Optional[threading.Thread] = None diff --git a/python/infinilm/llm/model_runner/model_runner.py b/python/infinilm/llm/model_runner/model_runner.py index 45f1eaee1..c611b1f30 100644 --- a/python/infinilm/llm/model_runner/model_runner.py +++ b/python/infinilm/llm/model_runner/model_runner.py @@ -59,7 +59,9 @@ def __init__(self, config: EngineConfig): ) elif config.cache_type == "paged": cache_config = PagedKVCacheConfig( - num_blocks=config.num_blocks, block_size=config.block_size + num_blocks=config.num_blocks, + block_size=config.block_size, + max_batch_size=config.max_batch_size, ) logger.info(f"Using Paged KV Cache with num_blocks={config.num_blocks}") else: diff --git a/python/infinilm/server/inference_server.py b/python/infinilm/server/inference_server.py index 242593e25..69d0829b2 100644 --- a/python/infinilm/server/inference_server.py +++ b/python/infinilm/server/inference_server.py @@ -117,6 +117,7 @@ def __init__( weight_load_mode: str = "async", ignore_eos: bool = False, kv_transfer_config: Optional[KVTransferConfig] = None, + moe_backend: str = "auto", ): """Initialize inference server. @@ -144,6 +145,7 @@ def __init__( weight_load_mode: Weight loading mode across tensor-parallel workers. ignore_eos: Whether to ignore EOS tokens during generation. kv_transfer_config: Optional configuration for the KV transfer mechanism. + moe_backend: MoE execution backend. """ self.model_path = model_path # vLLM-like served model id: directory name of model_path @@ -153,6 +155,7 @@ def __init__( self.tensor_parallel_size = tensor_parallel_size self.moe_ep_backend = moe_ep_backend self.moe_ep_size = moe_ep_size + self.moe_backend = moe_backend self.cache_type = cache_type self.max_tokens = max_tokens self.max_batch_size = max_batch_size @@ -192,6 +195,7 @@ 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, + moe_backend=self.moe_backend, cache_type=self.cache_type, max_batch_size=self.max_batch_size, max_tokens=self.max_tokens, @@ -607,6 +611,7 @@ def main(): tensor_parallel_size=cfg.tp, moe_ep_backend=moe_ep_backend, moe_ep_size=ep, + moe_backend=cfg.moe_backend, cache_type="paged" if cfg.enable_paged_attn else "static", max_tokens=cfg.max_new_tokens, max_batch_size=cfg.max_batch_size, diff --git a/scripts/run_qwen3_moe_contest.sh b/scripts/run_qwen3_moe_contest.sh new file mode 100755 index 000000000..ad7b30849 --- /dev/null +++ b/scripts/run_qwen3_moe_contest.sh @@ -0,0 +1,354 @@ +#!/usr/bin/env bash +# Run the T2-1-3 Qwen3-MoE correctness or serving benchmark profiles. +set -euo pipefail + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +REPO_DIR=$(cd "$SCRIPT_DIR/.." && pwd) + +usage() { + cat <<'EOF' +Usage: scripts/run_qwen3_moe_contest.sh PROFILE + +Profiles: + check Run the focused policy tests without starting a server. + smoke Start a small server and issue one real chat request. + c1-short Benchmark concurrency=1, input=32, output=256. + c64-short Benchmark concurrency=64, input=32, output=256. + c64-long Benchmark concurrency=64, input=256, output=1024. + +Environment: + MODEL_DIR Qwen3-30B-A3B-Thinking-2507 directory (required except for check) + INFINICORE_DIR Optional InfiniCore checkout; omit when installed in Python + INFINI_ROOT Installed InfiniCore prefix (default: $HOME/.infini) + PYTHON Server Python executable (default: python3) + BENCH_PYTHON Python with vLLM 0.9.2 client modules (default: PYTHON) + RESULT_DIR Log/result directory (default: current-directory benchmark_results/) + CUDA_VISIBLE_DEVICES Four TP devices (default: 0,1,2,3) + PORT Local service port (default: 8102) + BUILD=1 Rebuild/install InfiniLM before running the profile + DRY_RUN=1 Print the resolved commands without executing them + SHOW_COMMANDS=1 Print full commands during a real run + +Examples: + MODEL_DIR=/path/to/Qwen3-30B-A3B-Thinking-2507 \ + scripts/run_qwen3_moe_contest.sh smoke + + MODEL_DIR=/path/to/Qwen3-30B-A3B-Thinking-2507 \ + BENCH_PYTHON=/path/to/vllm-client/bin/python \ + scripts/run_qwen3_moe_contest.sh c64-long +EOF +} + +PROFILE=${1:-} +if [ -z "$PROFILE" ] || [ "$PROFILE" = "-h" ] || [ "$PROFILE" = "--help" ]; then + usage + exit 0 +fi + +MODEL_DIR=${MODEL_DIR:-} +INFINICORE_DIR=${INFINICORE_DIR:-} +INFINI_ROOT=${INFINI_ROOT:-${HOME:-$REPO_DIR}/.infini} +PYTHON=${PYTHON:-python3} +BENCH_PYTHON=${BENCH_PYTHON:-$PYTHON} +RESULT_DIR=${RESULT_DIR:-benchmark_results/qwen3_moe} +CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3} +HOST=${HOST:-127.0.0.1} +PORT=${PORT:-8102} +TP=${TP:-4} +BUILD=${BUILD:-0} +DRY_RUN=${DRY_RUN:-0} +SHOW_COMMANDS=${SHOW_COMMANDS:-0} +SERVER_START_TIMEOUT=${SERVER_START_TIMEOUT:-1800} + +MAX_BATCH_SIZE=1 +NUM_BLOCKS=18 +NUM_PROMPTS=1 +MAX_CONCURRENCY=1 +INPUT_LEN=32 +OUTPUT_LEN=128 +RUN_BENCH=0 + +case "$PROFILE" in + check) + ;; + smoke) + ;; + c1-short) + OUTPUT_LEN=256 + RUN_BENCH=1 + ;; + c64-short) + MAX_BATCH_SIZE=64 + NUM_BLOCKS=128 + NUM_PROMPTS=64 + MAX_CONCURRENCY=64 + OUTPUT_LEN=256 + RUN_BENCH=1 + ;; + c64-long) + MAX_BATCH_SIZE=64 + NUM_BLOCKS=120 + NUM_PROMPTS=64 + MAX_CONCURRENCY=64 + INPUT_LEN=256 + OUTPUT_LEN=1024 + RUN_BENCH=1 + ;; + *) + echo "Unknown profile: $PROFILE" >&2 + usage >&2 + exit 2 + ;; +esac + +if [ "$PROFILE" != "check" ] && [ -z "$MODEL_DIR" ]; then + echo "MODEL_DIR is required for profile $PROFILE" >&2 + exit 2 +fi + +if [ -n "$MODEL_DIR" ]; then + MODEL_NAME=${MODEL_NAME:-$(basename "$MODEL_DIR")} +else + MODEL_NAME=${MODEL_NAME:-Qwen3-30B-A3B-Thinking-2507} +fi +TIMESTAMP=$(date -u +%Y%m%dT%H%M%SZ) +RUN_NAME="${PROFILE}-${TIMESTAMP}" +SERVER_LOG="$RESULT_DIR/${RUN_NAME}-server.log" +CLIENT_LOG="$RESULT_DIR/${RUN_NAME}-client.log" +RESULT_FILE="$RUN_NAME.json" + +export INFINI_ROOT +export CUDA_VISIBLE_DEVICES +export LD_LIBRARY_PATH="$INFINI_ROOT/lib:${LD_LIBRARY_PATH:-}" +if [ -n "$INFINICORE_DIR" ]; then + export PYTHONPATH="$INFINICORE_DIR/python:$REPO_DIR/python:${PYTHONPATH:-}" +else + export PYTHONPATH="$REPO_DIR/python:${PYTHONPATH:-}" +fi +if [ "$MAX_BATCH_SIZE" -ge 32 ]; then + export INFINILM_MAX_NUM_BATCHED_TOKENS=4096 +else + unset INFINILM_MAX_NUM_BATCHED_TOKENS || true +fi + +SERVER_COMMAND=( + "$PYTHON" "$REPO_DIR/python/infinilm/server/inference_server.py" + --device nvidia + --model "$MODEL_DIR" + --dtype bfloat16 + --tp "$TP" + --dp 1 + --ep 1 + --moe-ep-backend disabled + --host "$HOST" + --port "$PORT" + --max-batch-size "$MAX_BATCH_SIZE" + --max-new-tokens 4096 + --num-blocks "$NUM_BLOCKS" + --block-size 256 + --temperature 1.0 + --top-p 0.8 + --top-k 1 + --enable-paged-attn + --attn paged-attn + --enable-graph + --moe-backend auto + --ignore-eos +) + +BENCH_COMMAND=( + "$BENCH_PYTHON" "$SCRIPT_DIR/run_vllm_bench_serve_client.py" + --backend openai-chat + --endpoint-type openai-chat + --model "$MODEL_NAME" + --tokenizer "$MODEL_DIR" + --endpoint /v1/chat/completions + --host "$HOST" + --port "$PORT" + --dataset-name random + --request-rate inf + --seed 714 + --num-prompts "$NUM_PROMPTS" + --max-concurrency "$MAX_CONCURRENCY" + --random-input-len "$INPUT_LEN" + --random-output-len "$OUTPUT_LEN" + --temperature 1.0 + --top-p 0.8 + --top-k 1 + --extra-body "{\"max_tokens\": $OUTPUT_LEN}" + --save-result + --save-detailed + --result-dir "$RESULT_DIR" + --result-filename "$RESULT_FILE" +) + +print_command() { + printf '%q ' "$@" + printf '\n' +} + +echo "Profile: $PROFILE" +if [ -n "$MODEL_DIR" ]; then + echo "Model: $MODEL_NAME" +fi + +if [ "$PROFILE" = "check" ]; then + CHECK_COMMAND=( + "$PYTHON" -m pytest -q "$REPO_DIR/test/llm/test_moe_policy.py" + ) + if [ "$DRY_RUN" = "1" ]; then + print_command "${CHECK_COMMAND[@]}" + exit 0 + fi + "${CHECK_COMMAND[@]}" + exit 0 +fi + +if [ "$DRY_RUN" = "1" ]; then + echo "Server command:" + print_command "${SERVER_COMMAND[@]}" + if [ "$RUN_BENCH" = "1" ]; then + echo "Client command:" + print_command "${BENCH_COMMAND[@]}" + fi + exit 0 +fi + +if [ ! -f "$MODEL_DIR/config.json" ]; then + echo "Model config not found: $MODEL_DIR/config.json" >&2 + exit 3 +fi +if [ -n "$INFINICORE_DIR" ] && [ ! -d "$INFINICORE_DIR/python/infinicore" ]; then + echo "InfiniCore Python package not found: $INFINICORE_DIR/python/infinicore" >&2 + exit 3 +fi +if ! command -v curl >/dev/null 2>&1; then + echo "curl is required" >&2 + exit 3 +fi + +mkdir -p "$RESULT_DIR" + +if [ "$BUILD" = "1" ]; then + echo "Building and installing InfiniLM..." + ( + cd "$REPO_DIR" + "$PYTHON" -m pip install -e . --no-build-isolation + ) +fi + +"$PYTHON" -c \ + 'import infinicore; from infinilm.lib import _infinilm; print("InfiniLM extension import: OK")' + +GPU_COUNT=$( + "$PYTHON" -c 'import torch; print(torch.cuda.device_count())' +) +if [ "$GPU_COUNT" -lt "$TP" ]; then + echo "Expected at least $TP visible GPUs, found $GPU_COUNT" >&2 + exit 4 +fi + +if [ "$RUN_BENCH" = "1" ]; then + "$BENCH_PYTHON" -c 'import vllm; print(f"vLLM client {vllm.__version__}")' +fi + +if curl -fsS "http://$HOST:$PORT/health" >/dev/null 2>&1; then + echo "A service is already listening at $HOST:$PORT; choose another PORT" >&2 + exit 5 +fi + +SERVER_PID= +cleanup() { + if [ -n "$SERVER_PID" ] && kill -0 "$SERVER_PID" >/dev/null 2>&1; then + echo "Stopping server process $SERVER_PID..." + kill "$SERVER_PID" >/dev/null 2>&1 || true + wait "$SERVER_PID" 2>/dev/null || true + fi +} +trap cleanup EXIT INT TERM + +echo "Starting InfiniLM server..." +if [ "$SHOW_COMMANDS" = "1" ]; then + print_command "${SERVER_COMMAND[@]}" +fi +"${SERVER_COMMAND[@]}" >"$SERVER_LOG" 2>&1 & +SERVER_PID=$! + +READY=0 +ATTEMPT=0 +START_SECONDS=$SECONDS +START_DEADLINE=$((START_SECONDS + SERVER_START_TIMEOUT)) +while [ "$SECONDS" -lt "$START_DEADLINE" ]; do + if curl -fsS "http://$HOST:$PORT/health" >/dev/null 2>&1; then + READY=1 + break + fi + if ! kill -0 "$SERVER_PID" >/dev/null 2>&1; then + echo "Server exited before becoming ready" >&2 + tail -n 100 "$SERVER_LOG" >&2 || true + exit 6 + fi + ATTEMPT=$((ATTEMPT + 1)) + if [ $((ATTEMPT % 6)) -eq 0 ]; then + echo "Waiting for model load and graph capture ($((SECONDS - START_SECONDS))s)..." + fi + sleep 5 +done + +if [ "$READY" != "1" ]; then + echo "Server did not become ready within ${SERVER_START_TIMEOUT}s" >&2 + tail -n 100 "$SERVER_LOG" >&2 || true + exit 7 +fi + +echo "Server is ready." +unset http_proxy https_proxy all_proxy ALL_PROXY || true + +if [ "$RUN_BENCH" = "1" ]; then + echo "Running the contest-style vLLM client..." + if [ "$SHOW_COMMANDS" = "1" ]; then + print_command "${BENCH_COMMAND[@]}" + fi + "${BENCH_COMMAND[@]}" 2>&1 | tee "$CLIENT_LOG" +else + SMOKE_FILE="$RESULT_DIR/${RUN_NAME}-smoke.json" + PAYLOAD=$(printf \ + '{"model":"%s","messages":[{"role":"user","content":"What is 15 + 27? Answer briefly."}],"max_tokens":%d,"temperature":0.0}' \ + "$MODEL_NAME" "$OUTPUT_LEN") + curl -fsS "http://$HOST:$PORT/v1/chat/completions" \ + -H 'Content-Type: application/json' \ + --data-binary "$PAYLOAD" | tee "$SMOKE_FILE" + printf '\n' + "$PYTHON" - "$SMOKE_FILE" <<'PY' +import json +import sys + +with open(sys.argv[1], encoding="utf-8") as response_file: + response = json.load(response_file) + +try: + choice = response["choices"][0] + content = choice["message"]["content"] + completion_tokens = response["usage"]["completion_tokens"] +except (KeyError, IndexError, TypeError) as error: + raise SystemExit(f"Smoke validation failed: malformed response: {error}") from error + +if not isinstance(content, str) or not content.strip(): + raise SystemExit("Smoke validation failed: empty completion") +if not isinstance(completion_tokens, int) or completion_tokens <= 0: + raise SystemExit("Smoke validation failed: no completion tokens") +if "42" not in content: + raise SystemExit("Smoke validation failed: arithmetic answer 42 not found") + +print( + "Smoke validation: PASS " + f"(completion_tokens={completion_tokens}, answer_contains_42=yes)" +) +PY +fi + +echo "Server log: $SERVER_LOG" +if [ "$RUN_BENCH" = "1" ]; then + echo "Client log: $CLIENT_LOG" + echo "Result JSON: $RESULT_DIR/$RESULT_FILE" +fi diff --git a/scripts/run_vllm_bench_serve_client.py b/scripts/run_vllm_bench_serve_client.py new file mode 100755 index 000000000..0cb5d11c3 --- /dev/null +++ b/scripts/run_vllm_bench_serve_client.py @@ -0,0 +1,43 @@ +#!/usr/bin/env python3 +"""Run vLLM 0.9.2's official serving benchmark client. + +This launcher imports only the serving benchmark modules, so the client does +not need to initialize a vLLM inference engine. It also supplies the generic +``--extra-body`` option that was added after vLLM 0.9.2. +""" + +import argparse +import json + +from vllm.benchmarks import endpoint_request_func +from vllm.benchmarks.serve import add_cli_args, main + + +def json_object(value: str) -> dict: + """Parse a JSON object for ``--extra-body``.""" + parsed = json.loads(value) + if not isinstance(parsed, dict): + raise argparse.ArgumentTypeError("--extra-body must be a JSON object") + return parsed + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="vLLM serving benchmark client") + add_cli_args(parser) + parser.add_argument("--extra-body", type=json_object, default=None) + args = parser.parse_args() + + if args.extra_body: + original = endpoint_request_func.ASYNC_REQUEST_FUNCS["openai-chat"] + + async def openai_chat_with_extra_body(request_func_input, pbar=None): + merged = dict(request_func_input.extra_body or {}) + merged.update(args.extra_body) + request_func_input.extra_body = merged + return await original(request_func_input, pbar) + + endpoint_request_func.ASYNC_REQUEST_FUNCS["openai-chat"] = ( + openai_chat_with_extra_body + ) + + main(args) diff --git a/test/llm/test_moe_policy.py b/test/llm/test_moe_policy.py new file mode 100644 index 000000000..e1af4dd2e --- /dev/null +++ b/test/llm/test_moe_policy.py @@ -0,0 +1,70 @@ +import json + +import pytest +from infinilm.llm.llm import ( + _default_max_num_batched_tokens, + _resolve_moe_backend, +) + + +def _model_dir(tmp_path, config): + (tmp_path / "config.json").write_text(json.dumps(config), encoding="utf-8") + return str(tmp_path) + + +def test_auto_selects_fused_qwen3_moe(tmp_path): + model_path = _model_dir(tmp_path, {"model_type": "qwen3_moe"}) + + assert _resolve_moe_backend(model_path, "cuda", "bfloat16", "auto", None) == ( + True, + "fused", + ) + + +@pytest.mark.parametrize( + ("config", "device", "dtype"), + [ + ({"model_type": "qwen3_moe"}, "cpu", "bfloat16"), + ({"model_type": "qwen3_moe"}, "cuda", "float32"), + ({"model_type": "llama"}, "cuda", "bfloat16"), + ( + { + "model_type": "qwen3_moe", + "quantization_config": {"quant_method": "gptq"}, + }, + "cuda", + "bfloat16", + ), + ], +) +def test_auto_preserves_legacy_for_unsupported_configs(tmp_path, config, device, dtype): + model_path = _model_dir(tmp_path, config) + + assert _resolve_moe_backend(model_path, device, dtype, "auto", None) == ( + False, + "legacy", + ) + + +def test_explicit_backend_and_legacy_override(tmp_path): + model_path = _model_dir(tmp_path, {"model_type": "llama"}) + + assert _resolve_moe_backend(model_path, "cpu", "float32", "fused", None) == ( + True, + "fused", + ) + assert _resolve_moe_backend(model_path, "cpu", "float32", "auto", False) == ( + False, + "legacy", + ) + with pytest.raises(ValueError, match="Conflicting MoE options"): + _resolve_moe_backend(model_path, "cpu", "float32", "legacy", True) + + +def test_qwen3_moe_prefill_budget(): + qwen_config = {"model_type": "qwen3_moe"} + + assert _default_max_num_batched_tokens(qwen_config, 31, 32768) == 32768 + assert _default_max_num_batched_tokens(qwen_config, 32, 32768) == 4096 + assert _default_max_num_batched_tokens(qwen_config, 64, 2048) == 2048 + assert _default_max_num_batched_tokens({"model_type": "llama"}, 64, 32768) == 32768