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, 01:00 UTC
Fix kv_cache_dtype handling for out-of-tree HPU plugin (#21302)
Signed-off-by: Konrad Zawora <kzawora@habana.ai> Signed-off-by: Chendi.Xue <chendi.xue@intel.com> Co-authored-by: Chendi.Xue <chendi.xue@intel.com>
Konrad Zawora committed
Jul 22, 2025 at 08:35 UTC
c17231e827991d5778e8ed258e7cdcb12c35b149
5 files changed
+30
-16
vllm/engine/arg_utils.py
+2
-16
@@ -1352,22 +1352,8 @@ class EngineArgs:
1352
1353
# No Fp8 KV cache so far.
1354
if self.kv_cache_dtype != "auto":
1355
- fp8_attention = self.kv_cache_dtype.startswith("fp8")
1356
- will_use_fa = (
1357
- current_platform.is_cuda()
1358
- and not envs.is_set("VLLM_ATTENTION_BACKEND")
1359
- ) or envs.VLLM_ATTENTION_BACKEND == "FLASH_ATTN_VLLM_V1"
1360
- supported = False
1361
- if (current_platform.is_rocm()
1362
- or (current_platform.is_cuda()
1363
- and current_platform.is_device_capability(100))
1364
- or current_platform.is_tpu()):
1365
- supported = True
1366
- elif fp8_attention and will_use_fa:
1367
- from vllm.attention.utils.fa_utils import (
1368
- flash_attn_supports_fp8)
1369
- supported = flash_attn_supports_fp8()
1370
-
1355
+ supported = current_platform.is_kv_cache_dtype_supported(
1356
+ self.kv_cache_dtype)
1357
if not supported:
1358
_raise_or_fallback(feature_name="--kv-cache-dtype",
1359
recommend_to_remove=False)
vllm/platforms/cuda.py
+13
@@ -586,6 +586,19 @@ class NonNvmlCudaPlatform(CudaPlatformBase):
586
" not found. Assuming no NVLink available.")
587
return False
588
589
+ @classmethod
590
+ def is_kv_cache_dtype_supported(cls, kv_cache_dtype: str) -> bool:
591
+ fp8_attention = kv_cache_dtype.startswith("fp8")
592
+ will_use_fa = (not envs.is_set("VLLM_ATTENTION_BACKEND")
593
+ ) or envs.VLLM_ATTENTION_BACKEND == "FLASH_ATTN_VLLM_V1"
594
+ supported = False
595
+ if cls.is_device_capability(100):
596
+ supported = True
597
+ elif fp8_attention and will_use_fa:
598
+ from vllm.attention.utils.fa_utils import flash_attn_supports_fp8
599
+ supported = flash_attn_supports_fp8()
600
+ return supported
601
+
602
603
# Autodetect either NVML-enabled or non-NVML platform
604
# based on whether NVML is available.
vllm/platforms/interface.py
+7
@@ -543,6 +543,13 @@ class Platform:
543
"""
544
raise RuntimeError(f"Unsupported torch distributed backend: {backend}")
545
546
+ @classmethod
547
+ def is_kv_cache_dtype_supported(cls, kv_cache_dtype: str) -> bool:
548
+ """
549
+ Returns if the kv_cache_dtype is supported by the current platform.
550
+ """
551
+ return False
552
+
553
554
class UnspecifiedPlatform(Platform):
555
_enum = PlatformEnum.UNSPECIFIED
vllm/platforms/rocm.py
+4
@@ -454,3 +454,7 @@ class RocmPlatform(Platform):
454
@classmethod
455
def device_count(cls) -> int:
456
return cuda_device_count_stateless()
457
+
458
+ @classmethod
459
+ def is_kv_cache_dtype_supported(cls, kv_cache_dtype: str) -> bool:
460
+ return True
\ No newline at end of file
vllm/platforms/tpu.py
+4
@@ -190,6 +190,10 @@ class TpuPlatform(Platform):
190
and params.sampling_type == SamplingType.RANDOM_SEED):
191
raise ValueError("Torch XLA does not support per-request seed.")
192
193
+ @classmethod
194
+ def is_kv_cache_dtype_supported(cls, kv_cache_dtype: str) -> bool:
195
+ return True
196
+
197
198
try:
199
from tpu_commons.platforms import TpuPlatform as TpuCommonsPlatform