[Bugfix][V1] Sync mamba_block_size via EngineCoreReadyResponse (#50809)

Signed-off-by: “新颐” <lxy537258@antgroup.com>
Co-authored-by: “新颐” <lxy537258@antgroup.com>

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

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

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,

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)