A high-throughput and memory-efficient inference and serving engine for LLMs
Remove redundant Triton KV cache dtype asserts and enforce architectural support (fp8 >= sm89) (#43914)
Signed-off-by: Mike G <180722391+mikekg@users.noreply.github.com> Co-authored-by: Michael Gschwind <mgschwind@nvidia.com>
Mike G committed
Jun 15, 2026 at 06:49 UTC
fa63bb9db6f48108077fb5497081f7189092d163
2 files changed
+33
-39
vllm/v1/attention/backends/triton_attn.py
+20
@@ -464,6 +464,26 @@ class TritonAttentionImpl(AttentionImpl):
464
else:
465
self.sliding_window = (sliding_window - 1, 0)
466
self.kv_cache_dtype = kv_cache_dtype
467
+ cap = current_platform.get_device_capability()
468
+ cap_str = cap.as_version_str() if cap is not None else "unknown"
469
+ dev = current_platform.get_device_name()
470
+ if self.kv_cache_dtype.startswith("fp8") and not (
471
+ current_platform.has_device_capability(89)
472
+ ):
473
+ suggested = "float16" if (cap is None or cap.to_int() < 80) else "bfloat16"
474
+ raise ValueError(
475
+ f"FP8 KV cache is not supported by the Triton attention backend "
476
+ f"on {dev} (compute capability {cap_str}); native FP8 (fp8e4nv) "
477
+ f"requires SM89+. Re-run with --kv-cache-dtype {suggested}."
478
+ )
479
+ if self.kv_cache_dtype == "bfloat16" and not (
480
+ current_platform.has_device_capability(80)
481
+ ):
482
+ raise ValueError(
483
+ f"bfloat16 KV cache is not supported on {dev} (compute capability "
484
+ f"{cap_str}); bfloat16 requires SM80+. Re-run with "
485
+ f"--kv-cache-dtype float16."
486
+ )
487
if logits_soft_cap is None:
488
# In flash-attn, setting logits_soft_cap as 0 means no soft cap.
489
logits_soft_cap = 0
vllm/v1/attention/ops/triton_reshape_and_cache_flash.py
+13
-39
@@ -17,9 +17,16 @@ _NATIVE_KV_CACHE_DTYPES = {"auto", "float16", "bfloat16", "float32", "half", "fl
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
- )
20
+ if not (
21
+ kv_cache_dtype in _NATIVE_KV_CACHE_DTYPES
22
+ or is_quantized_kv_cache(kv_cache_dtype)
23
+ ):
24
+ return False
25
+ if kv_cache_dtype.startswith("fp8"):
26
+ return current_platform.has_device_capability(89)
27
+ if kv_cache_dtype == "bfloat16":
28
+ return current_platform.has_device_capability(80)
29
+ return True
30
31
32
@triton.jit
@@ -359,7 +366,9 @@ def triton_reshape_and_cache_flash(
366
page_stride = key_cache.stride()[1]
367
368
assert _is_supported_kv_cache_dtype(kv_cache_dtype), (
362
- f"unsupported kv_cache_dtype (str), got {kv_cache_dtype}."
369
+ f"Triton reshape-and-cache cannot store kv_cache_dtype={kv_cache_dtype} "
370
+ f"on this device: an FP8 KV cache needs native fp8e4nv (SM89+). Use "
371
+ f"--kv-cache-dtype bfloat16 (or float16 on SM75)."
372
)
373
kv_cache_torch_dtype = (
374
current_platform.fp8_dtype()
@@ -374,23 +383,7 @@ def triton_reshape_and_cache_flash(
383
# (e.g. explicit cast to fp8e4m3fnuz is not supported in triton 3.4)
384
key_cache = key_cache.view(kv_cache_torch_dtype)
385
value_cache = value_cache.view(kv_cache_torch_dtype)
377
- assert kv_cache_dtype != torch.uint8, (
378
- "explicit fp8 cast and store to "
379
- "uint8 is not supported by triton reshape_and_cache_flash"
380
- )
381
-
386
FP8_KV_CACHE = is_quantized_kv_cache(kv_cache_dtype)
383
- assert (not FP8_KV_CACHE) or kv_cache_torch_dtype in [
384
- torch.float8_e4m3fn,
385
- torch.float8_e5m2,
386
- torch.uint8,
387
- torch.float8_e4m3fnuz,
388
- ], (
389
- "unsupported dtype of KV cache tensor, got "
390
- "{kv_cache_torch_dtype}. Supported kv cache dtypes: fp8e4m3fn, "
391
- "fp8e5m2, uint8, bfloat16, float16, float32, fp8e4m3fnuz."
392
- )
393
-
387
# heuristics instead of autotuning
388
TILE_SIZE = min(2048, triton.next_power_of_2(n))
389
if current_platform.is_rocm() or current_platform.is_xpu():
@@ -537,9 +530,6 @@ def triton_reshape_and_cache_flash_diffkv(
530
block_stride = kv_cache.stride()[0]
531
page_stride = kv_cache.stride()[1]
532
540
- assert _is_supported_kv_cache_dtype(kv_cache_dtype), (
541
- f"unsupported kv_cache_dtype (str), got {kv_cache_dtype}."
542
- )
533
kv_cache_torch_dtype = (
534
current_platform.fp8_dtype()
535
if is_quantized_kv_cache(kv_cache_dtype)
@@ -550,23 +540,7 @@ def triton_reshape_and_cache_flash_diffkv(
540
# to avoid erounous implicit cast in triton kernel (tl.store to uint8)
541
# (e.g. explicit cast to fp8e4m3fnuz is not supported in triton 3.4)
542
kv_cache = kv_cache.view(kv_cache_torch_dtype)
553
- assert kv_cache_dtype != torch.uint8, (
554
- "explicit fp8 cast and store to "
555
- "uint8 is not supported by triton reshape_and_cache_flash_diffkv"
556
- )
557
-
543
FP8_KV_CACHE = is_quantized_kv_cache(kv_cache_dtype)
559
- assert (not FP8_KV_CACHE) or kv_cache_torch_dtype in [
560
- torch.float8_e4m3fn,
561
- torch.float8_e5m2,
562
- torch.uint8,
563
- torch.float8_e4m3fnuz,
564
- ], (
565
- "unsupported dtype of KV cache tensor, got "
566
- "{kv_cache_torch_dtype}. Supported kv cache dtypes: fp8e4m3fn, "
567
- "fp8e5m2, uint8, bfloat16, float16, float32, fp8e4m3fnuz."
568
- )
569
-
544
# heuristics instead of autotuning
545
TILE_SIZE = max(head_size_k, head_size_v)
546
TILE_SIZE = triton.next_power_of_2(TILE_SIZE)