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 2, 06:00 UTC
[Bugfix] Sync block_size from EngineCore to frontend for hybrid Mamba… (#42967)
Signed-off-by: Amit Gruner <agruner@crusoe.ai> Co-authored-by: Amit Gruner <agruner@crusoe.ai> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
gruner committed
Jun 2, 2026 at 16:41 UTC
654bd2bca412c0ab72eb4ada1d0fc271c7aed44e
4 files changed
+32
-1
tests/v1/engine/test_engine_core_client.py
+26
-1
@@ -27,12 +27,13 @@ from vllm.platforms import current_platform
27
from vllm.pooling_params import LateInteractionParams, PoolingParams
28
from vllm.usage.usage_lib import UsageContext
29
from vllm.utils.torch_utils import set_default_torch_num_threads
30
-from vllm.v1.engine import EngineCoreRequest
30
+from vllm.v1.engine import EngineCoreReadyResponse, EngineCoreRequest
31
from vllm.v1.engine.core import EngineCore
32
from vllm.v1.engine.core_client import (
33
AsyncMPClient,
34
DPLBAsyncMPClient,
35
EngineCoreClient,
36
+ MPClient,
37
SyncMPClient,
38
)
39
from vllm.v1.engine.utils import CoreEngineProcManager
@@ -236,6 +237,30 @@ def test_dplb_non_late_interaction_still_uses_lb():
237
assert client.lb_engines[1][0] == 1
238
239
240
+def test_apply_ready_response_syncs_block_size():
241
+ import msgspec
242
+
243
+ client = object.__new__(MPClient)
244
+ client.vllm_config = SimpleNamespace(
245
+ cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0),
246
+ model_config=SimpleNamespace(max_model_len=8192),
247
+ )
248
+ client.stats_update_address = None
249
+
250
+ payload = msgspec.msgpack.encode(
251
+ EngineCoreReadyResponse(
252
+ max_model_len=8192,
253
+ num_gpu_blocks=100,
254
+ block_size=1056,
255
+ dp_stats_address=None,
256
+ dtype="bfloat16",
257
+ vllm_version="test",
258
+ )
259
+ )
260
+ client._apply_ready_response(payload)
261
+ assert client.vllm_config.cache_config.block_size == 1056
262
+
263
+
264
def loop_until_done(client: EngineCoreClient, outputs: dict):
265
while True:
266
engine_core_outputs = client.get_output().outputs
vllm/v1/engine/__init__.py
+1
@@ -74,6 +74,7 @@ class EngineCoreReadyResponse:
74
75
max_model_len: int
76
num_gpu_blocks: int
77
+ block_size: int
78
dp_stats_address: str | None
79
dtype: str
80
vllm_version: str
vllm/v1/engine/core.py
+1
@@ -1462,6 +1462,7 @@ class EngineCoreProc(EngineCore):
1462
ready_response = EngineCoreReadyResponse(
1463
max_model_len=self.vllm_config.model_config.max_model_len,
1464
num_gpu_blocks=self.vllm_config.cache_config.num_gpu_blocks or 0,
1465
+ block_size=self.vllm_config.cache_config.block_size,
1466
dp_stats_address=self.frontend_stats_publish_address,
1467
dtype=str(self.vllm_config.model_config.dtype).removeprefix("torch."),
1468
vllm_version=VLLM_VERSION,
vllm/v1/engine/core_client.py
+4
@@ -713,6 +713,10 @@ class MPClient(EngineCoreClient):
713
num_gpu_blocks += response.num_gpu_blocks
714
vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks
715
716
+ # Sync block_size: may be enlarged by _align_hybrid_block_size in the
717
+ # worker for hybrid Mamba models.
718
+ vllm_config.cache_config.block_size = response.block_size
719
+
720
# In external DP LB mode, the coordinator address that the
721
# front-end procs connect to is obtained by each engine via it's
722
# initial handshake with the rank 0 front-end.