wehub-resource-sync
248 lines
9.6 KiB
Python
248 lines
9.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Third Party
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache import torch_dev, torch_device_type
|
|
from lmcache.utils import EngineType
|
|
from lmcache.v1.config import LMCacheEngineConfig
|
|
from lmcache.v1.gpu_connector.gpu_connectors import GPUConnectorInterface
|
|
from lmcache.v1.gpu_connector.mock_gpu_connector import MockGPUConnector
|
|
from lmcache.v1.gpu_connector.utils import LayoutHints, need_gpu_interm_buffer
|
|
from lmcache.v1.metadata import LMCacheMetadata
|
|
|
|
# Boolean config flags whose underlying implementations exist only on a
|
|
# subset of accelerators. Each entry is ``(attr_name, human_label,
|
|
# supported_devices)``; the human label is what appears in the error message.
|
|
_DEVICE_SCOPED_VLLM_BOOL_FEATURES: tuple[tuple[str, str, frozenset[str]], ...] = (
|
|
("enable_blending", "config.enable_blending", frozenset({"cuda", "xpu"})),
|
|
("use_gpu_connector_v3", "config.use_gpu_connector_v3", frozenset({"cuda", "xpu"})),
|
|
)
|
|
|
|
|
|
def _validate_vllm_device_features(config: LMCacheEngineConfig) -> None:
|
|
"""Reject vLLM configurations that can't run on the active device.
|
|
|
|
Three classes of misconfiguration are rejected here, all *before* any
|
|
device-specific construction so the error surfaces as a plain
|
|
``ValueError`` instead of a deep ``torch.cuda.Stream()`` crash or a
|
|
``RuntimeError`` from ``torch.device('musa:0')`` on a non-MUSA build:
|
|
|
|
1. Any flag in :data:`_DEVICE_SCOPED_VLLM_BOOL_FEATURES` requested on
|
|
a device that has no implementation for that path.
|
|
2. ``config.use_layerwise=True`` on HPU — HPU only ships
|
|
:class:`VLLMPagedMemHPUConnectorV2`; previously, the dispatch silently
|
|
fell through to the CUDA layerwise connector.
|
|
|
|
Args:
|
|
config: The LMCache engine configuration.
|
|
|
|
Raises:
|
|
ValueError: For any of the unsupported combinations above.
|
|
"""
|
|
for attr, label, supported_devices in _DEVICE_SCOPED_VLLM_BOOL_FEATURES:
|
|
if torch_device_type not in supported_devices and getattr(config, attr):
|
|
supported = "/".join(sorted(supported_devices))
|
|
raise ValueError(
|
|
f"{label}=True is only supported on {supported}; the active "
|
|
f"device is '{torch_device_type}'. Set {attr}=False or run "
|
|
"on a supported accelerator build."
|
|
)
|
|
|
|
if torch_device_type == "hpu" and config.use_layerwise:
|
|
raise ValueError(
|
|
"config.use_layerwise=True is not supported on HPU; LMCache "
|
|
"ships no layerwise HPU connector. Set use_layerwise=False or "
|
|
"run on a CUDA-capable build."
|
|
)
|
|
|
|
|
|
def CreateGPUConnector(
|
|
config: LMCacheEngineConfig,
|
|
metadata: LMCacheMetadata,
|
|
engine: EngineType,
|
|
layout_hints: LayoutHints | None = None,
|
|
) -> GPUConnectorInterface:
|
|
"""
|
|
Create a GPU Connector based on the configuration and metadata.
|
|
|
|
Args:
|
|
config: The LMCache engine configuration.
|
|
metadata: The LMCache metadata.
|
|
engine: The serving engine type (EngineType.VLLM, EngineType.SGLANG,
|
|
EngineType.TRTLLM, or EngineType.MOCK).
|
|
layout_hints: Optional hints from the serving engine about KV cache
|
|
layout (e.g. ``{"kv_layout": "HND"}``).
|
|
"""
|
|
use_gpu = need_gpu_interm_buffer(config)
|
|
|
|
if engine == EngineType.SGLANG:
|
|
if torch_device_type == "musa":
|
|
raise ValueError(
|
|
"SGLang on MUSA is not supported; only the vLLM MUSA "
|
|
"connector is available."
|
|
)
|
|
|
|
num_layer, _, chunk_size, num_kv_head, head_dim = metadata.kv_shape
|
|
hidden_dim_size = num_kv_head * head_dim
|
|
local_worker_id = metadata.local_worker_id
|
|
torch_dev.set_device(local_worker_id)
|
|
device = torch.device(f"{torch_device_type}:{local_worker_id}")
|
|
kv_dtype = metadata.kv_dtype
|
|
|
|
if torch_device_type == "xpu":
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.xpu_connectors import (
|
|
SGLangLayerwiseXPUConnector,
|
|
SGLangXPUConnector,
|
|
)
|
|
|
|
if config.use_layerwise:
|
|
return SGLangLayerwiseXPUConnector(
|
|
hidden_dim_size,
|
|
num_layer,
|
|
use_gpu=use_gpu,
|
|
chunk_size=chunk_size,
|
|
dtype=kv_dtype,
|
|
device=device,
|
|
)
|
|
else:
|
|
return SGLangXPUConnector(
|
|
hidden_dim_size,
|
|
num_layer,
|
|
use_gpu=use_gpu,
|
|
chunk_size=chunk_size,
|
|
dtype=kv_dtype,
|
|
device=device,
|
|
)
|
|
else: # GPU for SGLang
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.gpu_connectors import (
|
|
SGLangGPUConnector,
|
|
SGLangLayerwiseGPUConnector,
|
|
)
|
|
|
|
if config.use_layerwise:
|
|
return SGLangLayerwiseGPUConnector(
|
|
hidden_dim_size,
|
|
num_layer,
|
|
use_gpu=use_gpu,
|
|
chunk_size=chunk_size,
|
|
dtype=kv_dtype,
|
|
device=device,
|
|
)
|
|
else:
|
|
return SGLangGPUConnector(
|
|
hidden_dim_size,
|
|
num_layer,
|
|
use_gpu=use_gpu,
|
|
chunk_size=chunk_size,
|
|
dtype=kv_dtype,
|
|
device=device,
|
|
)
|
|
elif engine == EngineType.VLLM:
|
|
_validate_vllm_device_features(config)
|
|
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.gpu_connectors import (
|
|
VLLMBufferLayerwiseGPUConnector,
|
|
VLLMPagedMemGPUConnectorV2,
|
|
VLLMPagedMemGPUConnectorV3,
|
|
VLLMPagedMemLayerwiseGPUConnector,
|
|
)
|
|
|
|
local_worker_id = metadata.local_worker_id
|
|
torch_dev.set_device(local_worker_id)
|
|
device = torch.device(f"{torch_device_type}:{local_worker_id}")
|
|
|
|
if torch_device_type == "cuda":
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.gpu_connectors import (
|
|
VLLMBufferLayerwiseGPUConnector,
|
|
VLLMPagedMemGPUConnectorV2,
|
|
VLLMPagedMemGPUConnectorV3,
|
|
VLLMPagedMemLayerwiseGPUConnector,
|
|
)
|
|
|
|
if config.use_layerwise:
|
|
if config.enable_blending:
|
|
return VLLMBufferLayerwiseGPUConnector.from_metadata(
|
|
metadata, use_gpu, device, layout_hints=layout_hints
|
|
)
|
|
else:
|
|
return VLLMPagedMemLayerwiseGPUConnector.from_metadata(
|
|
metadata, use_gpu, device, layout_hints=layout_hints
|
|
)
|
|
|
|
if config.use_gpu_connector_v3:
|
|
return VLLMPagedMemGPUConnectorV3.from_metadata(
|
|
metadata, use_gpu, device, layout_hints=layout_hints
|
|
)
|
|
else:
|
|
return VLLMPagedMemGPUConnectorV2.from_metadata(
|
|
metadata, use_gpu, device, layout_hints=layout_hints
|
|
)
|
|
elif torch_device_type == "xpu":
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.xpu_connectors import (
|
|
VLLMBufferLayerwiseXPUConnector,
|
|
VLLMPagedMemLayerwiseXPUConnector,
|
|
VLLMPagedMemXPUConnectorV2,
|
|
VLLMPagedMemXPUConnectorV3,
|
|
)
|
|
|
|
if config.use_layerwise:
|
|
if config.enable_blending:
|
|
return VLLMBufferLayerwiseXPUConnector.from_metadata(
|
|
metadata, use_gpu, device
|
|
)
|
|
else:
|
|
return VLLMPagedMemLayerwiseXPUConnector.from_metadata(
|
|
metadata, use_gpu, device
|
|
)
|
|
|
|
if config.use_gpu_connector_v3:
|
|
return VLLMPagedMemXPUConnectorV3.from_metadata(
|
|
metadata, use_gpu, device
|
|
)
|
|
else:
|
|
return VLLMPagedMemXPUConnectorV2.from_metadata(
|
|
metadata, use_gpu, device
|
|
)
|
|
elif torch_device_type == "musa":
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.musa_connectors import (
|
|
VLLMPagedMemLayerwiseMUSAConnector,
|
|
VLLMPagedMemMUSAConnectorV2,
|
|
)
|
|
|
|
if config.use_layerwise:
|
|
return VLLMPagedMemLayerwiseMUSAConnector.from_metadata(
|
|
metadata, use_musa=use_gpu, device=device
|
|
)
|
|
return VLLMPagedMemMUSAConnectorV2.from_metadata(metadata, use_gpu, device)
|
|
elif torch_device_type == "hpu":
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.hpu_connector import (
|
|
VLLMPagedMemHPUConnectorV2,
|
|
)
|
|
|
|
return VLLMPagedMemHPUConnectorV2.from_metadata(metadata, use_gpu, device)
|
|
else:
|
|
raise RuntimeError(f"No supported {torch_device_type} connector found.")
|
|
|
|
elif engine == EngineType.TRTLLM:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.gpu_connectors import TRTLLMGPUConnector
|
|
|
|
local_worker_id = metadata.local_worker_id
|
|
torch_dev.set_device(local_worker_id)
|
|
device = torch.device(f"{torch_device_type}:{local_worker_id}")
|
|
return TRTLLMGPUConnector.from_metadata(metadata, device=device)
|
|
|
|
elif engine == EngineType.MOCK:
|
|
kv_shape = metadata.kv_shape
|
|
return MockGPUConnector(kv_shape=kv_shape)
|
|
else:
|
|
raise RuntimeError(f"Unsupported engine: {engine}")
|