[FIX] fix mambaish model pp kv cache compute (#17334)

2026-01-20 23:59:05 +08:00
parent d97066d209
commit 16802fb6b2
1 changed files with 6 additions and 1 deletions
@@ -124,7 +124,12 @@ class ModelRunnerKVCacheMixin:
                self.num_effective_layers,
            )
        elif mambaish := self.mambaish_config:
-            num_layers = len(mambaish.full_attention_layer_ids)
+            effective_layer_ids = [
                i
                for i in mambaish.full_attention_layer_ids
                if self.start_layer <= i < self.end_layer
            ]
            num_layers = len(effective_layer_ids)
        else:
            num_layers = self.num_effective_layers