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
[XPU] Set consistent default KV cache layout (#24745)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Sep 15, 2025 at 12:09 UTC
2e41f5abca796c615f4db8d9a496d037fa385653
3 files changed
+23
-16
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+9
-6
@@ -56,9 +56,9 @@ except ImportError:
56
logger.warning("NIXL is not available")
57
NixlWrapper = None
58
59
-# Supported xPUs and types of kv transfer buffer.
60
-# {xPU: tuple of supported kv buffer types}
61
-_NIXL_SUPPORTED_XPUS = {
59
+# Supported platforms and types of kv transfer buffer.
60
+# {device: tuple of supported kv buffer types}
61
+_NIXL_SUPPORTED_DEVICE = {
62
"cuda": ("cuda", ),
63
"tpu": ("cpu", ),
64
"xpu": ("cpu", ),
@@ -458,9 +458,9 @@ class NixlConnectorWorker:
458
self.device_type = current_platform.device_type
459
self.kv_buffer_device: str = \
460
vllm_config.kv_transfer_config.kv_buffer_device
461
- if self.device_type not in _NIXL_SUPPORTED_XPUS:
461
+ if self.device_type not in _NIXL_SUPPORTED_DEVICE:
462
raise RuntimeError(f"{self.device_type} is not supported.")
463
- elif self.kv_buffer_device not in _NIXL_SUPPORTED_XPUS[
463
+ elif self.kv_buffer_device not in _NIXL_SUPPORTED_DEVICE[
464
self.device_type]:
465
raise RuntimeError(
466
f"{self.device_type} with {self.kv_buffer_device} kv_buffer "
@@ -468,7 +468,7 @@ class NixlConnectorWorker:
468
self.device_kv_caches: dict[str, torch.Tensor] = {}
469
470
# cpu kv buffer for xfer
471
- # used when xPU memory can not be registered under nixl
471
+ # used when device memory can not be registered under nixl
472
self.host_xfer_buffers: dict[str, torch.Tensor] = {}
473
self.use_host_buffer = self.kv_buffer_device == "cpu"
474
if self.kv_buffer_device == "cuda":
@@ -927,6 +927,9 @@ class NixlConnectorWorker:
927
if tp_ratio > 1:
928
# Heterogeneous TP expects same kv_cache_layout.
929
assert nixl_agent_meta.kv_cache_layout == self.kv_cache_layout
930
+ if self.device_type == "xpu":
931
+ raise ValueError(
932
+ "Heterogeneous TP is not supported on XPU")
933
934
assert nixl_agent_meta.block_len == self.block_len * tp_ratio, (
935
"Remote P worker KV layer cache must be of shape [2, N, "
vllm/platforms/xpu.py
+4
-6
@@ -9,6 +9,7 @@ import torch
9
import vllm.envs as envs
10
from vllm.logger import init_logger
11
from vllm.utils import DEFAULT_MAX_NUM_BATCHED_TOKENS
12
+from vllm.v1.attention.backends.utils import set_kv_cache_layout
13
14
from .interface import DeviceCapability, Platform, PlatformEnum, _Backend
15
@@ -164,12 +165,9 @@ class XPUPlatform(Platform):
165
vllm_config.scheduler_config.max_model_len,
166
DEFAULT_MAX_NUM_BATCHED_TOKENS)
167
167
- if (envs.VLLM_KV_CACHE_LAYOUT is None
168
- or envs.VLLM_KV_CACHE_LAYOUT != "NHD"):
169
- os.environ["VLLM_KV_CACHE_LAYOUT"] = "NHD"
170
- logger.info(
171
- "Setting VLLM_KV_CACHE_LAYOUT to 'NHD' for XPU; "
172
- "only NHD layout is supported by XPU attention kernels.")
168
+ set_kv_cache_layout("NHD")
169
+ logger.info("Setting VLLM_KV_CACHE_LAYOUT to 'NHD' for XPU; "
170
+ "only NHD layout is supported by XPU attention kernels.")
171
172
@classmethod
173
def is_pin_memory_available(cls):
vllm/v1/attention/backends/utils.py
+10
-4
@@ -5,8 +5,8 @@ import enum
5
import functools
6
from abc import abstractmethod
7
from dataclasses import dataclass, fields, make_dataclass
8
-from typing import (TYPE_CHECKING, Any, ClassVar, Generic, Optional, Protocol,
9
- TypeVar)
8
+from typing import (TYPE_CHECKING, Any, ClassVar, Generic, Literal, Optional,
9
+ Protocol, TypeVar, Union, get_args)
10
11
import numpy as np
12
import torch
@@ -30,7 +30,12 @@ from vllm.logger import init_logger
30
from vllm.v1.kv_cache_interface import AttentionSpec
31
32
logger = init_logger(__name__)
33
-_KV_CACHE_LAYOUT_OVERRIDE = None
33
+KVCacheLayoutType = Literal["NHD", "HND"]
34
+_KV_CACHE_LAYOUT_OVERRIDE: Union[KVCacheLayoutType, None] = None
35
+
36
+
37
+def is_valid_kv_cache_layout(value: str) -> bool:
38
+ return value in get_args(KVCacheLayoutType)
39
40
41
@dataclass
@@ -296,12 +301,13 @@ def get_kv_cache_layout():
301
if cache_layout is None:
302
cache_layout = get_kv_connector_cache_layout()
303
else:
304
+ assert is_valid_kv_cache_layout(cache_layout)
305
logger.info_once("`VLLM_KV_CACHE_LAYOUT` environment variable " \
306
"detected. Setting KV cache layout to %s.", cache_layout)
307
return cache_layout
308
309
304
-def set_kv_cache_layout(cache_layout: str):
310
+def set_kv_cache_layout(cache_layout: KVCacheLayoutType):
311
global _KV_CACHE_LAYOUT_OVERRIDE
312
_KV_CACHE_LAYOUT_OVERRIDE = cache_layout
313