feat: add Qwen3.6 MoE model support

csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.cpp → csrc/cache/hybrid_cache.cpp

@@ -1,40 +1,40 @@
-#include "qwen3_next_allocate_kv_cache_tensors.hpp"
-
-#include "../../global_state/global_state.hpp"
-#include "../../utils.hpp"
-#include "infinicore/context/context.hpp"
+#include "hybrid_cache.hpp"
 #include <algorithm>
 #include <stdexcept>
 #include <utility>
 #include <vector>
 
-namespace infinilm::models::qwen3_next {
+namespace infinilm::cache {
 
-AllocatedHybridCache qwen3_next_allocate_cache_tensors(
-    const cache::CacheConfig *cache_config,
-    const std::shared_ptr<infinilm::config::ModelConfig> &text_config,
+HybridCacheTensors allocate_hybrid_cache_tensors(
+    const CacheConfig *cache_config,
+    const std::shared_ptr<infinilm::config::ModelConfig> &model_config,
     const backends::AttentionBackend &attention_backend) {
     if (nullptr == cache_config) {
         return {};
     }
-    if (nullptr == text_config) {
-        throw std::runtime_error("infinilm::models::qwen3_next::qwen3_next_allocate_kv_cache_tensors: text_config is null");
+    if (nullptr == model_config) {
+        throw std::runtime_error("allocate_hybrid_cache_tensors: model_config is null");
     }
 
-    const size_t num_hidden_layers = text_config->get<size_t>("num_hidden_layers");
-    const size_t head_dim = text_config->get<size_t>("head_dim");
-    const size_t num_key_value_heads = text_config->get<size_t>("num_key_value_heads");
-    const size_t max_position_embeddings = text_config->get<size_t>("max_position_embeddings");
-
-    const size_t linear_conv_kernel_dim = text_config->get<size_t>("linear_conv_kernel_dim");
-    const size_t linear_key_head_dim = text_config->get<size_t>("linear_key_head_dim");
-    const size_t linear_num_key_heads = text_config->get<size_t>("linear_num_key_heads");
-    const size_t linear_num_value_heads = text_config->get<size_t>("linear_num_value_heads");
-    const size_t linear_value_head_dim = text_config->get<size_t>("linear_value_head_dim");
-
-    const auto &dtype{text_config->get_dtype()};
-    const auto &kv_cache_dtype{text_config->get_kv_cache_dtype()};
-    const std::vector<std::string> layer_types = text_config->get<std::vector<std::string>>("layer_types");
+    const size_t num_hidden_layers = model_config->get<size_t>("num_hidden_layers");
+    const size_t head_dim = model_config->get<size_t>("head_dim");
+    const size_t num_key_value_heads = model_config->get<size_t>("num_key_value_heads");
+    const size_t max_position_embeddings = model_config->get<size_t>("max_position_embeddings");
+
+    const size_t linear_conv_kernel_dim = model_config->get<size_t>("linear_conv_kernel_dim");
+    const size_t linear_key_head_dim = model_config->get<size_t>("linear_key_head_dim");
+    const size_t linear_num_key_heads = model_config->get<size_t>("linear_num_key_heads");
+    const size_t linear_num_value_heads = model_config->get<size_t>("linear_num_value_heads");
+    const size_t linear_value_head_dim = model_config->get<size_t>("linear_value_head_dim");
+
+    const auto &dtype{model_config->get_dtype()};
+    const auto &kv_cache_dtype{model_config->get_kv_cache_dtype()};
+    const std::vector<std::string> layer_types = model_config->get<std::vector<std::string>>("layer_types");
+    if (layer_types.size() != num_hidden_layers) {
+        throw std::runtime_error(
+            "allocate_hybrid_cache_tensors: layer_types size must match num_hidden_layers");
+    }
 
     std::vector<infinicore::Tensor> kv_cache_vec;
     std::vector<infinicore::Tensor> conv_state_vec;
@@ -43,16 +43,25 @@ AllocatedHybridCache qwen3_next_allocate_cache_tensors(
     conv_state_vec.reserve(num_hidden_layers);
     ssm_state_vec.reserve(num_hidden_layers);
 
+    size_t mamba_state_pool_size = 0;
     auto allocate_linear_attention_cache = [&](size_t layer_idx, size_t pool_size) {
-        auto conv_state = cache::MambaCache::create_layer_conv_state(
+        if (mamba_state_pool_size == 0) {
+            mamba_state_pool_size = pool_size;
+        } else if (mamba_state_pool_size != pool_size) {
+            throw std::runtime_error(
+                "allocate_hybrid_cache_tensors: inconsistent mamba state pool size at layer "
+                + std::to_string(layer_idx));
+        }
+
+        auto conv_state = MambaCache::create_layer_conv_state(
             linear_key_head_dim,
             linear_value_head_dim,
             linear_num_key_heads,
             linear_num_value_heads,
             linear_conv_kernel_dim,
             dtype,
             pool_size);
-        auto ssm_state = cache::MambaCache::create_layer_ssm_state(
+        auto ssm_state = MambaCache::create_layer_ssm_state(
             linear_key_head_dim,
             linear_value_head_dim,
             linear_num_key_heads,
@@ -65,8 +74,8 @@ AllocatedHybridCache qwen3_next_allocate_cache_tensors(
         ssm_state_vec.push_back(std::move(ssm_state));
     };
 
-    auto allocate_static_full_attention_cache = [&](size_t layer_idx, const cache::StaticKVCacheConfig &config) {
-        auto kv_cache = cache::StaticKVCache::create_layer_kv_cache(
+    auto allocate_static_full_attention_cache = [&](size_t layer_idx, const StaticKVCacheConfig &config) {
+        auto kv_cache = StaticKVCache::create_layer_kv_cache(
             head_dim,
             head_dim,
             num_key_value_heads,
@@ -80,8 +89,8 @@ AllocatedHybridCache qwen3_next_allocate_cache_tensors(
         ssm_state_vec.emplace_back();
     };
 
-    auto allocate_paged_full_attention_cache = [&](size_t layer_idx, const cache::PagedKVCacheConfig &config) {
-        auto kv_cache = cache::PagedKVCache::create_layer_kv_cache(
+    auto allocate_paged_full_attention_cache = [&](size_t layer_idx, const PagedKVCacheConfig &config) {
+        auto kv_cache = PagedKVCache::create_layer_kv_cache(
             head_dim,
             head_dim,
             num_key_value_heads,
@@ -96,9 +105,9 @@ AllocatedHybridCache qwen3_next_allocate_cache_tensors(
 
     switch (attention_backend) {
     case backends::AttentionBackend::STATIC_ATTN: {
-        auto static_kv_cache_config = dynamic_cast<const cache::StaticKVCacheConfig *>(cache_config);
+        auto static_kv_cache_config = dynamic_cast<const StaticKVCacheConfig *>(cache_config);
         if (nullptr == static_kv_cache_config) {
-            throw std::runtime_error("infinilm::models::qwen3_next::qwen3_next_allocate_kv_cache_tensors: invalid static kv cache config type");
+            throw std::runtime_error("allocate_hybrid_cache_tensors: invalid static kv cache config type");
         }
 
         for (size_t layer_idx = 0; layer_idx < num_hidden_layers; ++layer_idx) {
@@ -108,7 +117,7 @@ AllocatedHybridCache qwen3_next_allocate_cache_tensors(
             } else if ("full_attention" == layer_type) {
                 allocate_static_full_attention_cache(layer_idx, *static_kv_cache_config);
             } else {
-                throw std::runtime_error("infinilm::models::qwen3_next::qwen3_next_allocate_kv_cache_tensors: unsupported layer_type '" + layer_type + "' for layer " + std::to_string(layer_idx));
+                throw std::runtime_error("allocate_hybrid_cache_tensors: unsupported layer_type '" + layer_type + "' for layer " + std::to_string(layer_idx));
             }
         }
         break;
@@ -117,9 +126,9 @@ AllocatedHybridCache qwen3_next_allocate_cache_tensors(
         ;
     }
     case backends::AttentionBackend::PAGED_ATTN: {
-        auto paged_kv_cache_config = dynamic_cast<const cache::PagedKVCacheConfig *>(cache_config);
+        auto paged_kv_cache_config = dynamic_cast<const PagedKVCacheConfig *>(cache_config);
         if (nullptr == paged_kv_cache_config) {
-            throw std::runtime_error("infinilm::models::qwen3_next::qwen3_next_allocate_kv_cache_tensors: invalid paged kv cache config type");
+            throw std::runtime_error("allocate_hybrid_cache_tensors: invalid paged kv cache config type");
         }
         const size_t mamba_pool_size = std::max<size_t>(2, paged_kv_cache_config->num_blocks() / 4);
 
@@ -130,18 +139,19 @@ AllocatedHybridCache qwen3_next_allocate_cache_tensors(
             } else if ("full_attention" == layer_type) {
                 allocate_paged_full_attention_cache(layer_idx, *paged_kv_cache_config);
             } else {
-                throw std::runtime_error("infinilm::models::qwen3_next::qwen3_next_allocate_kv_cache_tensors: unsupported layer_type '" + layer_type + "' for layer " + std::to_string(layer_idx));
+                throw std::runtime_error("allocate_hybrid_cache_tensors: unsupported layer_type '" + layer_type + "' for layer " + std::to_string(layer_idx));
             }
         }
         break;
     }
     default:
-        throw std::runtime_error("infinilm::models::qwen3_next::qwen3_next_allocate_kv_cache_tensors: Unsupported attention backend: " + std::to_string(static_cast<int>(attention_backend)));
+        throw std::runtime_error("allocate_hybrid_cache_tensors: Unsupported attention backend: " + std::to_string(static_cast<int>(attention_backend)));
     }
-    return AllocatedHybridCache{
+    return HybridCacheTensors{
         std::move(kv_cache_vec),
         std::move(conv_state_vec),
-        std::move(ssm_state_vec)};
+        std::move(ssm_state_vec),
+        mamba_state_pool_size};
 }
 
-} // namespace infinilm::models::qwen3_next
+} // namespace infinilm::cache

csrc/cache/hybrid_cache.hpp

@@ -0,0 +1,26 @@
+#pragma once
+
+#include "../backends/attention_backends.hpp"
+#include "../config/model_config.hpp"
+#include "kv_cache.hpp"
+#include "mamba_cache.hpp"
+
+#include <cstddef>
+#include <memory>
+#include <vector>
+
+namespace infinilm::cache {
+
+struct HybridCacheTensors {
+    std::vector<infinicore::Tensor> kv_cache_tensors;
+    std::vector<infinicore::Tensor> conv_state_tensors;
+    std::vector<infinicore::Tensor> ssm_state_tensors;
+    size_t mamba_state_pool_size{0};
+};
+
+HybridCacheTensors allocate_hybrid_cache_tensors(
+    const CacheConfig *cache_config,
+    const std::shared_ptr<infinilm::config::ModelConfig> &model_config,
+    const backends::AttentionBackend &attention_backend);
+
+} // namespace infinilm::cache

csrc/config/hybrid_model_config.cpp

@@ -0,0 +1,63 @@
+#include "hybrid_model_config.hpp"
+
+#include <stdexcept>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace infinilm::config {
+
+void prepare_hybrid_model_config(
+    const std::shared_ptr<ModelConfig> &model_config) {
+    if (model_config == nullptr) {
+        throw std::runtime_error(
+            "prepare_hybrid_model_config: model_config is null");
+    }
+
+    auto &config_json = model_config->get_config_json();
+    const size_t num_hidden_layers = model_config->get<size_t>("num_hidden_layers");
+
+    if (!config_json.contains("layer_types")) {
+        const size_t full_attention_interval = model_config->get<size_t>("full_attention_interval");
+        if (full_attention_interval == 0) {
+            throw std::runtime_error(
+                "prepare_hybrid_model_config: full_attention_interval must be positive");
+        }
+
+        std::vector<std::string> layer_types;
+        layer_types.reserve(num_hidden_layers);
+        for (size_t layer_idx = 0; layer_idx < num_hidden_layers; ++layer_idx) {
+            layer_types.push_back(
+                (layer_idx + 1) % full_attention_interval == 0
+                    ? "full_attention"
+                    : "linear_attention");
+        }
+        config_json["layer_types"] = std::move(layer_types);
+    }
+
+    const auto &layer_types = config_json["layer_types"];
+    if (!layer_types.is_array()
+        || layer_types.size() != num_hidden_layers) {
+        throw std::runtime_error(
+            "prepare_hybrid_model_config: layer_types size must match num_hidden_layers");
+    }
+    for (size_t layer_idx = 0; layer_idx < num_hidden_layers; ++layer_idx) {
+        if (!layer_types[layer_idx].is_string()) {
+            throw std::runtime_error(
+                "prepare_hybrid_model_config: layer_types entries must be strings");
+        }
+        const auto &layer_type = layer_types[layer_idx].get_ref<const std::string &>();
+        if (layer_type != "full_attention"
+            && layer_type != "linear_attention") {
+            throw std::runtime_error(
+                "prepare_hybrid_model_config: unsupported layer_type '"
+                + layer_type + "' at layer " + std::to_string(layer_idx));
+        }
+    }
+
+    if (!config_json.contains("attention_bias")) {
+        config_json["attention_bias"] = false;
+    }
+}
+
+} // namespace infinilm::config

csrc/config/hybrid_model_config.hpp

@@ -0,0 +1,12 @@
+#pragma once
+
+#include "model_config.hpp"
+
+#include <memory>
+
+namespace infinilm::config {
+
+void prepare_hybrid_model_config(
+    const std::shared_ptr<ModelConfig> &model_config);
+
+} // namespace infinilm::config

csrc/engine/compiler/paged_compiler.cpp

@@ -4,22 +4,14 @@
 
 #include <algorithm>
 #include <cstdint>
+#include <stdexcept>
 #include <vector>
 
 namespace infinilm::engine {
 namespace {
 
 bool has_mamba_cache(const infinilm::global_state::ForwardContext &forward_context) {
-    auto has_state = [](const std::vector<infinicore::Tensor> &state_vec) {
-        for (const auto &state : state_vec) {
-            if (state) {
-                return true;
-            }
-        }
-        return false;
-    };
-
-    return has_state(forward_context.conv_state_vec) || has_state(forward_context.ssm_state_vec);
+    return forward_context.mamba_state_pool_size > 0;
 }
 
 } // namespace
@@ -45,8 +37,36 @@ void PagedCompiler::compile() {
         size_t nblocks = dynamic_cast<const cache::PagedKVCacheConfig *>(model_->get_cache_config())->num_blocks();
         auto &forward_context = infinilm::global_state::get_forward_context();
         const bool has_mamba_state = has_mamba_cache(forward_context);
-
+        const auto &model_config = model_->get_model_config();
+        const size_t position_id_axes = model_config == nullptr
+                                          ? 1
+                                          : model_config->get_or<size_t>("position_id_axes", 1);
+        if (position_id_axes == 0) {
+            throw std::runtime_error("PagedCompiler: position_id_axes must be positive");
+        }
+        auto compile_batch_sizes = decode_batch_sizes_;
         size_t max_batch_size = *std::max_element(decode_batch_sizes_.begin(), decode_batch_sizes_.end());
+        if (has_mamba_state) {
+            if (forward_context.mamba_state_pool_size < 2) {
+                throw std::runtime_error(
+                    "PagedCompiler: mamba state pool must reserve row 0 and at least one request row");
+            }
+            const size_t max_mamba_batch_size = std::min(
+                max_batch_size, forward_context.mamba_state_pool_size - 1);
+            compile_batch_sizes.erase(
+                std::remove_if(
+                    compile_batch_sizes.begin(),
+                    compile_batch_sizes.end(),
+                    [max_mamba_batch_size](size_t b) {
+                        return b > max_mamba_batch_size;
+                    }),
+                compile_batch_sizes.end());
+            if (compile_batch_sizes.empty()) {
+                return;
+            }
+            max_batch_size = *std::max_element(
+                compile_batch_sizes.begin(), compile_batch_sizes.end());
+        }
         compiled_map_decode_.clear();
         block_tables_holder_ = infinicore::Tensor::empty(
             {nblocks * max_batch_size}, infinicore::DataType::I32, infinicore::context::getDevice());
@@ -55,7 +75,14 @@ void PagedCompiler::compile() {
         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());
+            // Models declare their position-id axes explicitly. Single-axis
+            // models retain the traditional [b] layout.
+            input.position_ids = infinicore::Tensor::empty(
+                position_id_axes > 1
+                    ? std::vector<size_t>{position_id_axes, b}
+                    : std::vector<size_t>{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());
@@ -126,7 +153,7 @@ void PagedCompiler::compile() {
             infinicore::context::syncStream();
         }
 
-        for (size_t b : decode_batch_sizes_) {
+        for (size_t b : compile_batch_sizes) {
             auto input = make_decode_input(b);
 
             barrier_->wait();

csrc/global_state/forward_context.hpp

@@ -62,6 +62,14 @@ struct ForwardContext {
     std::vector<infinicore::Tensor> kv_cache_vec;
     std::vector<infinicore::Tensor> conv_state_vec;
     std::vector<infinicore::Tensor> ssm_state_vec;
+    size_t mamba_state_pool_size{0};
+
+    void clear_model_caches() {
+        kv_cache_vec.clear();
+        conv_state_vec.clear();
+        ssm_state_vec.clear();
+        mamba_state_pool_size = 0;
+    }
 };
 
 void initialize_forward_context(ForwardContext &forward_context);