From 0c8c3f41cff4a4d361d4a80804eddd8dfc2a56ce Mon Sep 17 00:00:00 2001 From: lxy <145345338+lxyxinyi@users.noreply.github.com> Date: Thu, 20 Aug 2026 00:44:36 +0800 Subject: [PATCH] [Bugfix][V1] Sync mamba_block_size via EngineCoreReadyResponse (#50809) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: “新颐” Co-authored-by: “新颐” --- tests/v1/engine/test_engine_core_client.py | 36 ++++++++++++++++++++++ vllm/v1/engine/__init__.py | 1 + vllm/v1/engine/core.py | 1 + vllm/v1/engine/core_client.py | 1 + 4 files changed, 39 insertions(+) diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index 1727c51b19..7424b76223 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -343,6 +343,42 @@ def test_apply_ready_response_syncs_block_size(): assert client.vllm_config.cache_config.block_size == 1056 +def test_apply_ready_response_syncs_mamba_block_size(): + import msgspec + + client = object.__new__(MPClient) + client.vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0), + model_config=SimpleNamespace(max_model_len=8192), + ) + client.stats_update_address = None + + payload = msgspec.msgpack.encode( + EngineCoreReadyResponse( + max_model_len=8192, + num_gpu_blocks=100, + block_size=1056, + dp_stats_address=None, + dtype="bfloat16", + vllm_version="test", + world_size=1, + data_parallel_size=1, + tensor_parallel_size=1, + pipeline_parallel_size=1, + decode_context_parallel_size=1, + data_parallel_rank=0, + max_num_seqs=256, + max_num_batched_tokens=8192, + instance_id="test-instance", + supports_lora=False, + max_loras=0, + mamba_block_size=1056, + ) + ) + client._apply_ready_response(payload) + assert client.vllm_config.cache_config.mamba_block_size == 1056 + + def loop_until_done(client: EngineCoreClient, outputs: dict): while True: engine_core_outputs = client.get_output().outputs diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index e07df08f3d..0a7f440ef7 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -90,6 +90,7 @@ class EngineCoreReadyResponse: instance_id: str supports_lora: bool max_loras: int + mamba_block_size: int | None = None # KV cache capacity (None for encoder-only/attention-free models). kv_cache_size_tokens: int | None = None kv_cache_max_concurrency: float | None = None diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 8946271aaf..96010465f3 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -1636,6 +1636,7 @@ class EngineCoreProc(EngineCore): max_model_len=self.vllm_config.model_config.max_model_len, num_gpu_blocks=self.vllm_config.cache_config.num_gpu_blocks or 0, block_size=self.vllm_config.cache_config.block_size, + mamba_block_size=self.vllm_config.cache_config.mamba_block_size, dp_stats_address=self.frontend_stats_publish_address, dtype=str(self.vllm_config.model_config.dtype).removeprefix("torch."), vllm_version=VLLM_VERSION, diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index a626af89e1..2bf7cf03c1 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -758,6 +758,7 @@ class MPClient(EngineCoreClient): # worker for hybrid Mamba models. cache_config = vllm_config.cache_config cache_config.block_size = response.block_size + cache_config.mamba_block_size = response.mamba_block_size # Keep these as per-engine cache_config_info values; do not sum across DP. cache_config.kv_cache_size_tokens = ( getattr(cache_config, "kv_cache_size_tokens", None)