[Bugfix][V1] Sync mamba_block_size via EngineCoreReadyResponse (#50809)
Signed-off-by: “新颐” <lxy537258@antgroup.com> Co-authored-by: “新颐” <lxy537258@antgroup.com>
lxy committed
Aug 20, 2026 at 00:44 UTC
0c8c3f41cff4a4d361d4a80804eddd8dfc2a56ce
4 files changed
+39
tests/v1/engine/test_engine_core_client.py
+36
@@ -343,6 +343,42 @@ def test_apply_ready_response_syncs_block_size():
343
assert client.vllm_config.cache_config.block_size == 1056
344
345
346
+def test_apply_ready_response_syncs_mamba_block_size():
347
+ import msgspec
348
+
349
+ client = object.__new__(MPClient)
350
+ client.vllm_config = SimpleNamespace(
351
+ cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0),
352
+ model_config=SimpleNamespace(max_model_len=8192),
353
+ )
354
+ client.stats_update_address = None
355
+
356
+ payload = msgspec.msgpack.encode(
357
+ EngineCoreReadyResponse(
358
+ max_model_len=8192,
359
+ num_gpu_blocks=100,
360
+ block_size=1056,
361
+ dp_stats_address=None,
362
+ dtype="bfloat16",
363
+ vllm_version="test",
364
+ world_size=1,
365
+ data_parallel_size=1,
366
+ tensor_parallel_size=1,
367
+ pipeline_parallel_size=1,
368
+ decode_context_parallel_size=1,
369
+ data_parallel_rank=0,
370
+ max_num_seqs=256,
371
+ max_num_batched_tokens=8192,
372
+ instance_id="test-instance",
373
+ supports_lora=False,
374
+ max_loras=0,
375
+ mamba_block_size=1056,
376
+ )
377
+ )
378
+ client._apply_ready_response(payload)
379
+ assert client.vllm_config.cache_config.mamba_block_size == 1056
380
+
381
+
382
def loop_until_done(client: EngineCoreClient, outputs: dict):
383
while True:
384
engine_core_outputs = client.get_output().outputs
vllm/v1/engine/__init__.py
+1
@@ -90,6 +90,7 @@ class EngineCoreReadyResponse:
90
instance_id: str
91
supports_lora: bool
92
max_loras: int
93
+ mamba_block_size: int | None = None
94
# KV cache capacity (None for encoder-only/attention-free models).
95
kv_cache_size_tokens: int | None = None
96
kv_cache_max_concurrency: float | None = None
vllm/v1/engine/core.py
+1
@@ -1636,6 +1636,7 @@ class EngineCoreProc(EngineCore):
1636
max_model_len=self.vllm_config.model_config.max_model_len,
1637
num_gpu_blocks=self.vllm_config.cache_config.num_gpu_blocks or 0,
1638
block_size=self.vllm_config.cache_config.block_size,
1639
+ mamba_block_size=self.vllm_config.cache_config.mamba_block_size,
1640
dp_stats_address=self.frontend_stats_publish_address,
1641
dtype=str(self.vllm_config.model_config.dtype).removeprefix("torch."),
1642
vllm_version=VLLM_VERSION,
vllm/v1/engine/core_client.py
+1
@@ -758,6 +758,7 @@ class MPClient(EngineCoreClient):
758
# worker for hybrid Mamba models.
759
cache_config = vllm_config.cache_config
760
cache_config.block_size = response.block_size
761
+ cache_config.mamba_block_size = response.mamba_block_size
762
# Keep these as per-engine cache_config_info values; do not sum across DP.
763
cache_config.kv_cache_size_tokens = (
764
getattr(cache_config, "kv_cache_size_tokens", None)