[Fix][Qwen3.5] Pass max_mamba_cache_size to mamba pool in disaggregation decode path (#19002)

This commit is contained in:
YAMY
2026-02-19 22:31:26 -08:00
committed by GitHub
parent 295bc17576
commit 8541b1118d
2 changed files with 6 additions and 1 deletions

View File

@@ -158,6 +158,7 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):
speculative_num_draft_tokens: int,
enable_mamba_extra_buffer: bool,
pre_alloc_size: int,
mamba_size: int = None,
):
DecodeReqToTokenPool.__init__(
self,
@@ -172,8 +173,11 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):
)
self.enable_mamba_extra_buffer = enable_mamba_extra_buffer
self.enable_memory_saver = enable_memory_saver
effective_mamba_size = (
mamba_size if mamba_size is not None else size
) + pre_alloc_size
self._init_mamba_pool(
size=size + pre_alloc_size,
size=effective_mamba_size,
mamba_spec_state_size=size + pre_alloc_size,
cache_params=cache_params,
device=device,

View File

@@ -445,6 +445,7 @@ class ModelRunnerKVCacheMixin:
speculative_num_draft_tokens=self.server_args.speculative_num_draft_tokens,
enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(),
pre_alloc_size=pre_alloc_size,
mamba_size=self.server_args.max_mamba_cache_size,
)
else:
self.req_to_token_pool = DecodeReqToTokenPool(