xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-09-09 14:57:06 +08:00
Code Issues Packages Projects Releases Wiki Activity

format

Browse Source 
Signed-off-by: inkcherry <mingzhi.liu@amd.com>
This commit is contained in:
inkcherry 2025-12-01 09:07:49 +00:00
parent 6fbeee78d1
commit bd6a340652
2 changed files with 29 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
vllm/distributed/kv_transfer/kv_connector/v1/moriio_connector.py
View File

@ -3,7 +3,6 @@
import contextlib import contextlib
import logging import logging
import math import math
import os
import queue import queue
import threading import threading
import time import time
@ -19,8 +18,8 @@ import msgspec
import numpy as np import numpy as np
import torch import torch
import zmq import zmq
from vllm import envs from vllm import envs
from vllm.attention.backends.registry import AttentionBackendEnum
from vllm.attention.selector import get_attn_backend from vllm.attention.selector import get_attn_backend
from vllm.config import VllmConfig from vllm.config import VllmConfig
from vllm.distributed.kv_transfer.kv_connector.v1.base import ( from vllm.distributed.kv_transfer.kv_connector.v1.base import (
@ -726,7 +725,7 @@ class MoRIIOWrapper:
return return
except (msgpack.exceptions.ExtraData, msgpack.exceptions.UnpackException): except (msgpack.exceptions.ExtraData, msgpack.exceptions.UnpackException):
logger.debug("Failed to decode msgpack message, will try as string. Error: %s") logger.debug("Failed to decode msgpack message, will try as string")
pass pass
try: try:
@ -809,9 +808,9 @@ class MoRIIOWrapper:
for path, sock in self.paths.items(): for path, sock in self.paths.items():
try: try:
sock.close(linger=0) sock.close(linger=0)
logger.debug(f"Closed ZMQ socket for path: {path}") logger.debug("Closed ZMQ socket for path: %s", path)
except Exception as e: except Exception as e:
logger.warning(f"Error closing ZMQ socket for path {path}: {e}") logger.warning("Error closing ZMQ socket for path %s: %s", path, e)
self.paths.clear() self.paths.clear()
@ -1292,6 +1291,15 @@ class MoRIIOConnectorScheduler:
return meta return meta
def shutdown(self):
for path, sock in self.paths.items():
try:
sock.close(linger=0)
logger.debug("Closed ZMQ socket for path: %s", path)
except Exception as e:
logger.warning("Error closing ZMQ socket for path %s: %s", path, e)
self.paths.clear()
def request_finished( def request_finished(
self, self,
request: "Request", request: "Request",
@ -1654,17 +1662,9 @@ class MoRIIOConnectorWorker:
index += 1 index += 1
def shutdown(self): def shutdown(self):
if hasattr(self, 'moriio_wrapper') and self.moriio_wrapper: if hasattr(self, "moriio_wrapper") and self.moriio_wrapper:
self.moriio_wrapper.shutdown() self.moriio_wrapper.shutdown()
for path, sock in self.paths.items():
try:
sock.close(linger=0)
logger.debug(f"Closed ZMQ socket for path: {path}")
except Exception as e:
logger.warning(f"Error closing ZMQ socket for path {path}: {e}")
self.paths.clear()
if hasattr(self, "_handshake_initiation_executor"): if hasattr(self, "_handshake_initiation_executor"):
self._handshake_initiation_executor.shutdown(wait=False) self._handshake_initiation_executor.shutdown(wait=False)
@ -1902,11 +1902,7 @@ class MoRIIOConnectorWorker:
caches_data = [] caches_data = []
for cache_or_caches in kv_caches.values(): for cache_or_caches in kv_caches.values():
cache_list = ( cache_list = [cache_or_caches] if use_mla else cache_or_caches
[cache_or_caches]
if use_mla
else cache_or_caches
)
for cache in cache_list: for cache in cache_list:
base_addr = cache.data_ptr() base_addr = cache.data_ptr()
region_len = self.num_blocks * self.block_len region_len = self.num_blocks * self.block_len

9
vllm/envs.py
View File

@ -1349,23 +1349,18 @@ environment_variables: dict[str, Callable[[], Any]] = {
), ),
# Controls the read mode for the Mori-IO connector # Controls the read mode for the Mori-IO connector
"VLLM_MORIIO_CONNECTOR_READ_MODE": lambda: ( "VLLM_MORIIO_CONNECTOR_READ_MODE": lambda: (
os.getenv("VLLM_MORIIO_CONNECTOR_READ_MODE", "False").lower() os.getenv("VLLM_MORIIO_CONNECTOR_READ_MODE", "False").lower() in ("true", "1")
in ("true", "1")
), ),
# Controls the QP (Queue Pair) per transfer configuration for the Mori-IO connector # Controls the QP (Queue Pair) per transfer configuration for the Mori-IO connector
"VLLM_MORIIO_QP_PER_TRANSFER": lambda: int( "VLLM_MORIIO_QP_PER_TRANSFER": lambda: int(
os.getenv("VLLM_MORIIO_QP_PER_TRANSFER", "1") os.getenv("VLLM_MORIIO_QP_PER_TRANSFER", "1")
), ),
# Controls the post-processing batch size for the Mori-IO connector # Controls the post-processing batch size for the Mori-IO connector
"VLLM_MORIIO_POST_BATCH_SIZE": lambda: int( "VLLM_MORIIO_POST_BATCH_SIZE": lambda: int(
os.getenv("VLLM_MORIIO_POST_BATCH_SIZE", "-1") os.getenv("VLLM_MORIIO_POST_BATCH_SIZE", "-1")
), ),
# Controls the number of workers for Mori operations for the Mori-IO connector # Controls the number of workers for Mori operations for the Mori-IO connector
"VLLM_MORIIO_NUM_WORKERS": lambda: int( "VLLM_MORIIO_NUM_WORKERS": lambda: int(os.getenv("VLLM_MORIIO_NUM_WORKERS", "1")),
os.getenv("VLLM_MORIIO_NUM_WORKERS", "1")
),
# Controls whether or not to use cudnn prefill # Controls whether or not to use cudnn prefill
"VLLM_USE_CUDNN_PREFILL": lambda: bool( "VLLM_USE_CUDNN_PREFILL": lambda: bool(
int(os.getenv("VLLM_USE_CUDNN_PREFILL", "0")) int(os.getenv("VLLM_USE_CUDNN_PREFILL", "0"))