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 14, 05:30 UTC
[Bugfix] Fix pooling models on CPU backend (#23392)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
Li, Jiang committed
Aug 22, 2025 at 17:47 UTC
88016c372a5962eb98f4dfc71243ccd64433710e
1 file changed
+18
-2
vllm/utils/__init__.py
+18
-2
@@ -1440,6 +1440,12 @@ def _patched_set_stream(stream: torch.cuda.Stream) -> None:
1440
torch.cuda.set_stream = _patched_set_stream
1441
1442
1443
+class _StreamPlaceholder:
1444
+
1445
+ def __init__(self):
1446
+ self.synchronize = lambda: None
1447
+
1448
+
1449
def current_stream() -> torch.cuda.Stream:
1450
"""
1451
replace `torch.cuda.current_stream()` with `vllm.utils.current_stream()`.
@@ -1459,8 +1465,18 @@ def current_stream() -> torch.cuda.Stream:
1465
# On ROCm using the default 0 stream in combination with RCCL
1466
# is hurting performance. Therefore creating a dedicated stream
1467
# per process
1462
- _current_stream_tls.value = torch.cuda.Stream(
1463
- ) if current_platform.is_rocm() else torch.cuda.current_stream()
1468
+ if current_platform.is_rocm():
1469
+ _current_stream_tls.value = torch.cuda.Stream()
1470
+ elif current_platform.is_cpu():
1471
+ _current_stream_tls.value = _StreamPlaceholder()
1472
+ else:
1473
+ current_stream = current_platform.current_stream
1474
+ if current_stream is not None:
1475
+ _current_stream_tls.value = current_stream()
1476
+ else:
1477
+ raise ValueError(
1478
+ "Fail to set current stream, current platform "
1479
+ "may not support current_stream with torch API")
1480
return _current_stream_tls.value
1481
1482