[2026春季][T2-1-3] surprisely

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

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<CacheConfig>
@@ -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

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<CacheConfig> 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 {

csrc/engine/compiler/paged_compiler.cpp

@@ -5,7 +5,20 @@
 namespace infinilm::engine {
 
 PagedCompiler::PagedCompiler(const std::shared_ptr<InfinilmModel> &model, RankBarrier *barrier)
-    : GraphCompiler(model, barrier) {
+    : GraphCompiler(model, barrier) {}
+
+void PagedCompiler::compile() {
+    compiled_map_decode_.clear();
+    const auto *config = dynamic_cast<const cache::PagedKVCacheConfig *>(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<InfinilmModel> &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<const cache::PagedKVCacheConfig *>(model_->get_cache_config())) {
-        size_t nblocks = dynamic_cast<const cache::PagedKVCacheConfig *>(model_->get_cache_config())->num_blocks();
-        size_t max_batch_size = *std::max_element(decode_batch_sizes_.begin(), decode_batch_sizes_.end());
-        compiled_map_decode_.clear();
-        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<int32_t> 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<int32_t> 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<size_t>(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<int32_t> 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<int32_t> 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<size_t>(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<InfinilmModel::Output>(
-                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<InfinilmModel::Output>(
+            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);

csrc/engine/compiler/static_batching_compiler.cpp

@@ -11,6 +11,7 @@ void StaticBatchingCompiler::compile() {
     if (model_->get_cache_config() != nullptr && dynamic_cast<const cache::StaticKVCacheConfig *>(model_->get_cache_config())) {
         size_t b = dynamic_cast<const cache::StaticKVCacheConfig *>(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<const cache::StaticKVCacheConfig *>(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));

csrc/engine/infer_engine.cpp

@@ -147,6 +147,11 @@ std::vector<std::string> 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<infinicore::Tensor> &t)
         -> std::optional<infinicore::Tensor> {
         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<std::vector<size_t>> 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<const int32_t *>(offsets_tensor->data());
+        if (offsets[0] != 0
+            || static_cast<size_t>(offsets[num_requests]) != input_ids.value()->size(1)) {
+            throw std::invalid_argument("input_offsets must cover the packed input");
+        }
+        std::vector<size_t> 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<size_t>(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,