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 15, 06:00 UTC
[Bugfix] Fix TP inference for Flex attention backend (#19657)
Signed-off-by: Isotr0py <2037008807@qq.com>
Isotr0py committed
Jun 16, 2025 at 19:21 UTC
1173804dca83bbb23301e9cf97da0a43bd5ef41c
5 files changed
+54
-2
tests/v1/engine/test_engine_core.py
+35
-1
@@ -19,7 +19,7 @@ from vllm.v1.executor.abstract import Executor, UniProcExecutor
19
from vllm.v1.kv_cache_interface import KVCacheConfig
20
from vllm.v1.outputs import ModelRunnerOutput
21
22
-from ...utils import create_new_process_for_each_test
22
+from ...utils import create_new_process_for_each_test, multi_gpu_test
23
24
if not current_platform.is_cuda():
25
pytest.skip(reason="V1 currently only supported on CUDA.",
@@ -378,3 +378,37 @@ def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):
378
# Odd steps schedules a new batch.
379
assert output is None
380
step += 1
381
+
382
+
383
+@multi_gpu_test(num_gpus=2)
384
+def test_engine_core_tp(monkeypatch: pytest.MonkeyPatch):
385
+ """
386
+ Test engine can initialize worker in tp properly
387
+ """
388
+
389
+ with monkeypatch.context() as m:
390
+ m.setenv("VLLM_USE_V1", "1")
391
+ """Setup the EngineCore."""
392
+ engine_args = EngineArgs(
393
+ model=MODEL_NAME,
394
+ tensor_parallel_size=2,
395
+ # Reduce startup time.
396
+ enforce_eager=True,
397
+ )
398
+ vllm_config = engine_args.create_engine_config()
399
+ executor_class = Executor.get_class(vllm_config)
400
+
401
+ with set_default_torch_num_threads(1):
402
+ engine_core = EngineCore(vllm_config=vllm_config,
403
+ executor_class=executor_class,
404
+ log_stats=True)
405
+
406
+ def get_worker_cache_config_field(worker, key: str):
407
+ return getattr(worker.cache_config, key)
408
+
409
+ num_gpu_blocks = engine_core.collective_rpc(
410
+ get_worker_cache_config_field, args=("num_gpu_blocks", ))
411
+ num_cpu_blocks = engine_core.collective_rpc(
412
+ get_worker_cache_config_field, args=("num_cpu_blocks", ))
413
+ assert all(x is not None for x in num_gpu_blocks)
414
+ assert all(x is not None for x in num_cpu_blocks)
vllm/v1/attention/backends/flex_attention.py
+7
-1
@@ -13,6 +13,7 @@ from torch.nn.attention.flex_attention import (BlockMask, _mask_mod_signature,
13
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
14
AttentionMetadata, AttentionType,
15
is_quantized_kv_cache)
16
+from vllm.distributed import get_tensor_model_parallel_world_size
17
from vllm.logger import init_logger
18
from vllm.platforms import current_platform
19
from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
@@ -236,7 +237,12 @@ class FlexAttentionMetadata:
237
238
def build_block_mask(self) -> BlockMask:
239
assert self.mask_mod is not None
239
- return create_block_mask_compiled(
240
+ # FIXME: With TP>1, create_block_mask_compiled will raise
241
+ # CUDA error: an illegal memory access was encountered
242
+ create_block_mask_fn = (create_block_mask_compiled
243
+ if get_tensor_model_parallel_world_size() == 1
244
+ else create_block_mask)
245
+ return create_block_mask_fn(
246
self.mask_mod,
247
None,
248
None,
vllm/v1/engine/core.py
+2
@@ -84,6 +84,8 @@ class EngineCore:
84
85
vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks
86
vllm_config.cache_config.num_cpu_blocks = num_cpu_blocks
87
+ self.collective_rpc("initialize_cache",
88
+ args=(num_gpu_blocks, num_cpu_blocks))
89
90
self.structured_output_manager = StructuredOutputManager(vllm_config)
91
vllm/v1/worker/gpu_worker.py
+5
@@ -112,6 +112,11 @@ class Worker(WorkerBase):
112
buffer.data.copy_(self._sleep_saved_buffers[name].data)
113
self._sleep_saved_buffers = {}
114
115
+ def initialize_cache(self, num_gpu_blocks: int,
116
+ num_cpu_blocks: int) -> None:
117
+ self.cache_config.num_gpu_blocks = num_gpu_blocks
118
+ self.cache_config.num_cpu_blocks = num_cpu_blocks
119
+
120
def init_device(self):
121
if self.device_config.device.type == "cuda":
122
# torch.distributed.all_reduce does not free the input tensor until
vllm/v1/worker/tpu_worker.py
+5
@@ -93,6 +93,11 @@ class TPUWorker:
93
if self.model_config.seed is None:
94
self.model_config.seed = 0
95
96
+ def initialize_cache(self, num_gpu_blocks: int,
97
+ num_cpu_blocks: int) -> None:
98
+ self.cache_config.num_gpu_blocks = num_gpu_blocks
99
+ self.cache_config.num_cpu_blocks = num_cpu_blocks
100
+
101
def init_device(self):
102
os.environ["PJRT_DEVICE"] = "TPU"
103
# Note: Currently the XLA compiler wrongly uses 2D ring strategy on 1D