A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 22:45 UTC
[Hybrid] Mamba2 prefix cache blocks freeing for running requests (#28047)
Signed-off-by: Stanislaw Wozniak <stw@zurich.ibm.com> Signed-off-by: Chen Zhang <zhangch99@outlook.com> Co-authored-by: Chen Zhang <zhangch99@outlook.com>
Stan Wozniak committed
Dec 25, 2025 at 21:54 UTC
f1531d9f2a60c347ea1b410bfb04850c92922de0
1 file changed
+8
vllm/v1/core/single_type_kv_cache_manager.py
+8
@@ -737,6 +737,14 @@ class MambaManager(SingleTypeKVCacheManager):
737
)
738
return super().allocate_new_blocks(request_id, num_tokens)
739
740
+ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
741
+ """
742
+ Get the number of tokens whose mamba state are not needed anymore. Mamba only
743
+ need to keep the state of the last computed token, so we return
744
+ num_computed_tokens - 1.
745
+ """
746
+ return num_computed_tokens - 1
747
+
748
749
class CrossAttentionManager(SingleTypeKVCacheManager):
750
"""Manager for cross-attention KV cache in encoder-decoder models."""