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 2, 13:30 UTC
Allow native KV cache dtype in Triton cache update (#43330)
Signed-off-by: Michael Gschwind <mgschwind@nvidia.com> Co-authored-by: Michael Gschwind <mgschwind@nvidia.com>
Mike G committed
May 28, 2026 at 09:51 UTC
53a20886752bdf8f26c4c5121458c514096da107
1 file changed
+10
-2
vllm/v1/attention/ops/triton_reshape_and_cache_flash.py
+10
-2
@@ -13,6 +13,14 @@ from vllm.utils.torch_utils import is_quantized_kv_cache
13
14
FP8_MIN, FP8_MAX = get_fp8_min_max()
15
16
+_NATIVE_KV_CACHE_DTYPES = {"auto", "float16", "bfloat16", "float32", "half", "float"}
17
+
18
+
19
+def _is_supported_kv_cache_dtype(kv_cache_dtype: str) -> bool:
20
+ return kv_cache_dtype in _NATIVE_KV_CACHE_DTYPES or is_quantized_kv_cache(
21
+ kv_cache_dtype
22
+ )
23
+
24
25
@triton.jit
26
def reshape_and_cache_kernel_flash(
@@ -350,7 +358,7 @@ def triton_reshape_and_cache_flash(
358
block_stride = key_cache.stride()[0]
359
page_stride = key_cache.stride()[1]
360
353
- assert kv_cache_dtype == "auto" or is_quantized_kv_cache(kv_cache_dtype), (
361
+ assert _is_supported_kv_cache_dtype(kv_cache_dtype), (
362
f"unsupported kv_cache_dtype (str), got {kv_cache_dtype}."
363
)
364
kv_cache_torch_dtype = (
@@ -529,7 +537,7 @@ def triton_reshape_and_cache_flash_diffkv(
537
block_stride = kv_cache.stride()[0]
538
page_stride = kv_cache.stride()[1]
539
532
- assert kv_cache_dtype == "auto" or is_quantized_kv_cache(kv_cache_dtype), (
540
+ assert _is_supported_kv_cache_dtype(kv_cache_dtype), (
541
f"unsupported kv_cache_dtype (str), got {kv_cache_dtype}."
542
)
543
kv_cache_torch_dtype = (