feat(mamba2): integrate model loading and recurrent inference

README.md

@@ -6,6 +6,71 @@
 
 当前版本依赖[`InfiniCore v0.2.9`](https://github.com/InfiniTensor/InfiniCore/releases/tag/v0.2.9)版本。
 
+### Mamba-2
+
+Mamba-2 requires an InfiniCore build containing `mamba2_scan`, the reduction
+graph support, and the graph allocator fixes accompanying this adaptation.
+Use matching InfiniLM/InfiniCore revisions; the v0.2.9 release alone is insufficient.
+On NVIDIA, enable Core's `nv-gpu`, `aten` (device dtype conversions), `ccl`
+(TP), and `graph` (Decode graphs) build options.
+On MetaX, use `metax-gpu=y use-mc=y aten=y ccl=y graph=y` and a matching
+MACA PyTorch build. Set `INFINIOP_METAX_ALLOW_TF32=0` before starting a process
+for strict FP32 comparisons; MetaX GEMM otherwise retains its TF32 default.
+
+The initial checkpoint is `state-spaces/mamba2-130m`. Prepare the native
+checkpoint and the `EleutherAI/gpt-neox-20b` tokenizer from local directories:
+
+```bash
+python scripts/prepare_mamba2_checkpoint.py \
+  --source /models/mamba2-130m-native \
+  --tokenizer /models/gpt-neox-tokenizer \
+  --output /models/mamba2-130m
+
+python examples/test_infer.py --device nvidia --model /models/mamba2-130m \
+  --enable-paged-attn --attn paged-attn --disable-prefix-caching \
+  --num-blocks 64 --max-new-tokens 32 --prompt "The capital of France is"
+```
+
+Add `--enable-graph` for eager Prefill plus Decode graphs, or `--tp 2` with
+two visible GPUs. The same prepared directory works with the existing service
+and benchmark entrypoints. This is a base language model for text continuation.
+For MetaX, replace `--device nvidia` with `--device metax`. The validated C500
+configuration is TP1 on a 50% compute / 32,000 MiB slice with MACA 3.5.3 and
+PyTorch 2.8.0+metax3.5.3.9; C500 TP2 has not been validated.
+
+The NVIDIA and MetaX implementations cover FP32, FP16 and BF16 activations with FP32
+residuals and SSM state. The preparation tool defaults to BF16 activation
+configuration and preserves source weight precision; setting `torch_dtype`
+in the prepared config selects FP16 or FP32. The loader retains the source
+precision of state parameters and normalization weights in FP32.
+
+The model uses the existing paged **request-state** interface, with separate
+convolution and SSM state rather than Attention KV pages. The state pool has
+`max(2, num_blocks // 4)` rows, including a reserved zero row, so its request
+capacity is one less than that value. For 130M, each row occupies about
+18.25 MiB in TP1 BF16. The scheduler still applies its logical page budget.
+
+Current scope is pure Mamba-2 with one B/C group, convolution width 4,
+head-wise D, gated RMSNorm after gating, and unbounded time steps. PP, hybrid
+Attention/SSM layers, quantization, prefix caching, remote state transfer,
+speculative rollback and scheduler chunked Prefill are outside this adaptation.
+The SSD kernel's internal chunks do not enable scheduler chunking.
+
+Device kernels live in InfiniCore under `src/infiniop/ops/mamba2_scan/`.
+NVIDIA and MetaX share the scan kernels through runtime-specific launch wrappers
+and the same indexed-state and workspace contract. Moore and Ascend backends
+are not implemented by this adaptation. Model math stays on the device.
+
+Targeted validation can be run with a prepared checkpoint:
+
+```bash
+INFINILM_MAMBA2_MODEL=/models/mamba2-130m \
+INFINILM_MAMBA2_TP=1 INFINILM_MAMBA2_GRAPH=1 \
+NVIDIA_TF32_OVERRIDE=0 \
+INFINIOP_METAX_ALLOW_TF32=0 \
+python -m pytest test/models/mamba2 -q
+```
+
 ## 使用方式
 #### 一、编译并安装 `InfiniCore`
 编译并安装 `InfiniCore`， 详情见 InfiniCore的 [`README`](https://github.com/InfiniTensor/InfiniCore) :

csrc/engine/compiler/paged_compiler.cpp

@@ -4,6 +4,7 @@
 
 #include <algorithm>
 #include <cstdint>
+#include <numeric>
 #include <stdexcept>
 #include <vector>
 
@@ -23,6 +24,43 @@ bool has_mamba_cache(const infinilm::global_state::ForwardContext &forward_conte
     return has_state(forward_context.conv_state_vec) || has_state(forward_context.ssm_state_vec);
 }
 
+class CacheStateGuard {
+public:
+    void save(const infinicore::Tensor &state, size_t rows) {
+        if (!state) {
+            return;
+        }
+        save_region(state->narrow({{0, 1, rows}}));
+    }
+
+    void save_region(const infinicore::Tensor &region) {
+        auto backup = infinicore::Tensor::empty(region->shape(), region->dtype(), region->device());
+        backup->copy_from(region);
+        saved_.emplace_back(region, backup);
+    }
+
+    void restore() {
+        for (auto &[region, backup] : saved_) {
+            region->copy_from(backup);
+        }
+        if (!saved_.empty()) {
+            infinicore::context::syncStream();
+            saved_.clear();
+        }
+    }
+
+    ~CacheStateGuard() {
+        try {
+            restore();
+        } catch (const std::exception &error) {
+            spdlog::error("Failed to restore request states after graph capture: {}", error.what());
+        }
+    }
+
+private:
+    std::vector<std::pair<infinicore::Tensor, infinicore::Tensor>> saved_;
+};
+
 } // namespace
 
 PagedCompiler::PagedCompiler(const std::shared_ptr<InfinilmModel> &model, RankBarrier *barrier)
@@ -70,8 +108,40 @@ void PagedCompiler::compile() {
             throw std::runtime_error("PagedCompiler: position_id_axes must be positive");
         }
 
-        size_t max_batch_size = *std::max_element(decode_batch_sizes_.begin(), decode_batch_sizes_.end());
+        infinicore::context::syncStream();
         compiled_map_decode_.clear();
+        if (decode_batch_sizes_.empty()) {
+            return;
+        }
+        size_t max_batch_size = *std::max_element(decode_batch_sizes_.begin(), decode_batch_sizes_.end());
+        if (has_mamba_state) {
+            for (const auto *states : {&forward_context.conv_state_vec, &forward_context.ssm_state_vec}) {
+                for (const auto &state : *states) {
+                    if (state) {
+                        const size_t capacity = state->size(0) == 0 ? 0 : state->size(0) - 1;
+                        max_batch_size = std::min(max_batch_size, capacity);
+                    }
+                }
+            }
+        }
+        if (max_batch_size == 0) {
+            return;
+        }
+        CacheStateGuard state_guard;
+        if (has_mamba_state) {
+            for (const auto *states : {&forward_context.conv_state_vec, &forward_context.ssm_state_vec}) {
+                for (const auto &state : *states) {
+                    state_guard.save(state, max_batch_size);
+                }
+            }
+        }
+        // Warmup and capture write into physical page zero. Preserve it so
+        // recapturing also remains safe while a request owns that page.
+        for (const auto &kv : forward_context.kv_cache_vec) {
+            if (kv) {
+                state_guard.save_region(kv->narrow({{1, 0, 1}}));
+            }
+        }
         block_tables_holder_ = infinicore::Tensor::empty(
             {nblocks * max_batch_size}, infinicore::DataType::I32, infinicore::context::getDevice());
         set_zeros(block_tables_holder_);
@@ -109,7 +179,8 @@ void PagedCompiler::compile() {
                 input.mamba_final_state_indices = infinicore::Tensor::empty(
                     {b}, infinicore::DataType::I32, infinicore::context::getDevice());
                 std::vector<int32_t> init_state_indices_vec(b, 0);
-                std::vector<int32_t> final_state_indices_vec(b, 1);
+                std::vector<int32_t> final_state_indices_vec(b);
+                std::iota(final_state_indices_vec.begin(), final_state_indices_vec.end(), 1);
                 infinicore::context::memcpyH2D(
                     input.mamba_init_state_indices.value()->data(),
                     init_state_indices_vec.data(),
@@ -155,6 +226,9 @@ void PagedCompiler::compile() {
         }
 
         for (size_t b : decode_batch_sizes_) {
+            if (b > max_batch_size) {
+                continue;
+            }
             auto input = make_decode_input(b);
 
             barrier_->wait();
@@ -176,6 +250,7 @@ void PagedCompiler::compile() {
 
             compiled_map_decode_[b] = CompiledResult{std::move(input), std::make_tuple(graph, shared_output)};
         }
+        state_guard.restore();
     }
 }
 

csrc/layers/quantization/none_quantization.cpp

@@ -82,12 +82,13 @@ std::vector<SplitParam> NoneQuantization::split_params(
     std::vector<SplitParam> result;
     auto weight_it = params.find("weight");
     auto bias_it = params.find("bias");
+    const int weight_dim = weight_prepacked_ ? 1 - narrow_dim : narrow_dim;
 
     for (const auto &s : splits) {
         result.push_back({s.prefix + ".weight",
                           infinicore::nn::Parameter(
-                              weight_it->second->narrow({{static_cast<size_t>(narrow_dim), s.start, s.size}}),
-                              narrow_dim, tp_rank, tp_size, s.num_shards)});
+                              weight_it->second->narrow({{static_cast<size_t>(weight_dim), s.start, s.size}}),
+                              weight_dim, tp_rank, tp_size, s.num_shards)});
         if (bias_it != params.end()) {
             result.push_back({s.prefix + ".bias",
                               infinicore::nn::Parameter(
@@ -104,7 +105,7 @@ std::shared_ptr<BaseQuantization> NoneQuantization::process_weights_after_loadin
     int /*split_dim*/) const {
 
     // Controlled by --pre-transpose CLI flag, default off.
-    if (!global_state::get_infinilm_config().pre_transpose) {
+    if (!global_state::get_infinilm_config().pre_transpose || weight_prepacked_) {
         return nullptr;
     }
 
@@ -115,15 +116,13 @@ std::shared_ptr<BaseQuantization> NoneQuantization::process_weights_after_loadin
         // subsequent forwards can feed it directly to GEMM.
         params["weight"] = weight_it->second->permute({1, 0})->contiguous();
 
-        // Mark as pre-packed so forward() uses linear_packed.
-        weight_prepacked_ = true;
+        // A quantization object may be shared by several unprocessed linears.
+        auto packed = std::make_shared<NoneQuantization>(get_config());
+        packed->weight_prepacked_ = true;
+        return packed;
     }
 
-    // Must return non-null so that BaseLinear::process_weights_after_loading
-    // writes the modified params back into parameters_.
-    // Returning shared_from_this() triggers the "quantization changed" path
-    // which calls parameters_.clear() + re-insert from params.
-    return std::const_pointer_cast<BaseQuantization>(shared_from_this());
+    return nullptr;
 }
 
 } // namespace infinilm::quantization

csrc/layers/quantization/none_quantization.hpp

@@ -6,7 +6,7 @@ namespace infinilm::quantization {
 class NoneQuantization : public BaseQuantization {
 public:
     explicit NoneQuantization(const nlohmann::json &quant_config)
-        : BaseQuantization(quant_config){};
+        : BaseQuantization(quant_config) {};
 
     NoneQuantization();
 
@@ -40,15 +40,14 @@ class NoneQuantization : public BaseQuantization {
         int narrow_dim,
         int tp_rank, int tp_size, int tp_num_heads) const override;
 
-    // Ascend: pre-pack weight to [IC, OC] after loading to skip runtime permute.
-    // Returns shared_from_this() only on Ascend; nullptr otherwise (no-op).
+    // Pre-pack to [IC, OC] when enabled and return a per-linear layout state.
     std::shared_ptr<BaseQuantization> process_weights_after_loading(
         ParamsMap &params,
         const infinicore::Device &device,
         int split_dim = -1) const override;
 
 private:
-    mutable bool weight_prepacked_ = false; // true when weight was pre-packed for Ascend
+    bool weight_prepacked_ = false;
 };
 
 } // namespace infinilm::quantization

csrc/models/mamba2/mamba2_for_causal_lm.cpp

@@ -0,0 +1,132 @@
+#include "mamba2_for_causal_lm.hpp"
+#include "../models_registry.hpp"
+#include <algorithm>
+#include <cmath>
+
+namespace infinilm::models::mamba2 {
+namespace {
+
+class TiedLMHead final : public layers::linear::ReplicatedLinear {
+public:
+    explicit TiedLMHead(const infinicore::Tensor &weight)
+        : layers::linear::ReplicatedLinear(weight->size(1), weight->size(0), false, weight->dtype(), weight->device()) {
+        register_parameter("weight", infinicore::nn::Parameter(weight));
+    }
+
+    // Preserve the embedding's shared layout when other projections are packed.
+    void process_weights_after_loading() override {}
+};
+
+} // namespace
+
+std::shared_ptr<config::ModelConfig> create_mamba2_model_config(std::shared_ptr<config::ModelConfig> config) {
+    auto &j = config->get_config_json();
+    if (j.at("model_type") != "mamba2") {
+        throw std::runtime_error("Expected `model_type=mamba2`.");
+    }
+    for (const auto *key : {"hidden_size", "num_hidden_layers", "vocab_size", "num_heads", "head_dim", "state_size"}) {
+        if (config->get<int64_t>(key) <= 0) {
+            throw std::runtime_error(std::string("Mamba-2 requires positive `") + key + "`.");
+        }
+    }
+    const auto hidden = config->get<size_t>("hidden_size");
+    const auto intermediate = hidden * config->get_or<size_t>("expand", 2);
+    const auto heads = config->get<size_t>("num_heads");
+    const auto head_dim = config->get<size_t>("head_dim");
+    if (config->get_or<int64_t>("expand", 2) <= 0 || heads * head_dim != intermediate
+        || config->get<size_t>("state_size") > 256
+        || config->get_or<std::string>("hidden_act", "silu") != "silu"
+        || config->get_or<size_t>("intermediate_size", intermediate) != intermediate
+        || config->get_or<size_t>("n_groups", 1) != 1
+        || config->get_or<size_t>("conv_kernel", 4) != 4
+        || config->get_or<bool>("use_bias", false)
+        || !config->get_or<bool>("use_conv_bias", true)
+        || config->get_or<bool>("norm_before_gate", false)
+        || !config->get_or<bool>("residual_in_fp32", true)
+        || !config->get_or<bool>("rms_norm", true)
+        || !config->get_or<bool>("rmsnorm", true)
+        || config->get_or<bool>("D_has_hdim", false)
+        || config->get_or<size_t>("d_ssm", intermediate) != intermediate
+        || config->get_or<size_t>("d_intermediate", 0) != 0
+        || (j.contains("attn_layer_idx") && !j["attn_layer_idx"].empty())
+        || (j.contains("quantization_config") && !j["quantization_config"].empty())) {
+        throw std::runtime_error("Unsupported Mamba-2 configuration; use the checkpoint preparation tool.");
+    }
+    for (const auto *key : {"dt_limit", "time_step_limit"}) {
+        if (j.contains(key)) {
+            throw std::runtime_error("Explicit time-step limits are not supported; omit the field for the unbounded Mamba-2 scan.");
+        }
+    }
+    j["intermediate_size"] = intermediate;
+    j["layer_norm_epsilon"] = j.value("layer_norm_epsilon", 1e-5);
+    j["rms_norm_eps"] = j["layer_norm_epsilon"];
+    const double epsilon = config->get<double>("layer_norm_epsilon");
+    if (!std::isfinite(epsilon) || epsilon <= 0) {
+        throw std::runtime_error("Mamba-2 requires a finite positive normalization epsilon.");
+    }
+    return config;
+}
+
+Mamba2ForCausalLM::Mamba2ForCausalLM(std::shared_ptr<config::ModelConfig> config, const infinicore::Device &device)
+    : TextCausalLM(config, device) {
+    if (config->get_or<bool>("tie_word_embeddings", true)) {
+        lm_head_ = register_module<TiedLMHead>("lm_head", model_->embedding_weight());
+    }
+}
+
+Mamba2Model::Mamba2Model(std::shared_ptr<config::ModelConfig> config, const infinicore::Device &device)
+    : dtype_(config->get_dtype()) {
+    const auto &rank = global_state::get_tensor_model_parallel_rank_info();
+    if (rank.pp_size != 1) {
+        throw std::runtime_error("Mamba-2 currently requires `pp_size=1`.");
+    }
+    const auto hidden = config->get<size_t>("hidden_size");
+    INFINICORE_NN_MODULE_INIT(embeddings, config->get<size_t>("vocab_size"), hidden, std::nullopt, dtype_, device);
+    for (size_t i = 0; i < config->get<size_t>("num_hidden_layers"); ++i) {
+        layers_.push_back(register_module<Mamba2Block>("layers." + std::to_string(i), config, i, device));
+    }
+    INFINICORE_NN_MODULE_INIT(norm_f, hidden, config->get<double>("layer_norm_epsilon"), infinicore::DataType::F32, device);
+}
+
+infinicore::Tensor Mamba2Model::forward(const InfinilmModel::Input &input) const {
+    if (!input.input_offsets || !input.mamba_init_state_indices || !input.mamba_final_state_indices) {
+        throw std::runtime_error("Mamba-2 requires packed offsets and request state indices.");
+    }
+    auto ids = input.input_ids.value()->view({1, input.input_ids.value()->numel()});
+    auto residual = cast_activation(embeddings_->forward(ids), infinicore::DataType::F32);
+    for (const auto &layer : layers_) {
+        residual = layer->forward(residual);
+    }
+    return cast_activation(norm_f_->forward(residual), dtype_);
+}
+
+void Mamba2ForCausalLM::reset_cache(const cache::CacheConfig *config) {
+    auto &context = global_state::get_forward_context();
+    context.kv_cache_vec.clear();
+    context.conv_state_vec.clear();
+    context.ssm_state_vec.clear();
+    cache_config_ = config ? config->unique_copy() : nullptr;
+    if (config == nullptr) {
+        return;
+    }
+    const auto *paged = dynamic_cast<const cache::PagedKVCacheConfig *>(config);
+    if (paged == nullptr) {
+        throw std::runtime_error("Mamba-2 requires the paged request-state cache interface.");
+    }
+    const size_t pool = std::max<size_t>(2, paged->num_blocks() / 4);
+    const auto heads = model_config_->get<size_t>("num_heads") / global_state::get_tensor_model_parallel_world_size();
+    const auto head_dim = model_config_->get<size_t>("head_dim");
+    const auto state_size = model_config_->get<size_t>("state_size");
+    const auto conv_dim = heads * head_dim + 2 * state_size;
+    const auto device = infinicore::context::getDevice();
+    for (size_t i = 0; i < model_config_->get<size_t>("num_hidden_layers"); ++i) {
+        context.conv_state_vec.push_back(infinicore::Tensor::zeros({pool, conv_dim, 3}, model_config_->get_dtype(), device));
+        context.ssm_state_vec.push_back(infinicore::Tensor::zeros({pool, heads, head_dim, state_size}, infinicore::DataType::F32, device));
+    }
+}
+
+} // namespace infinilm::models::mamba2
+
+namespace {
+INFINILM_REGISTER_CAUSAL_LM_MODEL(mamba2, infinilm::models::mamba2::Mamba2ForCausalLM, infinilm::models::mamba2::create_mamba2_model_config);
+} // namespace