mimo v2.5 support

rtp_llm/config/model_config.py

@@ -14,7 +14,7 @@
     init_quant_config,
 )
 from rtp_llm.multimodal.multimodal_mixin_register import get_multimodal_mixin_cls
-from rtp_llm.ops import DataType, KvCacheDataType
+from rtp_llm.ops import DataType, HybridAttentionType, KvCacheDataType
 from rtp_llm.ops import ModelConfig as CppModelConfig
 from rtp_llm.ops import TaskType
 from rtp_llm.utils.base_model_datatypes import VitParameters
@@ -78,6 +78,7 @@ class ModelConfig(CppModelConfig):
         "phy2log_path",
         "lora_infos",
         "headwise_config",
+        "attention_value_scale",
     }
 
     # Known C++ ModelConfig members (from ModelConfig.h)
@@ -251,6 +252,9 @@ def _eval_kv_cache_mem_size(self) -> float:
         # Get kv_cache_dtype from attn_config
         kv_cache_dtype_enum = self.attn_config.kv_cache_dtype
         kv_cache_bytes = 1 if kv_cache_dtype_enum == KvCacheDataType.FP8 else 2
+        hybrid_config = self.hybrid_attention_config
+        if hybrid_config.enable_hybrid_attention:
+            return self._eval_hybrid_kv_cache_mem_size(kv_cache_bytes)
         kv_cache_size = (
             2
             * self.num_layers
@@ -261,6 +265,43 @@ def _eval_kv_cache_mem_size(self) -> float:
         )
         return kv_cache_size
 
+    def _eval_hybrid_kv_cache_mem_size(self, kv_cache_bytes: int) -> float:
+        """KV cache size for a model whose layers do not share one attention shape.
+
+        Global and sliding-window layers differ in KV head count, K and V may differ in
+        head dimension, and a windowed layer never holds more than its window, so the
+        homogeneous ``2 * num_layers * kv_head_num * size_per_head`` estimate is wrong on
+        all three counts.
+        """
+        swa_config = self.hybrid_attention_config.swa_attention_config
+        pattern = self.hybrid_attention_config.hybrid_attention_types
+        ga_layers = sum(1 for t in pattern if t != HybridAttentionType.SLIDING_WINDOW)
+        swa_layers = len(pattern) - ga_layers
+
+        k_head_size = self.attn_config.size_per_head
+        v_head_size = self.attn_config.v_size_per_head or k_head_size
+        kv_head_size = k_head_size + v_head_size
+
+        swa_kv_head_num = swa_config.swa_kv_head_num or self.attn_config.kv_head_num
+        # A windowed layer only ever keeps the last window_size tokens resident.
+        swa_tokens = (
+            min(swa_config.window_size, self.max_seq_len)
+            if swa_config.window_size > 0
+            else self.max_seq_len
+        )
+
+        ga_bytes = (
+            ga_layers
+            * self.attn_config.kv_head_num
+            * kv_head_size
+            * kv_cache_bytes
+            * self.max_seq_len
+        )
+        swa_bytes = (
+            swa_layers * swa_kv_head_num * kv_head_size * kv_cache_bytes * swa_tokens
+        )
+        return ga_bytes + swa_bytes
+
     def _eval_runtime_buffer_mem_size(self) -> float:
         """Evaluate runtime buffer memory size."""
         input_buffer = self.max_seq_len * self.hidden_size
@@ -566,6 +607,10 @@ def __init__(self, *args, **kwargs):
         self.mm_related_params = VitParameters()
         self.quant_config = None
 
+        # Checkpoint-provided scale applied to V before attention. None means the
+        # checkpoint does not scale V, which is the case for every model but MiMo V2.5.
+        self.attention_value_scale: Optional[float] = None
+
     def apply_override_args(self, json_model_override_args: str) -> None:
         """Apply model override arguments to ModelConfig.
 

rtp_llm/cpp/cache/CacheConfig.cc

@@ -28,6 +28,7 @@ std::string cacheGroupPolicySummary(const CacheGroupPolicy& policy) {
        << ", charge_to_paged_budget=" << policy.charge_to_paged_budget
        << ", memory_placement=" << static_cast<int>(policy.memory_placement)
        << ", active_tail_blocks=" << policy.active_tail_blocks
+       << ", prefix_reuse_window_tokens=" << policy.prefix_reuse_window_tokens
        << ", validate_tail_blocks=" << policy.validate_tail_blocks
        << ", cp_mapping=" << static_cast<int>(policy.cp_mapping) << ", cp_slice=" << static_cast<int>(policy.cp_slice)
        << '}';
@@ -130,8 +131,10 @@ bool CacheConfig::samePolicy(const CacheGroupPolicy& lhs, const CacheGroupPolicy
            && lhs.evict_policy == rhs.evict_policy && lhs.reservable == rhs.reservable
            && lhs.explicit_block_num == rhs.explicit_block_num
            && lhs.charge_to_paged_budget == rhs.charge_to_paged_budget && lhs.memory_placement == rhs.memory_placement
-           && lhs.active_tail_blocks == rhs.active_tail_blocks && lhs.validate_tail_blocks == rhs.validate_tail_blocks
-           && lhs.cp_mapping == rhs.cp_mapping && lhs.cp_slice == rhs.cp_slice;
+           && lhs.active_tail_blocks == rhs.active_tail_blocks
+           && lhs.prefix_reuse_window_tokens == rhs.prefix_reuse_window_tokens
+           && lhs.validate_tail_blocks == rhs.validate_tail_blocks && lhs.cp_mapping == rhs.cp_mapping
+           && lhs.cp_slice == rhs.cp_slice;
 }
 
 void CacheConfig::setTopology(std::vector<GroupBase> new_groups, std::vector<LayerBase> new_layers) {

rtp_llm/cpp/cache/CacheGroupType.h

@@ -55,9 +55,13 @@ struct CacheGroupPolicy {
     bool                 charge_to_paged_budget = false;
     CacheMemoryPlacement memory_placement       = CacheMemoryPlacement::DEVICE;
     uint32_t             active_tail_blocks     = 0;
-    bool                 validate_tail_blocks   = true;
-    CpBlockMappingMode   cp_mapping             = CpBlockMappingMode::NONE;
-    CpBlockSliceMode     cp_slice               = CpBlockSliceMode::NONE;
+    // Number of token positions covered by the SWA prefix-reuse window. This
+    // is intentionally separate from active_tail_blocks: the latter controls
+    // allocation/retention, while this value controls cache-key matching.
+    uint32_t           prefix_reuse_window_tokens = 0;
+    bool               validate_tail_blocks       = true;
+    CpBlockMappingMode cp_mapping                 = CpBlockMappingMode::NONE;
+    CpBlockSliceMode   cp_slice                   = CpBlockSliceMode::NONE;
 };
 
 // One cache-store registration step: pair a cache key from the full logical

rtp_llm/cpp/cache/HybridKVCacheAllocator.cc

@@ -1,6 +1,7 @@
 #include "rtp_llm/cpp/cache/HybridKVCacheAllocator.h"
 
 #include <algorithm>
+#include <cstdint>
 #include <unordered_map>
 #include <unordered_set>
 
@@ -56,6 +57,32 @@ BlockIndicesType validBlocksAfter(const BlockIndicesType& blocks, size_t begin)
     return valid;
 }
 
+// Return the first canonical cache-key block needed by the SWA window when
+// `end` is the last reused block.  Cache keys are in canonical units under CP,
+// so one key covers cp_scale raw blocks.  A zero window keeps the historical
+// single-tail behavior used by generic SWA groups that do not declare a
+// prefix-reuse window; MiMo supplies the explicit window in its descriptor.
+int swaMatchBegin(const KVCacheGroup& group, int end, int cp_scale) {
+    const auto window_tokens = group.policy().prefix_reuse_window_tokens;
+    if (window_tokens == 0) {
+        return end;
+    }
+    if (window_tokens == 1) {
+        return end + 1;
+    }
+
+    const int     block_tokens    = std::max(1, group.seqSizePerBlock() * std::max(cp_scale, 1));
+    const int64_t first_new_token = static_cast<int64_t>(end + 1) * block_tokens;
+    const int64_t history_tokens  = static_cast<int64_t>(window_tokens - 1);
+    const int64_t history_start   = std::max<int64_t>(0, first_new_token - history_tokens);
+    return std::max(0, static_cast<int>(history_start / block_tokens));
+}
+
+struct SwaMatch {
+    int              begin = 0;
+    BlockIndicesType blocks;
+};
+
 }  // namespace
 
 bool HybridKVCacheAllocator::skipReuseCacheGroup(int gid) const {
@@ -102,14 +129,14 @@ int HybridKVCacheAllocator::reuseCache(const CacheKeysType&                 cach
         full_matched_blocks[static_cast<size_t>(gid)] = std::move(match_result.block_indices);
     }
 
-    int                           pos = min_full_reuse_blocks - 1;
-    std::vector<BlockIdxType>     linear_tail_blocks(linear_group_ids_.size(), NULL_BLOCK_IDX);
-    std::vector<BlockIndicesType> swa_tail_blocks(swa_group_ids_.size());
-    const bool                    has_tail_groups = !linear_group_ids_.empty() || !swa_group_ids_.empty();
+    int                       pos = min_full_reuse_blocks - 1;
+    std::vector<BlockIdxType> linear_tail_blocks(linear_group_ids_.size(), NULL_BLOCK_IDX);
+    std::vector<SwaMatch>     swa_tail_matches(swa_group_ids_.size());
+    const bool                has_tail_groups = !linear_group_ids_.empty() || !swa_group_ids_.empty();
     for (; pos >= 0 && has_tail_groups; --pos) {
-        bool                          all_tail_groups_matched = true;
-        std::vector<BlockIdxType>     candidate_linear_tail_blocks(linear_group_ids_.size(), NULL_BLOCK_IDX);
-        std::vector<BlockIndicesType> candidate_swa_tail_blocks(swa_group_ids_.size());
+        bool                      all_tail_groups_matched = true;
+        std::vector<BlockIdxType> candidate_linear_tail_blocks(linear_group_ids_.size(), NULL_BLOCK_IDX);
+        std::vector<SwaMatch>     candidate_swa_tail_matches(swa_group_ids_.size());
         for (size_t i = 0; i < linear_group_ids_.size(); ++i) {
             const int gid = linear_group_ids_[i];
             auto      result =
@@ -128,17 +155,28 @@ int HybridKVCacheAllocator::reuseCache(const CacheKeysType&                 cach
             if (skipReuseCacheGroup(gid)) {
                 continue;
             }
-            auto result =
-                kv_cache_groups_[static_cast<size_t>(gid)]->matchSingleKey(cache_keys[static_cast<size_t>(pos)]);
-            if (result.block_indices.empty()) {
-                all_tail_groups_matched = false;
+            const auto& group = *kv_cache_groups_[static_cast<size_t>(gid)];
+            const int   begin = swaMatchBegin(group, pos, cp_scale);
+            auto&       match = candidate_swa_tail_matches[i];
+            match.begin       = begin;
+            if (begin <= pos) {
+                match.blocks.reserve(static_cast<size_t>(pos - begin + 1));
+                for (int key_pos = begin; key_pos <= pos; ++key_pos) {
+                    auto result = group.matchSingleKey(cache_keys[static_cast<size_t>(key_pos)]);
+                    if (result.block_indices.empty()) {
+                        all_tail_groups_matched = false;
+                        break;
+                    }
+                    match.blocks.push_back(result.block_indices[0]);
+                }
+            }
+            if (!all_tail_groups_matched) {
                 break;
             }
-            candidate_swa_tail_blocks[i].push_back(result.block_indices[0]);
         }
         if (all_tail_groups_matched) {
             linear_tail_blocks = std::move(candidate_linear_tail_blocks);
-            swa_tail_blocks    = std::move(candidate_swa_tail_blocks);
+            swa_tail_matches   = std::move(candidate_swa_tail_matches);
             break;
         }
     }
@@ -176,10 +214,17 @@ int HybridKVCacheAllocator::reuseCache(const CacheKeysType&                 cach
         if (skipReuseCacheGroup(gid)) {
             continue;
         }
-        const size_t tail_begin =
-            static_cast<size_t>(std::max(group_reuse_len - static_cast<int>(swa_tail_blocks[i].size()), 0));
-        for (size_t j = 0; j < swa_tail_blocks[i].size(); ++j) {
-            kv_resource.mutableBlockIds(0, gid).setAt(tail_begin + j, swa_tail_blocks[i][j]);
+        const auto& match = swa_tail_matches[i];
+        for (size_t j = 0; j < match.blocks.size(); ++j) {
+            const int canonical_pos = match.begin + static_cast<int>(j);
+            // Compact-last-rank SWA uses one slot per canonical key.  The
+            // non-compact layout keeps cp_size logical slots per key, and the
+            // canonical key owns the last slot in that group.
+            const int logical_pos =
+                cpCompactSwaGroup(gid, cp_mapper) ? canonical_pos : (canonical_pos + 1) * cp_scale - 1;
+            if (logical_pos >= 0 && logical_pos < group_reuse_len) {
+                kv_resource.mutableBlockIds(0, gid).setAt(static_cast<size_t>(logical_pos), match.blocks[j]);
+            }
         }
     }
     return reuse_blocks_len;
@@ -259,19 +304,15 @@ MallocResult HybridKVCacheAllocator::initMallocForCommonLen(const MallocInfo& ma
         original_sizes[static_cast<size_t>(gid)] = kv_resource->blocksNum(0, gid);
     }
     for (int gid = 0; gid < kv_resource->groupNums(); ++gid) {
-        auto&     block_ids_0   = kv_resource->mutableBlockIds(0, gid);
-        const int group_seq_len = cpEffectiveSeqLenForGroup(cp_mapper, config_, gid, common_seq_len);
-        const auto& group = kv_cache_groups_[static_cast<size_t>(gid)];
+        auto&       block_ids_0   = kv_resource->mutableBlockIds(0, gid);
+        const int   group_seq_len = cpEffectiveSeqLenForGroup(cp_mapper, config_, gid, common_seq_len);
+        const auto& group         = kv_cache_groups_[static_cast<size_t>(gid)];
         // Snapshot the slot count before the call so a failure can report this
         // group's exact physical request in the error_code=602 record.
         const int blocks_before = static_cast<int>(block_ids_0.blocksNum());
         if (!group->malloc(block_ids_0, group_seq_len, malloc_info.reuse_cache, 0)) {
-            logMallocFailure(malloc_info,
-                             "init_group_malloc",
-                             0,
-                             gid,
-                             false,
-                             group->needBlocksNum(group_seq_len, blocks_before, 0));
+            logMallocFailure(
+                malloc_info, "init_group_malloc", 0, gid, false, group->needBlocksNum(group_seq_len, blocks_before, 0));
             rollbackInitMalloc(*kv_resource, referenced_blocks, original_sizes);
             return {false, 0};
         }

rtp_llm/cpp/cache/HybridPoolConfigCreator.cc

@@ -70,10 +70,16 @@ void validateHybridPoolDescs(const ModelConfig& model_config, uint32_t kernel_to
     }
 }
 
-uint32_t mhaLocalKvHeadNum(const ModelConfig& model_config, const ParallelismConfig& parallelism_config) {
+uint32_t mhaLocalKvHeadNum(const KVCacheSpecDesc&   desc,
+                           const ModelConfig&       model_config,
+                           const ParallelismConfig& parallelism_config) {
     const auto     attn_tp = std::max<int64_t>(1, parallelism_config.get_attn_tp_size());
     const uint32_t tp      = static_cast<uint32_t>(attn_tp);
-    const uint32_t kv      = static_cast<uint32_t>(model_config.attn_config.kv_head_num);
+    // Must match MHAKVCacheSpec::build(): a per-desc override describes a layer kind whose
+    // KV head count differs from the model-wide one, and the group's local head count has
+    // to agree with the spec that sized its blocks.
+    const uint32_t kv = desc.kv_head_num_override != 0 ? desc.kv_head_num_override :
+                                                         static_cast<uint32_t>(model_config.attn_config.kv_head_num);
     RTP_LLM_CHECK_WITH_INFO(kv > 0, "local kv head num requires positive kv_head_num");
     return (kv % tp == 0) ? kv / tp : kv / std::gcd(kv, tp);
 }
@@ -98,7 +104,7 @@ uint32_t localKvHeadNumForDesc(const KVCacheSpecDesc&   desc,
                                const ParallelismConfig& parallelism_config) {
     switch (desc.cache_type) {
         case KVCacheSpecType::MultiHeadAttention:
-            return mhaLocalKvHeadNum(model_config, parallelism_config);
+            return mhaLocalKvHeadNum(desc, model_config, parallelism_config);
         case KVCacheSpecType::LinearAttention:
             return linearLocalKvHeadNum(model_config, parallelism_config);
         case KVCacheSpecType::MultiHeadLatentAttention:
@@ -239,7 +245,14 @@ void setupIndependentPoolSizes(CacheConfig& config, bool is_mtp) {
         group_kv_block_stride_bytes[gid] = kv_stride;
         group_kv_scale_stride_bytes[gid] = scale_stride;
         const auto type                  = config.typeForGroup(gid);
-        const bool is_paged_group        = type == CacheGroupType::FULL || type == CacheGroupType::LINEAR;
+        // A sliding-window group counts too when it is an ordinary paged pool: it draws
+        // its block count from the same global budget (finalizeBlockNums gives it
+        // global_block_num / linear_step), so its bytes belong in the per-block cost. DSv4
+        // instead backs its window with a fixed-allocation state cache, which is sized
+        // outside the paged budget and must stay excluded.
+        const bool is_state_cache = spec->type == KVCacheSpecType::OpaqueState;
+        const bool is_paged_group = type == CacheGroupType::FULL || type == CacheGroupType::LINEAR
+                                    || (type == CacheGroupType::SWA && !is_state_cache);
         if (is_paged_group && !config.usesExplicitIndependentBlocks(gid)) {
             total_kv_block_bytes += static_cast<size_t>(layer_count) * kv_stride;
             total_scale_block_bytes += static_cast<size_t>(layer_count) * scale_stride;

rtp_llm/cpp/cache/KVCacheSpecDesc.cc

@@ -72,6 +72,9 @@ CacheGroupPolicy SpecBuilder::groupPolicy(const KVCacheSpecDesc& desc) {
         if (desc.tail->active_tail_blocks.has_value()) {
             policy.active_tail_blocks = *desc.tail->active_tail_blocks;
         }
+        if (desc.tail->prefix_reuse_window_tokens.has_value()) {
+            policy.prefix_reuse_window_tokens = *desc.tail->prefix_reuse_window_tokens;
+        }
         if (desc.tail->validate_tail_blocks.has_value()) {
             policy.validate_tail_blocks = *desc.tail->validate_tail_blocks;
         }

rtp_llm/cpp/cache/KVCacheSpecDesc.h

@@ -43,6 +43,7 @@ struct CacheMemoryPolicyDesc {
 struct CacheTailPolicyDesc {
     std::optional<uint32_t> active_tail_blocks;
     std::optional<bool>     validate_tail_blocks;
+    std::optional<uint32_t> prefix_reuse_window_tokens;
 };
 
 struct CacheCpPolicyDesc {
@@ -72,6 +73,14 @@ struct KVCacheSpecDesc {
     size_t   block_stride_bytes_alignment       = 0;
     uint32_t block_stride_alignment_min_entries = 0;
 
+    // MHA geometry overrides; 0 means "take it from SpecBuildContext.attn_config".
+    // attn_config carries one global head count and one head dimension, which is not
+    // enough for models whose layer kinds disagree. MiMo V2.5 needs both: 4 KV heads on
+    // its global-attention layers versus 8 on its sliding-window layers, and V head dim
+    // 128 against QK 192.
+    uint32_t kv_head_num_override     = 0;
+    uint32_t v_size_per_head_override = 0;
+
     std::optional<CacheGroupType>          group_type;
     std::optional<CacheReusePolicyDesc>    reuse;
     std::optional<CacheCapacityPolicyDesc> capacity;